1import re
2
3
[docs]
4class LazyGetMixin:
5 """
6 OmegaStore mixin to support chunked lazy get via name
7
8 Usage:
9
10 equivalent of ``om.datasets.get('foo', lazy=True).iterchunks()``::
11
12 mdf = om.datasets.get('foo#')
13 mdf = om.datasets.get('foo#iterchunks')
14
15 equivalent of ``om.datasets.get('foo', lazy=True).iterchunks(chunksize=10)``::
16
17 mdf = om.datasets.get('foo#iterchunks:chunksize=10')
18
19 equivalent of ``om.datasets.get('foo', lazy=True).iloc[0:10]``::
20
21 mdf = om.datasets.get('foo#rows:start=1,end=10')
22 """
23
24 # requires a trailing ; to work in all cases, see https://regex101.com/r/lYeKAw/1
25 ops_pattern = re.compile(r"(?P<name>.*)#(?P<opspec>.*?);(.*)$")
26
27 def metadata(self, name, *args, **kwargs):
28 if isinstance(name, str):
29 name, opspec = self._extract_opspec(name)
30 return super().metadata(name, *args, **kwargs)
31
32 def get(self, name, *args, **kwargs):
33 name, opspec = self._extract_opspec(name or '')
34 if opspec is not None:
35 kwargs = {**kwargs, **dict(lazy=True)}
36 lazy = super().get(name, *args, **kwargs)
37 if ':' in opspec:
38 op, op_kwargs_specs = opspec.split(':', 1)
39 op_kwargs = {}
40 for kw in op_kwargs_specs.split(','):
41 k, v = kw.split('=')
42 op_kwargs[k] = v
43 else:
44 op, op_kwargs = self._default_op(name, lazy)
45 meth = getattr(lazy, op, lambda *args, **kwargs: value)
46 value = meth(**op_kwargs)
47 else:
48 value = super().get(name, *args, **kwargs)
49 return value
50
51 def _extract_opspec(self, name):
52 match = self.ops_pattern.match(name + ';') if '#' in name else None
53 opspec = None
54 if match is not None:
55 name, opspec, _ = match.groups()
56 return name, opspec
57
58 def _default_op(self, name, lazy):
59 from omegaml.mdataframe import MDataFrame
60
61 if isinstance(lazy, MDataFrame):
62 op = 'iterchunks'
63 opkwargs = {}
64 else:
65 op = '__repr__'
66 opkwargs = {}
67 return op, opkwargs