1import re
2
3import pandas as pd
4
5
[docs]
6class ProjectedMixin:
7 """
8 A OmegaStore mixin to process column specifications in dataset name
9 """
10
11 colspec_pattern = re.compile(r"(?P<name>.*)\[(?P<colspec>.*)\].*$")
12
13 def metadata(self, name, *args, **kwargs):
14 if isinstance(name, str):
15 name, colspec = self._extract_column_specs(name)
16 return super().metadata(name, *args, **kwargs)
17
18 def get(self, name, *args, **kwargs):
19 """
20 Return a projected dataset given a name of form name[colspec]
21
22 colspec can be any of
23
24 * a comma separated list of columns, e.g. ``foo[a,b]``
25 * an open-ended slice, e.g. ``foo[a:]`` => all columns following a, inclusive
26 * an closed slice, e.g. ``foo[a:b]`` => all columns between a,b, inclusive
27 * a close-ended slice, e.g. ``foo[:b]`` => all columns up to b, inclusive
28 * an empty slice, e.g. ``foo[:]`` => all columns
29 * a list of columns to exclude, e.g. ``foo[^b]`` => all columns except b
30
31 :param name: (str) the name of the dataset, optionally including a
32 column specification
33 :return: the dataset with projected columns
34 """
35 # split base name from specs, get metadata
36 name, colspec = self._extract_column_specs(name)
37 if colspec is None:
38 # no column spec in name, avoid projection
39 data = super().get(name, *args, **kwargs)
40 else:
41 # column specs in name, get projected data
42 data = self._get_data_from_projection(name, colspec, *args, **kwargs)
43 return data
44
45 def _extract_column_specs(self, name):
46 colspec_pattern = self.colspec_pattern
47 match = colspec_pattern.match(name)
48 colspec = None
49 if match is not None:
50 name, colspec = match.groups()
51 return name, colspec
52
53 def _get_data_from_projection(self, name, colspec, *args, **kwargs):
54 # see if we can get columns from metadata
55 # if so we can specify the columns before getting the data
56 meta = self.metadata(name)
57 if 'columns' in meta.kind_meta:
58 colmap = meta.kind_meta['columns']
59 if isinstance(colmap, dict):
60 all_columns = list(colmap.keys())[1]
61 else:
62 # colmap is list of tuples (colname, storedname)
63 all_columns = list(zip(*colmap))[1]
64 columns = columnset(colspec, all_columns)
65 kwargs['columns'] = columns
66 data = super().get(name, *args, **kwargs)
67 else:
68 # we don't have columns in metadata, get the data first
69 # only subset on dataframes
70 data = super().get(name, *args, **kwargs)
71 if isinstance(data, pd.DataFrame):
72 all_columns = data.columns
73 name, columns = columnset(colspec, all_columns)
74 data = data[columns]
75 return data
76
77
78def columnset(colspec, all_columns):
79 """
80 find the specified columns in data[colspec] in list of all columns
81
82 colspec can be any of
83
84 * a comma separated list of columns, e.g. foo[a,b]
85 * an open-ended slice, e.g. foo[a:] => all columns following a, inclusive
86 * an closed slice, e.g. foo[a:b] => all columns between a,b, inclusive
87 * a close-ended slice, e.g. foo[:b] => all columns up to b, inclusive
88 * an empty slice, e.g. foo[:] => all columns
89 * a list of columns to exclude, e.g. foo[^b] => all columns except b
90 """
91 if colspec is not None:
92 if ':' in colspec:
93 from_col, to_col = colspec.split(':')
94 from_i = all_columns.index(from_col) if from_col else 0
95 to_i = (all_columns.index(to_col) if to_col else len(all_columns)) + 1
96 columns = all_columns[from_i:to_i]
97 elif colspec.startswith('^'):
98 columns = [col for col in all_columns if col not in colspec[1:].split(',')]
99 else:
100 columns = colspec.split(',')
101 else:
102 columns = all_columns
103 return columns