1from io import BytesIO
2
3import numpy as np
4
5from omegaml.backends.basedata import BaseDataBackend
6
7
[docs]
8class NumpyNDArrayBackend(BaseDataBackend):
9 """
10 Store numpy NDArray of any shape or size
11
12 The NDArray is serialized to a byte string and stored as a BLOB.
13 Thus it can have arbitrary size and dimensions, ideal for image
14 data and Tensors.
15 """
16
17 KIND = 'ndarray.bin'
18
19 _save_method = 'np.save'
20
[docs]
21 @classmethod
22 def supports(self, obj, name, as_pydata=False, **kwargs):
23 return not as_pydata and isinstance(obj, np.ndarray)
24
[docs]
25 def put(self, obj, name, attributes=None, allow_pickle=False, **kwargs):
26 # TODO associate meta.gridfile with actual fs file
27 kind_meta = {
28 'dtype': obj.dtype.name,
29 'shape': obj.shape,
30 'save_method': self._save_method,
31 'allow_pickle': allow_pickle,
32 }
33 buf = BytesIO()
34 np.save(buf, obj, allow_pickle=allow_pickle)
35 buf.seek(0)
36 fn = self.data_store.object_store_key(name, 'np', hashed=True)
37 gridfile = self._store_to_file(self.data_store, buf, fn)
38 return self.data_store.make_metadata(
39 name,
40 self.KIND,
41 attributes=attributes,
42 kind_meta=kind_meta,
43 gridfile=gridfile,
44 **kwargs,
45 ).save()
46
[docs]
47 def get(self, name, version=-1, force_python=False, lazy=False, **kwargs):
48 # TODO read file from meta.gridfile if exists, fallback to .fs only if not
49 meta = self.data_store.metadata(name)
50 fin = meta.gridfile
51 if meta.kind_meta.get('save_method') == self._save_method:
52 allow_pickle = meta.kind_meta.get('allow_pickle')
53 loaded = self._load_from_npsave(fin, allow_pickle)
54 else:
55 dtype_name = meta.kind_meta['dtype']
56 dtype = getattr(np, dtype_name, None)
57 shape = meta.kind_meta['shape']
58 loaded = self._load_from_stringbuffer(fin, dtype, shape)
59 return loaded
60
61 def _load_from_stringbuffer(self, fin, dtype, shape):
62 return np.frombuffer(fin.read(), dtype=dtype).reshape(*shape)
63
64 def _load_from_npsave(self, fin, allow_pickle):
65 buf = BytesIO()
66 buf.write(fin.read())
67 fin.close()
68 buf.seek(0)
69 return np.load(buf, allow_pickle=allow_pickle)