1from pymongo.collection import Collection
2
3from omegaml.backends.basedata import BaseDataBackend
4from omegaml.mdataframe import MDataFrame
5from omegaml.util import PickableCollection, json_normalize
6
7
[docs]
8class PandasRawDictBackend(BaseDataBackend):
9 """
10 OmegaStore backend to support arbitrary collections
11
12 Usage::
13
14 # store any collection as part of metadata
15 coll = db['some_collection']
16 om.datasets.put(coll, 'foo')
17 => Metadata(name='foo', collection='some_collection', ...)
18 # parse the collection using pandas.io.json_normalize
19 df = om.datasets.get('foo')
20 # use an alternate parser that accepts dict|list(dict)
21 df = om.datasets.get('foo', parser=some_fn)
22 # get a MDataFrame
23 om.datasets.getl('foo')
24 # preserve all document keys, including _id
25 om.datasets.getl('foo', raw=True)
26 """
27
28 KIND = 'pandas.rawdict'
29
[docs]
30 @classmethod
31 def supports(self, obj, name, as_raw=None, data_store=None, **kwargs):
32 new_as_raw = as_raw and isinstance(obj, (dict, list, tuple))
33 new_as_collection = isinstance(obj, (Collection, PickableCollection))
34 exists_as_dict = not (new_as_raw or new_as_collection) and (name and data_store.metadata(name) is not None)
35 return new_as_raw or new_as_collection or exists_as_dict
36
[docs]
37 def get(self, name, version=-1, lazy=False, raw=False, parser=None, filter=None, **kwargs):
38 collection = self.data_store.collection(name)
39 # json_normalize needs a list of dicts to work, not a generator
40 json_normalizer = lambda v: json_normalize([r for r in v])
41 parser = parser or json_normalizer
42 query = filter or kwargs
43 mdf = MDataFrame(collection, query=query, parser=parser, raw=raw, **kwargs)
44 return mdf if lazy else mdf.value
45
[docs]
46 def put(self, obj, name, attributes=None, as_raw=None, **kwargs):
47 if isinstance(obj, (Collection, PickableCollection)):
48 # already a collection, import it to metadata
49 collection = obj
50 elif isinstance(obj, dict):
51 # actual data, a single document, just insert
52 collection = self.data_store.collection(name)
53 collection.insert_one(obj)
54 elif isinstance(obj, (list, tuple)) or hasattr(obj, '__iter__'):
55 # actual data, multiple documents, insert many
56 collection = self.data_store.collection(name)
57 collection.insert_many(obj)
58 else:
59 raise ValueError(f'cannot insert object of type {type(obj)}')
60 meta = self.data_store._make_metadata(
61 name,
62 kind=self.KIND,
63 collection=collection.name,
64 attributes=attributes,
65 **kwargs.get('meta_kwargs', {}),
66 )
67 return meta.save()