1from uuid import uuid4
2
3import numpy as np
4import pandas as pd
5from bson import Code
6from numpy import isscalar
7from pymongo.collection import Collection
8
9from omegaml.store import qops
10from omegaml.store.filtered import FilteredCollection
11from omegaml.store.query import Filter, MongoQ
12from omegaml.store.queryops import MongoQueryOps
13from omegaml.util import (
14 PickableCollection,
15 cursor_to_dataframe,
16 ensure_index,
17 extend_instance,
18 json_normalize,
19 make_list,
20 make_tuple,
21 restore_index,
22 restore_index_columns_order,
23)
24
25INSPECT_CACHE = []
26
27
[docs]
28class MGrouper:
29 """
30 a Grouper for MDataFrames
31 """
32
33 STATS_MAP = {'std': 'stdDevSamp', 'mean': 'avg'}
34
35 def __init__(self, mdataframe, collection, columns, sort=True):
36 self.mdataframe = mdataframe
37 self.collection = collection
38 self.columns = make_tuple(columns)
39 self.should_sort = sort
40
41 def __getattr__(self, attr):
42 if attr in self.columns:
43 return MSeriesGroupby(self, self.collection, attr)
44
45 def statfunc():
46 columns = self.columns or self._non_group_columns()
47 return self.agg(dict.fromkeys(columns, attr))
48
49 return statfunc
50
51 def __getitem__(self, item):
52 return self.__getattr__(item)
53
[docs]
54 def agg(self, specs):
55 """
56 shortcut for .aggregate
57 """
58 return self.aggregate(specs)
59
[docs]
60 def aggregate(self, specs, **kwargs):
61 """
62 aggregate by given specs
63
64 See the following link for a list of supported operations.
65 https://docs.mongodb.com/manual/reference/operator/aggregation/group/
66
67 :param specs: a dictionary of { column : function | list[functions] }
68 pairs.
69 """
70
71 def add_stats(specs, column, stat):
72 specs['%s_%s' % (column, stat)] = {'$%s' % MGrouper.STATS_MAP.get(stat, stat): '$%s' % column}
73
74 # generate $group command
75 _specs = {}
76 for column, stats in specs.items():
77 stats = make_tuple(stats)
78 for stat in stats:
79 add_stats(_specs, column, stat)
80 groupby = qops.GROUP(columns=self.columns, **_specs)
81 # execute and return a dataframe
82 pipeline = self._amend_pipeline([groupby])
83 data = self.collection.aggregate(pipeline, allowDiskUse=True)
84
85 def get_data():
86 # we need this to build a pipeline for from_records
87 # to process, otherwise the cursor will be exhausted already
88 for group in data:
89 _id = group.pop('_id')
90 if isinstance(_id, dict):
91 group.update(_id)
92 yield group
93
94 df = pd.DataFrame.from_records(get_data())
95 columns = make_list(self.columns)
96 if columns:
97 df = df.set_index(columns, drop=True)
98 return df
99
100 def _amend_pipeline(self, pipeline):
101 """amend pipeline with default ops on coll.aggregate() calls"""
102 if self.should_sort:
103 sort = qops.SORT(**dict(qops.make_sortkey('_id')))
104 pipeline.append(sort)
105 return pipeline
106
107 def _non_group_columns(self):
108 """get all columns in mdataframe that is not in columns"""
109 return [
110 col
111 for col in self.mdataframe.columns
112 if col not in self.columns and col != '_id' and not col.startswith('_idx') and not col.startswith('_om#')
113 ]
114
115 def _count(self):
116 count_columns = self._non_group_columns()
117 if len(count_columns) == 0:
118 count_columns.append('_'.join(self.columns) + '_count')
119 groupby = {"$group": {"_id": {k: "$%s" % k for k in self.columns}}}
120 for k in count_columns:
121 groupby['$group']['%s' % k] = {"$sum": 1}
122 pipeline = self._amend_pipeline([groupby])
123 if self.should_sort:
124 sort = qops.SORT(**dict(qops.make_sortkey('_id')))
125 pipeline.append(sort)
126 return list(self.collection.aggregate(pipeline, allowDiskUse=True))
127
[docs]
128 def count(self):
129 """return counts by group columns"""
130 counts = self._count()
131 # remove mongo object _id
132 for group in counts:
133 group.update(group.pop('_id'))
134 # transform results to dataframe, then return as pandas would
135 resultdf = pd.DataFrame(counts).set_index(make_list(self.columns), drop=True)
136 return resultdf
137
138 def __iter__(self):
139 """for each group returns the key and a Filter object"""
140 # reduce count to only one column
141 groups = getattr(self, self.columns[0])._count()
142 for group in groups:
143 keys = group.get('_id')
144 data = self.mdataframe._clone(query=keys)
145 yield keys, data
146
147
[docs]
148class MLocIndexer:
149 """
150 implements the LocIndexer for MDataFrames
151 """
152
153 def __init__(self, mdataframe, positional=False):
154 self.mdataframe = mdataframe
155 # if positional, any loc[spec] will be applied on the rowid only
156 self.positional = positional
157 # indicator will be set true if loc specs are from a range type (list, tuple, np.ndarray)
158 self._from_range = False
159
[docs]
160 def __getitem__(self, specs):
161 """
162 access by index
163
164 use as mdf.loc[specs] where specs is any of
165
166 * a list or tuple of scalar index values, e.g. .loc[(1,2,3)]
167 * a slice of values e.g. .loc[1:5]
168 * a list of slices, e.g. .loc[1:5, 2:3]
169
170 :return: the sliced part of the MDataFrame
171 """
172 filterq, projection = self._get_filter(specs)
173 df = self.mdataframe
174 if filterq:
175 df = self.mdataframe.query(filterq)
176 df.from_loc_indexer = True
177 df.from_loc_range = self._from_range
178 if projection:
179 df = df[projection]
180 if isinstance(self.mdataframe, MSeries):
181 df = df._as_mseries(df.columns[0])
182 if getattr(df, 'immediate_loc', False):
183 df = df.value
184 return df
185
186 def __setitem__(self, specs, value):
187 raise NotImplementedError()
188
189 def _get_filter(self, specs):
190 filterq = []
191 projection = None
192 if self.positional:
193 idx_cols = ['_om#rowid']
194 else:
195 idx_cols = self.mdataframe._get_frame_index()
196 flt_kwargs = {}
197 enumerable_types = (list, tuple, np.ndarray)
198 if isinstance(specs, np.ndarray):
199 specs = specs.tolist()
200 if (
201 isinstance(specs, enumerable_types)
202 and isscalar(specs[0])
203 and len(idx_cols) == 1
204 and not any(isinstance(s, slice) for s in specs)
205 ):
206 # single column index with list of scalar values
207 if self.positional and isinstance(specs, tuple) and len(specs) == 2 and all(isscalar(v) for v in specs):
208 # iloc[int, int] is a cell access
209 flt_kwargs[idx_cols[0]] = specs[0]
210 projection = self._get_projection(specs[1])
211 else:
212 flt_kwargs[f'{idx_cols[0]}__in'] = specs
213 self._from_range = True
214 elif isinstance(specs, (int, str)):
215 flt_kwargs[idx_cols[0]] = specs
216 else:
217 specs = make_tuple(specs)
218 # list/tuple of slices or scalar values, or MultiIndex
219 for i, spec in enumerate(specs):
220 if i < len(idx_cols):
221 col = idx_cols[i]
222 if isinstance(spec, slice):
223 self._from_range = True
224 start, stop = spec.start, spec.stop
225 if start is not None:
226 flt_kwargs[f'{col}__gte'] = start
227 if stop is not None:
228 if isinstance(stop, int):
229 stop -= int(self.positional)
230 flt_kwargs[f'{col}__lte'] = stop
231 elif isinstance(spec, enumerable_types) and isscalar(spec[0]):
232 self._from_range = True
233 # single column index with list of scalar values
234 # -- convert to list for PyMongo serialization
235 if isinstance(spec, np.ndarray):
236 spec = spec.tolist()
237 flt_kwargs[f'{col}__in'] = spec
238 elif isscalar(col):
239 flt_kwargs[col] = spec
240 else:
241 # we're out of index columns, let's look at columns
242 cur_proj = self._get_projection(spec)
243 # if we get a single column, that's the projection
244 # if we had a single column, now need to add more, create a list
245 # if we have a list already, extend that
246 if projection is None:
247 projection = cur_proj
248 elif isinstance(projection, list):
249 projection.extend(cur_proj)
250 else:
251 projection = [projection, cur_proj]
252 if flt_kwargs:
253 filterq.append(MongoQ(**flt_kwargs))
254 finalq = None
255 for q in filterq:
256 if finalq:
257 finalq |= q
258 else:
259 finalq = q
260 return finalq, projection
261
262 def _get_projection(self, spec):
263 columns = self.mdataframe.columns
264 if np.isscalar(spec):
265 return [spec]
266 if isinstance(spec, (tuple, list)):
267 assert all(columns.index(col) for col in columns)
268 return spec
269 if isinstance(spec, slice):
270 start, stop = spec.start, spec.stop
271 if all(isinstance(v, int) for v in (start, stop)):
272 start, stop, step = spec.indices(len(columns))
273 else:
274 start = columns.index(start) if start is not None else 0
275 stop = columns.index(stop) + 1 if stop is not None else len(columns)
276 return columns[slice(start, stop)]
277 raise IndexError
278
279
[docs]
280class MPosIndexer(MLocIndexer):
281 """
282 implements the position-based indexer for MDataFrames
283 """
284
285 def __init__(self, mdataframe):
286 super().__init__(mdataframe, positional=True)
287
288 def _get_projection(self, spec):
289 columns = self.mdataframe.columns
290 if np.isscalar(spec):
291 return columns[spec]
292 if isinstance(spec, (tuple, list)):
293 return [col for i, col in enumerate(spec) if i in spec]
294 if isinstance(spec, slice):
295 start, stop = slice.start, slice.stop
296 if start and not isinstance(start, int):
297 start = 0
298 if stop and not isinstance(stop, int):
299 # sliced ranges are inclusive
300 stop = len(columns)
301 return columns[slice(start, stop)]
302 raise IndexError
303
304
305class MSeriesGroupby(MGrouper):
306 """
307 like a MGrouper but limited to one column
308 """
309
310 def count(self):
311 """
312 return series count
313
314 :return: counts by group
315 """
316 # MGrouper will insert a _count column, see _count(). we remove
317 # that column again and return a series named as the group column
318 resultdf = super().count()
319 count_column = [col for col in resultdf.columns if col.endswith('_count')][0]
320 new_column = count_column.replace('_count', '')
321 resultdf = resultdf.rename(columns={count_column: new_column})
322 return resultdf[new_column]
323
324
[docs]
325class MDataFrame:
326 """
327 A DataFrame for mongodb
328
329 Performs out-of-core, lazy computOation on a mongodb cluster.
330 Behaves like a pandas DataFrame. Actual results are returned
331 as pandas DataFrames.
332 """
333
334 STATFUNCS = ['mean', 'std', 'min', 'max', 'sum', 'var']
335
336 def __init__(
337 self,
338 collection,
339 columns=None,
340 query=None,
341 limit=None,
342 skip=None,
343 sort_order=None,
344 force_columns=None,
345 immediate_loc=False,
346 auto_inspect=False,
347 normalize=False,
348 raw=False,
349 parser=None,
350 preparefn=None,
351 from_loc_range=False,
352 metadata=None,
353 **kwargs,
354 ):
355 self.collection = PickableCollection(collection)
356 # columns in frame
357 self.columns = make_tuple(columns) if columns else self._get_fields(raw=raw)
358 self.columns = [str(col) for col in self.columns]
359 # columns to sort by, defaults to not sorted
360 self.sort_order = sort_order
361 # top n documents to fetch
362 self.head_limit = limit
363 # top n documents to skip before returning
364 self.skip_topn = skip
365 # filter criteria
366 self.filter_criteria = query or {}
367 # force columns -- on output add columns not present
368 self.force_columns = force_columns or []
369 # was this created from the loc indexer?
370 self.from_loc_indexer = kwargs.get('from_loc_indexer', False)
371 # was the loc index used a range? Else a single value
372 self.from_loc_range = from_loc_range
373 # setup query for filter criteries, if provided
374 if self.filter_criteria:
375 # make sure we have a filtered collection with the criteria given
376 if isinstance(self.filter_criteria, dict):
377 self.query_inplace(**self.filter_criteria)
378 elif isinstance(self.filter_criteria, Filter):
379 self.query_inplace(self.filter_criteria)
380 else:
381 raise ValueError(f'Invalid query specification of type {type(self.filter_criteria)}')
382 # if immediate_loc is True, .loc and .iloc always evaluate
383 self.immediate_loc = immediate_loc
384 # __array__ will return this value if it is set, set it otherwise
385 self._evaluated = None
386 # set true to automatically capture inspects on .value. retrieve using .inspect(cached=True)
387 self.auto_inspect = auto_inspect
388 self._inspect_cache = INSPECT_CACHE
389 # apply mixins
390 self._applyto = str(self.__class__)
391 self._apply_mixins()
392 # parser to parse documents to dataframe
393 self._parser = json_normalize if normalize else parser
394 # prepare function to be applied just before returning from .value
395 self._preparefn = preparefn
396 # keep technical fields like _id, _idx etc
397 self._raw = raw
398 # metadata stored by omegaml (equiv. of metadata.kind_meta)
399 self.metadata = metadata or dict()
400
401 def _apply_mixins(self, *args, **kwargs):
402 """
403 apply mixins in defaults.OMEGA_MDF_MIXINS
404 """
405 from omegaml import settings
406
407 defaults = settings()
408 for mixin, applyto in defaults.OMEGA_MDF_MIXINS:
409 if any(v in self._applyto for v in applyto.split(',')):
410 extend_instance(self, mixin, *args, **kwargs)
411
412 def __getstate__(self):
413 # pickle support. note that the hard work is done in PickableCollection
414 data = dict(self.__dict__)
415 data.update(_evaluated=None)
416 data.update(_inspect_cache=None)
417 data.update(auto_inspect=self.auto_inspect)
418 data.update(_preparefn=self._preparefn)
419 data.update(_parser=self._parser)
420 data.update(_raw=self._raw)
421 data.update(collection=self.collection)
422 return data
423
424 def __reduce__(self):
425 state = self.__getstate__()
426 args = (self.collection,)
427 return _mdf_remake, args, state
428
429 def __setstate__(self, state):
430 # pickle support. note that the hard work is done in PickableCollection
431 for k, v in state.items():
432 setattr(self, k, v)
433
434 def _getcopy_kwargs(self, without=None):
435 """return all parameters required on a copy of this MDataFrame"""
436 kwargs = dict(
437 columns=self.columns,
438 sort_order=self.sort_order,
439 limit=self.head_limit,
440 skip=self.skip_topn,
441 from_loc_indexer=self.from_loc_indexer,
442 from_loc_range=self.from_loc_range,
443 immediate_loc=self.immediate_loc,
444 metadata=self.metadata,
445 query=self.filter_criteria,
446 auto_inspect=self.auto_inspect,
447 parser=self._parser,
448 preparefn=self._preparefn,
449 )
450 [kwargs.pop(k) for k in make_tuple(without or [])]
451 return kwargs
452
453 def __array__(self, dtype=None):
454 # FIXME inefficient. make MDataFrame a drop-in replacement for any numpy ndarray
455 # this evaluates every single time
456 if self._evaluated is None:
457 self._evaluated = array = self.value.values
458 else:
459 array = self._evaluated
460 return array
461
462 def __getattr__(self, attr):
463 if attr in MDataFrame.STATFUNCS:
464 return self.statfunc(attr)
465 if attr in self.columns:
466 kwargs = self._getcopy_kwargs()
467 kwargs.update(columns=attr)
468 return MSeries(self.collection, **kwargs)
469 raise AttributeError(attr)
470
471 def __getitem__(self, cols_or_slice):
472 """
473 select and project by column, columns, slice, masked-style filter
474
475 Masked-style filters work similar to pd.DataFrame/Series masks
476 but do not actually return masks but an instance of Filter. A
477 Filter is a delayed evaluation on the data frame.
478
479 # select all rows where any column is == 5
480 mdf = MDataFrame(coll)
481 flt = mdf == 5
482 mdf[flt]
483 =>
484
485 :param cols_or_slice: single column (str), multi-columns (list),
486 slice to select columns or a masked-style
487 :return: filtered MDataFrame or MSeries
488 """
489 if isinstance(cols_or_slice, str):
490 # column name => MSeries
491 return self._as_mseries(cols_or_slice)
492 elif isinstance(cols_or_slice, int):
493 # column number => MSeries
494 column = self.columns[cols_or_slice]
495 return self._as_mseries(column)
496 elif isinstance(cols_or_slice, (tuple, list)):
497 # list of column names => MDataFrame subset on columns
498 kwargs = self._getcopy_kwargs()
499 kwargs.update(columns=cols_or_slice)
500 return MDataFrame(self.collection, **kwargs)
501 elif isinstance(cols_or_slice, Filter):
502 kwargs = self._getcopy_kwargs()
503 kwargs.update(query=cols_or_slice.query)
504 return MDataFrame(self.collection, **kwargs)
505 elif isinstance(cols_or_slice, np.ndarray):
506 return self.iloc[cols_or_slice]
507 raise ValueError('unknown accessor type %s' % type(cols_or_slice))
508
509 def __setitem__(self, column, value):
510 # True for any scalar type, numeric, bool, string
511 if np.isscalar(value):
512 result = self.collection.update_many(filter=self.filter_criteria, update=qops.SET(column, value))
513 self.columns.append(column)
514 return self
515
516 def _clone(self, collection=None, **kwargs):
517 # convenience method to clone itself with updates
518 collection = collection if collection is not None else self.collection
519 return self.__class__(collection, **kwargs, **self._getcopy_kwargs(without=list(kwargs.keys())))
520
521 def statfunc(self, stat):
522 aggr = MGrouper(self, self.collection, [], sort=False)
523 return getattr(aggr, stat)
524
[docs]
525 def groupby(self, columns, sort=True):
526 """
527 Group by a given set of columns
528
529 :param columns: the list of columns
530 :param sort: if True sort by group key
531 :return: MGrouper
532 """
533 return MGrouper(self, self.collection, columns, sort=sort)
534
535 def _get_fields(self, raw=False):
536 result = []
537 doc = self.collection.find_one()
538 if doc is not None:
539 if raw:
540 result = list(doc.keys())
541 else:
542 result = [
543 str(col)
544 for col in doc.keys()
545 if col != '_id' and not col.startswith('_idx') and not col.startswith('_om#')
546 ]
547 return result
548
549 def _get_frame_index(self):
550 """return the dataframe's index columns"""
551 doc = self.collection.find_one()
552 if doc is None:
553 result = []
554 else:
555 result = restore_index_columns_order(doc.keys())
556 return result
557
558 def _get_frame_om_fields(self):
559 """return the dataframe's omega special fields columns"""
560 doc = self.collection.find_one()
561 if doc is None:
562 result = []
563 else:
564 result = [k for k in list(doc.keys()) if k.startswith('_om#')]
565 return result
566
567 def _as_mseries(self, column):
568 kwargs = self._getcopy_kwargs()
569 kwargs.update(columns=make_tuple(column))
570 return MSeries(self.collection, **kwargs)
571
[docs]
572 def inspect(self, explain=False, cached=False, cursor=None, raw=False):
573 """
574 inspect this dataframe's actual mongodb query
575
576 :param explain: if True explains access path
577 """
578 if not cached:
579 if isinstance(self.collection, FilteredCollection):
580 query = self.collection.query
581 else:
582 query = ('*',)
583 if explain:
584 cursor = cursor or self._get_cursor()
585 explain = cursor.explain()
586 data = {'projection': self.columns, 'query': query, 'explain': explain or 'specify explain=True'}
587 else:
588 data = self._inspect_cache
589 if not (raw or explain):
590 data = pd.DataFrame(json_normalize(data))
591 return data
592
593 def count(self):
594 """
595 projected number of rows when resolving
596 """
597 nrows = len(self)
598 counts = pd.Series(dict.fromkeys(self.columns, nrows), index=self.columns)
599 return counts
600
[docs]
601 def __len__(self):
602 """
603 the projected number of rows when resolving
604 """
605 # we reduce to just 1 column to reduce speed
606 short = self._clone()
607 short = short[self.columns[0]] if self.columns else short
608 return sum(1 for d in short._get_cursor())
609
610 @property
611 def shape(self):
612 """
613 return shape of dataframe
614 """
615 return len(self), len(self.columns)
616
617 @property
618 def ndim(self):
619 return len(self.shape)
620
621 @property
622 def value(self):
623 """
624 resolve the query and return a Pandas DataFrame
625
626 :return: the result of the query as a pandas DataFrame
627 """
628 cursor = self._get_cursor()
629 df = self._get_dataframe_from_cursor(cursor)
630 if self.auto_inspect:
631 self._inspect_cache.append(self.inspect(explain=True, cursor=cursor, raw=True))
632 # this ensures the equiv. of pandas df.loc[n] is a Series
633 if self.from_loc_indexer:
634 if len(df) == 1 and not self.from_loc_range:
635 idx = df.index
636 df = df.T
637 df = df[df.columns[0]]
638 if df.ndim == 1 and len(df) == 1 and not isinstance(idx, pd.MultiIndex):
639 # single row single dimension, numeric index only
640 df = df.iloc[0]
641 elif (df.ndim == 1 or df.shape[1] == 1) and not self.from_loc_range:
642 df = df[df.columns[0]]
643 if self._preparefn:
644 df = self._preparefn(df)
645 return df
646
647 def reset(self):
648 # TODO if head(), tail(), query(), .loc/iloc should return a new MDataFrame instance to avoid having a reset need
649 self.head_limit = None
650 self.skip_topn = None
651 self.filter_criteria = {}
652 self.force_columns = []
653 self.sort_order = None
654 self.from_loc_indexer = False
655 return self
656
657 def _get_dataframe_from_cursor(self, cursor):
658 """
659 from the given cursor return a DataFrame
660 """
661 df = cursor_to_dataframe(cursor, parser=self._parser)
662 df = self._restore_dataframe_proper(df)
663 return df
664
665 @property
666 def _index_meta(self):
667 return self.metadata.get('idx_meta') or dict()
668
669 def _restore_dataframe_proper(self, df):
670 df = restore_index(df, self._index_meta)
671 if '_id' in df.columns and not self._raw:
672 df.drop('_id', axis=1, inplace=True)
673 if self.force_columns:
674 missing = set(self.force_columns) - set(self.columns)
675 for col in missing:
676 df[col] = np.nan
677 return df
678
679 def _get_cursor(self):
680 projection = make_tuple(self.columns)
681 projection += make_tuple(self._get_frame_index())
682 if not self.sort_order:
683 # implicit sort
684 projection += make_tuple(self._get_frame_om_fields())
685 cursor = self.collection.find(projection=projection)
686 if self.sort_order:
687 cursor.sort(qops.make_sortkey(make_tuple(self.sort_order)))
688 if self.head_limit:
689 cursor.limit(self.head_limit)
690 if self.skip_topn:
691 cursor.skip(self.skip_topn)
692 return cursor
693
[docs]
694 def sort(self, columns):
695 """
696 sort by specified columns
697
698 :param columns: str of single column or a list of columns. Sort order
699 is specified as the + (ascending) or - (descending)
700 prefix to the column name. Default sort order is
701 ascending.
702 :return: the MDataFrame
703 """
704 self._evaluated = None
705 self.sort_order = make_tuple(columns)
706 return self
707
[docs]
708 def head(self, limit=10):
709 """
710 return up to limit numbers of rows
711
712 :param limit: the number of rows to return. Defaults to 10
713 :return: the MDataFrame
714 """
715 return self._clone(limit=limit)
716
717 def tail(self, limit=10):
718 """
719 return up to limit number of rows from last inserted values
720
721 :param limit:
722 :return:
723 """
724 tail_n = self.skip(len(self) - limit)
725 return self._clone(skip=tail_n)
726
[docs]
727 def skip(self, topn):
728 """
729 skip the topn number of rows
730
731 :param topn: the number of rows to skip.
732 :return: the MDataFrame
733 """
734 return self._clone(skip=topn)
735
[docs]
736 def merge(
737 self,
738 right,
739 on=None,
740 left_on=None,
741 right_on=None,
742 how='inner',
743 target=None,
744 suffixes=('_x', '_y'),
745 sort=False,
746 inspect=False,
747 filter=None,
748 ):
749 """
750 merge this dataframe with another dataframe. only left outer joins
751 are currently supported. the output is saved as a new collection,
752 target name (defaults to a generated name if not specified).
753
754 :param right: the other MDataFrame
755 :param on: the list of key columns to merge by
756 :param left_on: the list of the key columns to merge on this dataframe
757 :param right_on: the list of the key columns to merge on the other
758 dataframe
759 :param how: the method to merge. supported are left, inner, right.
760 Defaults to inner
761 :param target: the name of the collection to store the merge results
762 in. If not provided a temporary name will be created.
763 :param suffixes: the suffixes to apply to identical left and right
764 columns
765 :param sort: if True the merge results will be sorted. If False the
766 MongoDB natural order is implied.
767 :returns: the MDataFrame to the target MDataFrame
768 """
769 # validate input
770 supported_how = ["left", 'inner', 'right']
771 assert how in supported_how, "only %s merges are currently supported" % supported_how
772 for key in [on, left_on, right_on]:
773 if key:
774 assert isinstance(key, str), "only single column merge keys are supported (%s)" % key
775 if isinstance(right, (Collection, PickableCollection, FilteredCollection)):
776 right = MDataFrame(right)
777 assert isinstance(right, MDataFrame), "both must be MDataFrames, got right=%" % type(right)
778 if how == 'right':
779 # A right B == B left A
780 return right.merge(
781 self, on=on, left_on=right_on, right_on=left_on, how='left', target=target, suffixes=suffixes
782 )
783 # generate lookup parameters
784 on = on or '_id'
785 right_name = self._get_collection_name_of(right, right)
786 target_name = self._get_collection_name_of(target, '_temp.merge.%s' % uuid4().hex)
787 target_field = "%s_%s" % (right_name.replace('.', '_'), right_on or on)
788 """
789 TODO enable filter criteria on right dataframe. requires changing LOOKUP syntax from
790 equitly to arbitray match
791
792 if right.filter_criteria:
793 right_filter = [qops.MATCH(self._get_filter_criteria(**right.filter_criteria))]
794 else:
795 right_filter = None
796 """
797 right_filter = None
798 lookup = qops.LOOKUP(right_name, key=on, left_key=left_on, right_key=right_on, target=target_field)
799 # unwind merged documents from arrays to top-level document fields
800 unwind = qops.UNWIND(target_field, preserve=how != 'inner')
801 # get all fields from left, right
802 project = {}
803 for left_col in self.columns:
804 source_left_col = left_col
805 if left_col == '_id':
806 project[left_col] = 1
807 continue
808 if left_col.startswith('_idx'):
809 continue
810 if left_col.startswith('_om#'):
811 continue
812 if left_col != (on or left_on) and left_col in right.columns:
813 left_col = '%s%s' % (left_col, suffixes[0])
814 project[left_col] = "$%s" % source_left_col
815 for right_col in right.columns:
816 if right_col == '_id':
817 continue
818 if right_col.startswith('_idx'):
819 continue
820 if right_col.startswith('_om#'):
821 continue
822 if right_col == (on or right_on) and right_col == (on or left_on):
823 # if the merge field is the same in both frames, we already
824 # have it from left
825 continue
826 if right_col in self.columns:
827 left_col = '%s%s' % (right_col, suffixes[1])
828 else:
829 left_col = '%s' % right_col
830 project[left_col] = '$%s.%s' % (target_field, right_col)
831 expected_columns = list(project.keys())
832 if '_id' not in project:
833 project['_id'] = 0 # never copy objectids to avoid duplicate keys, unless requested
834 project = {"$project": project}
835 # store merged documents and return an MDataFrame to it
836 out = qops.OUT(target_name)
837 pipeline = [lookup, unwind, project]
838 if filter:
839 query = qops.MATCH(self._get_filter_criteria(**filter))
840 pipeline.append(query)
841 if sort:
842 sort_cols = make_list(on or [left_on, right_on])
843 sort_key = qops.make_sortkey(sort_cols)
844 sort = qops.SORT(**dict(sort_key))
845 pipeline.append(sort)
846 pipeline.append(out)
847 if inspect:
848 result = pipeline
849 else:
850 result = self.collection.aggregate(pipeline, allowDiskUse=True)
851 result = MDataFrame(self.collection.database[target_name], force_columns=expected_columns)
852 return result
853
854 def append(self, other):
855 if isinstance(other, Collection):
856 other = MDataFrame(other)
857 assert isinstance(other, MDataFrame), f"both must be MDataFrames, got other={type(other)}"
858 outname = self.collection.name
859 mrout = {'merge': outname, 'nonAtomic': True}
860 mapfn = Code(
861 """
862 function() {
863 this._id = ObjectId();
864 if(this['_om#rowid']) {
865 this['_om#rowid'] += %s;
866 }
867 emit(this._id, this);
868 }
869 """
870 % len(self)
871 )
872 reducefn = Code("""
873 function(key, value) {
874 return value;
875 }
876 """)
877 finfn = Code("""
878 function(key, value) {
879 return value;
880 }
881 """)
882 other.collection.map_reduce(mapfn, reducefn, mrout, finalize=finfn, jsMode=True)
883 unwind = {"$replaceRoot": {"newRoot": {"$ifNull": ["$value", "$$CURRENT"]}}}
884 output = qops.OUT(outname)
885 pipeline = [unwind, output]
886 self.collection.aggregate(pipeline, allowDiskUse=True)
887 return self
888
889 def _get_collection_name_of(self, some, default=None):
890 """
891 determine the collection name of the given parameter
892
893 returns the collection name if some is a MDataFrame, a Collection
894 or a string_type. Otherwise returns default
895 """
896 if isinstance(some, MDataFrame):
897 name = some.collection.name
898 elif isinstance(some, Collection):
899 name = some.name
900 else:
901 name = default
902 return name
903
904 def _get_filter_criteria(self, *args, **kwargs):
905 """
906 return mongo query from filter specs
907
908 this uses a Filter to produce the query from the kwargs.
909
910 :param args: a Q object or logical combination of Q objects
911 (optional)
912 :param kwargs: all AND filter criteria
913 """
914 if len(args) > 0:
915 q = args[0]
916 if isinstance(q, MongoQ):
917 filter_criteria = Filter(self.collection, q).query
918 elif isinstance(q, Filter):
919 filter_criteria = Filter(self.collection, q.q).query
920 else:
921 filter_criteria = Filter(self.collection, **kwargs).query
922 return filter_criteria
923
[docs]
924 def query_inplace(self, *args, **kwargs):
925 """
926 filters this MDataFrame and returns it.
927
928 Any subsequent operation on the dataframe will have the filter
929 applied. To reset the filter call .reset() without arguments.
930
931 :param args: a Q object or logical combination of Q objects
932 (optional)
933 :param kwargs: all AND filter criteria
934 :return: self
935 """
936 self._evaluated = None
937 self.filter_criteria = self._get_filter_criteria(*args, **kwargs)
938 self.collection = FilteredCollection(self.collection, query=self.filter_criteria)
939 return self
940
[docs]
941 def query(self, *args, **kwargs):
942 """
943 return a new MDataFrame with a filter criteria
944
945 Any subsequent operation on the new dataframe will have the filter
946 applied. To reset the filter call .reset() without arguments.
947
948 Note: Unlike pandas DataFrames, a filtered MDataFrame operates
949 on the same collection as the original DataFrame
950
951 :param args: a Q object or logical combination of Q objects
952 (optional)
953 :param kwargs: all AND filter criteria
954 :return: a new MDataFrame with the filter applied
955 """
956 effective_filter = dict(self.filter_criteria)
957 filter_criteria = self._get_filter_criteria(*args, **kwargs)
958 if '$and' in effective_filter:
959 effective_filter['$and'].extend(filter_criteria.get('$and'))
960 else:
961 effective_filter.update(filter_criteria)
962 coll = FilteredCollection(self.collection, query=effective_filter)
963 return self._clone(collection=coll, query=effective_filter)
964
[docs]
965 def create_index(self, keys, **kwargs):
966 """
967 create and index the easy way
968 """
969 keys, kwargs = MongoQueryOps().make_index(keys)
970 result = ensure_index(self.collection, keys, **kwargs)
971 return result
972
973 def list_indexes(self):
974 """
975 list all indices in database
976 """
977 return cursor_to_dataframe(self.collection.list_indexes())
978
979 def iterchunks(self, chunksize=100):
980 """
981 return an iterator
982
983 Args:
984 chunksize (int): number of rows in each chunk
985
986 Returns:
987 a dataframe of max. length chunksize
988 """
989 chunksize = int(chunksize)
990 i = 0
991 while True:
992 chunkdf = self.skip(i).head(chunksize).value
993 if len(chunkdf) == 0:
994 break
995 i += chunksize
996 yield chunkdf
997
998 def itertuples(self, chunksize=1000):
999 chunksize = int(chunksize)
1000 __doc__ = pd.DataFrame.itertuples.__doc__
1001
1002 for chunkdf in self.iterchunks(chunksize=chunksize):
1003 for row in chunkdf.iterrows():
1004 yield row
1005
1006 def iterrows(self, chunksize=1000):
1007 chunksize = int(chunksize)
1008 __doc__ = pd.DataFrame.iterrows.__doc__
1009
1010 for chunkdf in self.iterchunks(chunksize=chunksize):
1011 if isinstance(chunkdf, pd.DataFrame):
1012 for row in chunkdf.iterrows():
1013 yield row
1014 else:
1015 # Series does not have iterrows
1016 for i in range(0, len(chunkdf), chunksize):
1017 yield chunkdf.iloc[i : i + chunksize]
1018
1019 def iteritems(self):
1020 if not hasattr(pd.DataFrame, 'iteritems'):
1021 raise NotImplementedError('MDataFrame.iteritems has been removed since Pandas 2.0. Use .items instead.')
1022 __doc__ = pd.DataFrame.iteritems.__doc__
1023 return self.items()
1024
1025 def items(self):
1026 __doc__ = pd.DataFrame.items.__doc__
1027
1028 for col in self.columns:
1029 yield col, self[col].value
1030
1031 @property
1032 def loc(self):
1033 """
1034 Access by index
1035
1036 Use as mdf.loc[index_value]
1037
1038 :return: MLocIndexer
1039 """
1040 self._evaluated = None
1041 indexer = MLocIndexer(self)
1042 return indexer
1043
1044 @property
1045 def iloc(self):
1046 self._evaluated = None
1047 indexer = MPosIndexer(self)
1048 return indexer
1049
1050 def rows(self, start=None, end=None, chunksize=1000):
1051 # equivalent to .iloc[start:end].iteritems(),
1052 start, end, chunksize = (int(v) for v in (start, end, chunksize))
1053 return self.iloc[slice(start, end)].iterchunks(chunksize)
1054
1055 def __repr__(self):
1056 kwargs = ', '.join(f'{k}={v}' for k, v in self._getcopy_kwargs().items())
1057 return f"MDataFrame(collection={self.collection.name}, {kwargs})"
1058
1059
[docs]
1060class MSeries(MDataFrame):
1061 """
1062 Series implementation for MDataFrames
1063
1064 behaves like a DataFrame but limited to one column.
1065 """
1066
1067 def __init__(self, *args, **kwargs):
1068 super().__init__(*args, **kwargs)
1069 # true if only unique values apply
1070 self.is_unique = False
1071 # apply mixins
1072 self._applyto = str(self.__class__)
1073 self._apply_mixins(*args, **kwargs)
1074
1075 def __getitem__(self, cols_or_slice):
1076 if isinstance(cols_or_slice, Filter):
1077 return MSeries(self.collection, columns=self.columns, query=cols_or_slice.query)
1078 return super().__getitem__(cols_or_slice)
1079
1080 @property
1081 def name(self):
1082 return self.columns[0]
1083
[docs]
1084 def unique(self):
1085 """
1086 return the unique set of values for the series
1087
1088 :return: MSeries
1089 """
1090 self.is_unique = True
1091 return self
1092
1093 def _get_cursor(self):
1094 if self.is_unique:
1095 # this way indexes get applied
1096 cursor = self.collection.distinct(make_tuple(self.columns)[0])
1097 else:
1098 cursor = super()._get_cursor()
1099 return cursor
1100
1101 @property
1102 def value(self):
1103 """
1104 return the value of the series
1105
1106 this is a Series unless unique() was called. If unique()
1107 only distinct values are returned as an array, matching
1108 the behavior of a Series
1109
1110 :return: pandas.Series
1111 """
1112 cursor = self._get_cursor()
1113 column = make_tuple(self.columns)[0]
1114 if self.is_unique:
1115 # the .distinct() cursor returns a list of values
1116 # this is to make sure we return the same thing as pandas
1117 val = [v for v in cursor]
1118 else:
1119 val = self._get_dataframe_from_cursor(cursor)
1120 val = val[column]
1121 val.name = self.name
1122 if len(val) == 1 and self.from_loc_indexer:
1123 val = val.iloc[0]
1124 if self.auto_inspect:
1125 self._inspect_cache.append(self.inspect(explain=True, cursor=cursor, raw=True))
1126 if self._preparefn:
1127 df = self._preparefn(val)
1128 return val
1129
1130 def __repr__(self):
1131 kwargs = ', '.join(f'{k}={v}' for k, v in self._getcopy_kwargs().items())
1132 return f"MSeries(collection={self.collection.name}, {kwargs})"
1133
1134 @property
1135 def shape(self):
1136 return (len(self),)
1137
1138
1139def _mdf_remake(collection):
1140 # recreate a pickled MDF
1141 mdf = MDataFrame(collection)
1142 return mdf