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 KIND = 'ndarray.bin'
17
18 _save_method = 'np.save'
19
[docs]
20 @classmethod
21 def supports(self, obj, name, as_pydata=False, **kwargs):
22 return not as_pydata and isinstance(obj, np.ndarray)
23
[docs]
24 def put(self, obj, name, attributes=None, allow_pickle=False, **kwargs):
25 # TODO associate meta.gridfile with actual fs file
26 kind_meta = {
27 'dtype': obj.dtype.name,
28 'shape': obj.shape,
29 'save_method': self._save_method,
30 'allow_pickle': allow_pickle,
31 }
32 buf = BytesIO()
33 np.save(buf, obj, allow_pickle=allow_pickle)
34 buf.seek(0)
35 fn = self.data_store.object_store_key(name, 'np', hashed=True)
36 gridfile = self._store_to_file(self.data_store, buf, fn)
37 return self.data_store.make_metadata(name, self.KIND, attributes=attributes,
38 kind_meta=kind_meta,
39 gridfile=gridfile,
40 **kwargs).save()
41
[docs]
42 def get(self, name, version=-1, force_python=False, lazy=False, **kwargs):
43 # TODO read file from meta.gridfile if exists, fallback to .fs only if not
44 meta = self.data_store.metadata(name)
45 fin = meta.gridfile
46 if meta.kind_meta.get('save_method') == self._save_method:
47 allow_pickle = meta.kind_meta.get('allow_pickle')
48 loaded = self._load_from_npsave(fin, allow_pickle)
49 else:
50 dtype_name = meta.kind_meta['dtype']
51 dtype = getattr(np, dtype_name, None)
52 shape = meta.kind_meta['shape']
53 loaded = self._load_from_stringbuffer(fin, dtype, shape)
54 return loaded
55
56 def _load_from_stringbuffer(self, fin, dtype, shape):
57 return np.frombuffer(fin.read(), dtype=dtype).reshape(*shape)
58
59 def _load_from_npsave(self, fin, allow_pickle):
60 buf = BytesIO()
61 buf.write(fin.read())
62 fin.close()
63 buf.seek(0)
64 return np.load(buf, allow_pickle=allow_pickle)