1import hashlib
2import json
3from itertools import product
4from uuid import uuid4
5
6import pandas as pd
7
8from omegaml.documents import make_QueryCache
9from omegaml.mdataframe import MDataFrame, MSeries
10from omegaml.store import qops
11from omegaml.store.filtered import FilteredCollection
12from omegaml.util import extend_instance, make_tuple
13
14
[docs]
15class ApplyMixin:
16 """
17 Implements the apply() mixin supporting arbitrary functions to build aggregation pipelines
18
19 Note that .apply() does not execute immediately. Instead it builds an aggregation pipeline
20 that is executed on MDataFrame.value. Note that .apply() calls cannot be cascaded yet, i.e.
21 a later .apply() will override a previous.apply().
22
23 See ApplyContext for usage examples.
24 """
25
26 def __init__(self, *args, **kwargs):
27 super().__init__(*args, **kwargs)
28 self._init_mixin(*args, **kwargs)
29
30 def _init_mixin(self, *args, **kwargs):
31 self.apply_fn = kwargs.get('apply_fn', None)
32 # set to True if the pipeline is a facet operation
33 self.is_from_facet = kwargs.get('is_from_facet', False)
34 # index columns
35 self.index_columns = kwargs.get('index_columns', [])
36 # db alias
37 self._db_alias = kwargs.get('db_alias', self._ensure_db_connection())
38 # cache used on persist()
39 self.cache = kwargs.get('cache', ApplyCache(self._db_alias))
40
41 def _ensure_db_connection(self):
42 # must import _dbs, _connections locally to ensure mongoshim has been applied
43 from mongoengine.connection import _connections, _dbs
44
45 seek_db = self.collection.database
46 for alias, db in _dbs.items():
47 if db is seek_db:
48 self._db_alias = alias
49 break
50 else:
51 # fake connection register
52 alias = self._db_alias = f'omega-{uuid4().hex}'
53 _connections[alias] = seek_db.client
54 _dbs[alias] = seek_db
55 return self._db_alias
56
57 def nocache(self):
58 self.cache = None
59 return self
60
61 def reset_cache(self, full=False):
62 """
63 Reset the apply cache
64
65 :param full: if True will reset all caches for the collection, if False will only remove
66 the cache for the specific .apply operations
67 :return:
68 """
69 QueryCache = make_QueryCache(db_alias=self._db_alias)
70 if full:
71 QueryCache.objects.filter(value__collection=self.collection.name).delete()
72 else:
73 pipeline = self._build_pipeline()
74 key = self._make_cache_key(self.collection, pipeline)
75 QueryCache.objects.filter(key=key).delete()
76 return self
77
78 def _make_cache_key(self, collection, pipeline):
79 # remove random output value
80 if '$out' in pipeline[-1] and pipeline[-1]['$out'].startswith('cache'):
81 pipeline = list(pipeline)[:-1]
82 spipeline = json.dumps(pipeline, sort_keys=True)
83 data = f'{collection.name}_{spipeline}'.encode()
84 # SEC: CWE-916
85 # - status: wontfix
86 # - reason: hashcode is used purely for name resolution, not a security function
87 key = hashlib.md5(data).hexdigest()
88 return key
89
90 def _getcopy_kwargs(self, **kwargs):
91 kwargs = super()._getcopy_kwargs(**kwargs)
92 kwargs.update(
93 is_from_facet=self.is_from_facet,
94 index_columns=self.index_columns,
95 cache=self.cache,
96 apply_fn=self.apply_fn,
97 )
98 return kwargs
99
100 def noapply(self):
101 self.apply_fn = None
102 return self
103
104 def apply(self, fn, inplace=False, preparefn=None):
105 if inplace:
106 obj = self
107 else:
108 kwargs = self._getcopy_kwargs()
109 kwargs.update(preparefn=preparefn)
110 if isinstance(self, MSeries):
111 obj = MSeries(self.collection, **kwargs)
112 else:
113 obj = MDataFrame(self.collection, **kwargs)
114 obj.apply_fn = fn
115 return obj
116
117 def persist(self):
118 """
119 Execute and store results in cache
120
121 Any pipeline of the same operations, in the same order, on
122 the same collection will return the same result.
123 """
124 # generate a cache key
125 pipeline = self._build_pipeline()
126 key = self._make_cache_key(self.collection, pipeline)
127 outname = f'cache_{uuid4().hex}'
128 value = {
129 'collection': self.collection.name,
130 'result': outname,
131 }
132 # do usual processing, store result
133 # -- note we pass pipeline to avoid processing iterators twice
134 pipeline.append(
135 {
136 '$out': outname,
137 },
138 )
139 cursor = self._get_cursor(pipeline=pipeline, use_cache=False)
140 # consume cursor to store output (via $out)
141 for v in cursor:
142 pass
143 # set cache
144 self.cache.set(key, value)
145 return key
146
147 def set_index(self, columns):
148 self.index_columns = make_tuple(columns)
149 return self
150
151 def inspect(self, explain=False, *args, **kwargs):
152 if self.apply_fn:
153 details = {
154 'pipeline': self._build_pipeline(),
155 }
156 if explain:
157 details.update(self.__dict__)
158 return details
159 return super().inspect(*args, explain=explain, **kwargs)
160
161 def _execute(self):
162 ctx = ApplyContext(self, columns=self.columns)
163 try:
164 result = self.apply_fn(ctx)
165 except Exception as e:
166 msg = [repr(stage) for stage in ctx.stages] + [repr(e)]
167 raise RuntimeError(msg)
168 if result is None or isinstance(result, ApplyContext):
169 result = result or ctx
170 self.index_columns = self.index_columns or result.index_columns
171 return result
172 elif isinstance(result, list):
173 return result
174 elif isinstance(result, dict):
175 # expect a mapping of col=ApplyContext each with its own list of stages
176 # -- build a combined context by adding each expression
177 # this ensures any multi-stage projections are carried forward
178 facets = {}
179 for col, expr in result.items():
180 if isinstance(expr, ApplyContext):
181 facets[col] = list(expr)
182 project = {
183 '$project': {
184 col: '$' + expr.columns[0],
185 }
186 }
187 facets[col].append(project)
188 else:
189 facets[col] = expr
190 facet = {
191 '$facet': facets,
192 }
193 self.is_from_facet = True
194 return [facet]
195 raise ValueError(f'Cannot build pipeline from apply result of type {type(result)}')
196
197 def _build_pipeline(self):
198 pipeline = []
199 stages = self._execute()
200 pipeline.extend(stages)
201 self._amend_pipeline(pipeline)
202 return pipeline
203
204 def _amend_pipeline(self, pipeline):
205 """amend pipeline with default ops on coll.aggregate() calls"""
206 if self.sort_order:
207 sort = qops.SORT(**dict(qops.make_sortkey(self.sort_order)))
208 pipeline.append(sort)
209 return pipeline
210
211 def _get_cached_cursor(self, pipeline=None, use_cache=True):
212 pipeline = pipeline or self._build_pipeline()
213 if use_cache and self.cache:
214 key = self._make_cache_key(self.collection, pipeline)
215 entry = self.cache.get(key)
216 if entry is not None:
217 # read result
218 outname = entry.value['result']
219 return self.collection.database[outname].find()
220
221 def _get_cursor(self, pipeline=None, use_cache=True):
222 # for apply functions, call the apply function, expecting a pipeline in return
223 if self.apply_fn:
224 pipeline = pipeline or self._build_pipeline()
225 cursor = self._get_cached_cursor(pipeline=pipeline, use_cache=use_cache)
226 if cursor is None:
227 filter_criteria = self._get_filter_criteria()
228 cursor = FilteredCollection(self.collection).aggregate(
229 pipeline, filter=filter_criteria, allowDiskUse=True
230 )
231 else:
232 cursor = super()._get_cursor()
233 return cursor
234
235 def _get_dataframe_from_cursor(self, cursor):
236 df = super()._get_dataframe_from_cursor(cursor)
237 if self.is_from_facet:
238 # if this was from a facet pipeline (i.e. multi-column mapping), combine
239 # $facet returns one document for each stage.
240 frames = []
241 for col in df.columns:
242 coldf = pd.DataFrame(df[col].iloc[0]).set_index('_id')
243 frames.append(coldf)
244 df = pd.concat(frames, axis=1).reset_index()
245 df = self._restore_dataframe_proper(df)
246 # TODO write a unit test for this condition
247 if self.index_columns and all(col in df.columns for col in self.index_columns):
248 df.set_index(list(self.index_columns), inplace=True)
249 return df
250
251
[docs]
252class ApplyContext:
253 """
254 Enable apply functions
255
256 .apply(fn) will call fn(ctx) where ctx is an ApplyContext.
257 The context supports methods to apply functions in a Pandas-style apply manner. ApplyContext is extensible
258 by adding an extension class to defaults.OMEGA_MDF_APPLY_MIXINS.
259
260 Note that unlike a Pandas DataFrame, ApplyContext does not itself contain any data.
261 Rather it is part of an expression tree, i.e. the aggregation pipeline. Thus any
262 expressions applied are translated into operations on the expression tree. The expression
263 tree is evaluated on MDataFrame.value, at which point the ApplyContext nor the function
264 that created it are active.
265
266 Examples::
267
268 mdf.apply(lambda v: v * 5 ) => multiply every column in dataframe
269 mdf.apply(lambda v: v['foo'].dt.week) => get week of date for column foo
270 mdf.apply(lambda v: dict(a=v['foo'].dt.week,
271 b=v['bar'] * 5) => run multiple pipelines and get results
272
273 The callable passed to apply can be any function. It can either return None,
274 the context passed in or a list of pipeline stages.
275
276 # apply any of the below functions
277 mdf.apply(customfn)
278
279 # same as lambda v: v.dt.week
280 def customfn(ctx):
281 return ctx.dt.week
282
283 # simple pipeline
284 def customfn(ctx):
285 ctx.project(x={'$multiply: ['$x', 5]},)
286 ctx.project(y={'$divide: ['$x', 2]},)
287
288 # complex pipeline
289 def customfn(ctx):
290 return [
291 { '$match': ... },
292 { '$project': ... },
293 ]
294 """
295
296 def __init__(self, caller, columns=None, index=None):
297 self.caller = caller
298 self.columns = columns
299 self.index_columns = index or []
300 self.computed = []
301 self.stages = []
302 self.expressions = []
303 self._apply_mixins()
304
305 def _apply_mixins(self):
306 """
307 apply mixins in defaults.OMEGA_MDF_APPLY_MIXINS
308 """
309 from omegaml import settings
310
311 defaults = settings()
312 for mixin, applyto in defaults.OMEGA_MDF_APPLY_MIXINS:
313 if any(v in self.caller._applyto for v in applyto.split(',')):
314 extend_instance(self, mixin)
315
316 def __iter__(self):
317 # return pipeline stages
318 for stage in self.stages:
319 if isinstance(stage, ApplyContext):
320 for sub_stage in stage:
321 yield sub_stage
322 else:
323 yield stage
324
325 def __getitem__(self, sel):
326 """
327 return a stage subset on a column
328 """
329 subctx = ApplyContext(self.caller, columns=make_tuple(sel), index=self.index_columns)
330 self.add(subctx)
331 return subctx
332
333 def __setitem__(self, sel, val):
334 """
335 add a projection to a sub context
336
337 ctx['col'] = value-expression
338 """
339 mapping = {col: v for (col, v) in zip(make_tuple(sel), make_tuple(val))}
340 self.project(mapping)
341
342 def __repr__(self):
343 return f'ApplyContext(stages={self.stages}, expressions={self.expressions},)'
344
345 def add(self, stage):
346 """
347 Add a processing stage to the pipeline
348
349 see https://docs.mongodb.com/manual/meta/aggregation-quick-reference/
350 """
351 self.stages.append(stage)
352 return self
353
354 def project_keeper_columns(self):
355 # keep index, computed
356 index = {col: '$' + col for col in self.index_columns}
357 computed = {col: '$' + col for col in self.computed}
358 keep = {}
359 keep.update(index)
360 keep.update(computed)
361 project = self.project(keep, keep=True)
362 return project
363
364 def _getLastStageKind(self, kind):
365 # see if there is already an open projection stage
366 for stage in self.stages[::-1]:
367 if kind in stage:
368 return stage
369
370 def _getProjection(self, append=False):
371 stage = self._getLastStageKind('$project')
372 if stage is None or append:
373 stage = {
374 '$project': {
375 '_id': 1,
376 }
377 }
378 self.stages.append(stage)
379 return stage
380
381 def _getGroupBy(self, by=None, append=False):
382 stage = self._getLastStageKind('$group')
383 if stage and stage['$group']['_id'] != by and by != '$$last':
384 # if a different groupby criteria, add a new one
385 stage = None
386 if stage is None and by == '$$last':
387 by = None
388 if stage is None or append:
389 stage = {
390 '$group': {
391 '_id': by,
392 },
393 }
394 self.stages.append(stage)
395 return stage
396
397 def groupby(self, by, expr=None, append=None, **kwargs):
398 """
399 add a groupby accumulation using $group
400
401 :param by: the groupby columns, if provided as a list will be transformed
402 :param expr:
403 :param append:
404 :param kwargs:
405 :return:
406
407 """
408 by = make_tuple(by)
409 self.index_columns = self.index_columns + list(by)
410 # define groupby
411 by = {col: '$' + col for col in by}
412 stage = self._getGroupBy(by)
413 groupby = stage['$group']
414 # add acccumulators
415 expr = expr or {col: colExpr for col, colExpr in kwargs.items()}
416 groupby.update(expr)
417 # add a projection to extract groupby values
418 extractId = {col: '$_id.' + col for col in by}
419 # add a projection to keep accumulator columns
420 keepCols = dict.fromkeys(expr, 1)
421 keepCols.update(extractId)
422 self.project(keepCols, append=True)
423 # sort by groupby keys
424 self.add(
425 {
426 '$sort': dict.fromkeys(by, 1),
427 },
428 )
429 return self
430
431 def project(self, expr=None, append=False, keep=False, **kwargs):
432 """
433 add a projection using $project
434
435 :param expr: the column-operator mapping
436 :param append: if True add a $project stage, otherwise add to existing
437 :param kwargs: if expr is None, the column-operator mapping as kwargs
438 :return: ApplyContext
439
440 """
441 # get last $project stage in pipeline
442 stage = self._getProjection(append=append)
443 expr = expr or kwargs
444 self.expressions.append(expr)
445 for k, v in expr.items():
446 # only append to stage if no other column projection was there
447 project = stage.get('$project')
448 if k not in project:
449 project.update(
450 {k: v},
451 )
452 elif not keep:
453 # if a column is already projected, add a new projection stage
454 stage = self._getProjection(append=True)
455 project = stage.get('$project')
456 project.update(
457 {k: v},
458 )
459 return self
460
461
[docs]
462class ApplyArithmetics:
463 """
464 Math operators for ApplyContext
465
466 * :code:`__mul__` (*)
467 * :code:`__add__` (+)
468 * :code:`__sub__` (-)
469 * :code:`__div__` (/)
470 * :code:`__floordiv__` (//)
471 * :code:`__mod__` (%)
472 * :code:`__pow__` (pow)
473 * :code:`__ceil__` (ceil)
474 * :code:`__floor__` (floor)
475 * :code:`__trunc__` (trunc)
476 * :code:`__abs__` (abs)
477 * :code:`sqrt` (math.sqrt)
478
479 """
480
481 def __arithmop__(op, wrap_op=None):
482 """
483 return a pipeline $project stage math operator as
484 { col:
485 { '$operator': [ values, ...] }
486 ...
487 }
488
489 If wrap_op is specified, will wrap the $operator clause as
490 { col:
491 { '$wrap_op': { '$operator': [ values, ...] } }0
492 ...
493 }
494 """
495
496 def inner(self, other):
497 terms = []
498 for term in make_tuple(other):
499 if isinstance(term, str):
500 term = '$' + term
501 terms.append(term)
502
503 def wrap(expr):
504 if wrap_op is not None:
505 expr = {wrap_op: expr}
506 return expr
507
508 mapping = {
509 col: wrap(
510 {op: ['$' + col] + terms},
511 )
512 for col in self.columns
513 }
514 keepCols = {col: '$' + col for col in self.index_columns}
515 mapping.update(keepCols)
516 self.project(mapping)
517 return self
518
519 return inner
520
521 #: multiply
522 __mul__ = __arithmop__('$multiply')
523 #: add
524 __add__ = __arithmop__('$add')
525 #: subtract
526 __sub__ = __arithmop__('$subtract')
527 #: divide
528 __div__ = __arithmop__('$divide')
529 __truediv__ = __arithmop__('$divide')
530 #: divide integer
531 __floordiv__ = __arithmop__('$divide', wrap_op='$floor')
532 #: modulo (%)
533 __mod__ = __arithmop__('$mod')
534 #: pow
535 __pow_ = __arithmop__('$pow')
536 #: ceil
537 __ceil__ = __arithmop__('$ceil')
538 #: floor
539 __floor__ = __arithmop__('$floor')
540 #: truncate
541 __trunc__ = __arithmop__('$trunc')
542 #: absolute
543 __abs__ = __arithmop__('$abs')
544 #: square root
545 sqrt = __arithmop__('sqrt')
546
547
[docs]
548class ApplyDateTime:
549 """
550 Datetime operators for ApplyContext
551 """
552
553 @property
554 def dt(self):
555 return self
556
557 def __dtop__(op):
558 """
559 return a datetime $project operator as
560 { col:
561 { '$operator': '$col} }
562 ...
563 }
564 """
565
566 def inner(self, columns=None):
567 columns = make_tuple(columns or self.columns)
568 mapping = {col: {op: '$' + col} for col in columns}
569 self.project(mapping)
570 return self
571
572 inner.__doc__ = op.replace('$', '')
573 return inner
574
575 # mongodb mappings
576 _year = __dtop__('$year')
577 _month = __dtop__('$month')
578 _week = __dtop__('$week')
579 _dayOfWeek = __dtop__('$dayOfWeek')
580 _dayOfMonth = __dtop__('$dayOfMonth')
581 _dayOfYear = __dtop__('$dayOfYear')
582 _hour = __dtop__('$hour')
583 _minute = __dtop__('$minute')
584 _second = __dtop__('$second')
585 _millisecond = __dtop__('$millisecond')
586 _isoDayOfWeek = __dtop__('$isoDayOfWeek')
587 _isoWeek = __dtop__('$isoWeek')
588 _isoWeekYear = __dtop__('$isoWeekYear')
589
590 # .dt accessor convenience similar to pandas.dt
591 # see https://pandas.pydata.org/pandas-docs/stable/api.html#datetimelike-properties
592 year = property(_year)
593 month = property(_month)
594 day = property(_dayOfMonth)
595 hour = property(_hour)
596 minute = property(_minute)
597 second = property(_second)
598 millisecond = property(_millisecond)
599 week = property(_isoWeek)
600 dayofyear = property(_dayOfYear)
601 dayofweek = property(_dayOfWeek)
602
603
[docs]
604class ApplyString:
605 """
606 String operators
607 """
608
609 @property
610 def str(self):
611 return self
612
613 def __strexpr__(op, unwind=False, base=None, max_terms=None):
614 """
615 return a pipeline $project string operator as
616 { col:
617 { '$operator': [ values, ...] }
618 ...
619 }
620 """
621
622 def inner(self, other, *args):
623 # get all values passed and build terms from them
624 values = list(make_tuple(other) + args)
625 terms = []
626 for term in values:
627 if isinstance(term, str):
628 # if the term is a column name, add as a column name
629 if term in self.columns:
630 term = '$' + term
631 # allow to specify values explicitely by $$<value> => <value>
632 term = term.replace('$$', '')
633 terms.append(term)
634 # limit number of terms if requested
635 if max_terms:
636 terms = terms[:max_terms]
637 # add projection of output columns to operator
638 mapping = {
639 col: {
640 op: terms if base is None else ['$' + col] + terms,
641 }
642 for col in self.columns
643 }
644 self.project(mapping)
645 # unwind all columns if requested
646 if unwind:
647 exprs = [{'$unwind': {'path': '$' + col}} for col in self.columns]
648 self.stages.extend(exprs)
649 return self
650
651 inner.__doc__ = op.replace('$', '')
652 return inner
653
654 def __strunary__(op, unwind=False):
655 """
656 return a datetime $project operator as
657 { col:
658 { '$operator': '$col} }
659 ...
660 }
661 """
662
663 def inner(self, columns=None):
664 columns = make_tuple(columns or self.columns)
665 mapping = {col: {op: '$' + col} for col in columns}
666 self.project(mapping)
667 if unwind:
668 self.stages.append(
669 {
670 '$unwind': {''},
671 },
672 )
673 return self
674
675 inner.__doc__ = op.replace('$', '')
676
677 return inner
678
679 def isequal(self, other):
680 self.strcasecmp(other)
681 # strcasecmp returns 0 for equality, 1 and -1 for greater/less than
682 # https://docs.mongodb.com/manual/reference/operator/aggregation/strcasecmp/
683 mapping = {
684 col: {
685 '$cond': {
686 'if': {'$eq': ['$' + col, 0]},
687 'then': True,
688 'else': False,
689 },
690 }
691 for col in self.columns
692 }
693 self.project(mapping)
694
695 concat = __strexpr__('$concat', base=True)
696 split = __strexpr__('$split', unwind=True, base=True, max_terms=2)
697 usplit = __strexpr__('$split', unwind=False, base=True, max_terms=2)
698 upper = __strunary__('$toUpper')
699 lower = __strunary__('$toLower')
700 substr = __strexpr__('$substr', base=True)
701 strcasecmp = __strexpr__('$strcasecmp', base=True)
702 len = __strunary__('$strLenBytes')
703 index = __strexpr__('$indexOfBytes', base=True)
704
705
[docs]
706class ApplyAccumulators:
707 def agg(self, map=None, **kwargs):
708 stage = self._getGroupBy(by='$$last')
709 specs = map or kwargs
710 for col, colExpr in specs.items():
711 if isinstance(colExpr, dict):
712 # specify an arbitrary expression
713 groupby = stage['$group']
714 groupby[col] = colExpr
715 elif isinstance(colExpr, str):
716 # specify some known operator
717 if hasattr(self, colExpr):
718 method = getattr(self, colExpr)
719 method(col)
720 else:
721 raise SyntaxError(f'{colExpr} is not known')
722 elif isinstance(colExpr, (tuple, list)):
723 # specify a list of some known operators
724 for statExpr in colExpr:
725 if hasattr(self, statExpr):
726 method = getattr(self, statExpr)
727 method(col)
728 else:
729 raise SyntaxError(f'{statExpr} is not known')
730 elif callable(colExpr):
731 # specify a callable that returns an expression
732 groupby = stage['$group']
733 groupby[col] = colExpr(col)
734 else:
735 SyntaxError(f'{colExpr} on column {col} is unknown or invalid')
736 return self
737
738 def __statop__(op, opname=None):
739 opname = opname or op.replace('$', '')
740
741 def inner(self, columns=None):
742 columns = make_tuple(columns or self.columns)
743 stage = self._getGroupBy(by='$$last')
744 groupby = stage['$group']
745 groupby.update(
746 {
747 f'{col}_{opname}': {
748 op: '$' + col,
749 }
750 for col in columns
751 },
752 )
753 self.computed.extend(groupby.keys())
754 self.project_keeper_columns()
755 return self
756
757 return inner
758
759 sum = __statop__('$sum')
760 avg = __statop__('$avg')
761 mean = __statop__('$avg')
762 min = __statop__('$min')
763 max = __statop__('$max')
764 std = __statop__('$stdDevSamp', 'std')
765
766
767class ApplyCache:
768 """
769 A Cache that works on collections and pipelines
770 """
771
772 def __init__(self, db_alias):
773 self._db_alias = db_alias
774
775 def set(self, key, value):
776 # https://stackoverflow.com/a/22003440/890242
777 QueryCache = make_QueryCache(self._db_alias)
778 QueryCache.objects(key=key).update_one(
779 set__key=f"{key}",
780 set__value=value,
781 upsert=True,
782 )
783
784 def get(self, key):
785 QueryCache = make_QueryCache(self._db_alias)
786 try:
787 result = QueryCache.objects.get(key=key)
788 except:
789 result = None
790 return result
791
792
793class ApplyStatistics:
794 def quantile(self, q=0.5):
795 def preparefn(val):
796 return val.pivot(columns='var', index='percentile', values='value')
797
798 return self.apply(self._percentile(q), preparefn=preparefn)
799
800 def cov(self):
801 def preparefn(val):
802 val = val.pivot(columns='y', index='x', values='cov')
803 val.index.name = None
804 val.columns.name = None
805 return val
806
807 return self.apply(self._covariance, preparefn=preparefn)
808
809 def corr(self):
810 def preparefn(val):
811 val = val.pivot(columns='y', index='x', values='rho')
812 val.index.name = None
813 val.columns.name = None
814 return val
815
816 return self.apply(self._pearson, preparefn=preparefn)
817
818 def _covariance(self, ctx):
819 # this works
820 # source http://ci.columbia.edu/ci/premba_test/c0331/s7/s7_5.html
821 facets = {}
822 means = {}
823 unwinds = []
824 count = len(ctx.caller.noapply()) - 1
825 for x, y in product(ctx.columns, ctx.columns):
826 xcol = '$' + x
827 ycol = '$' + y
828 # only calculate the same column's mean once
829 if xcol not in means:
830 means[xcol] = ctx.caller[x].noapply().mean().values[0, 0]
831 if ycol not in means:
832 means[ycol] = ctx.caller[y].noapply().mean().values[0, 0]
833 sumands = {
834 xcol: {
835 '$subtract': [xcol, means[xcol]],
836 },
837 ycol: {'$subtract': [ycol, means[ycol]]},
838 }
839 multiply = {
840 '$multiply': [sumands[xcol], sumands[ycol]],
841 }
842 agg = {
843 '$group': {
844 '_id': None,
845 'value': {
846 '$sum': multiply,
847 },
848 }
849 }
850 project = {
851 '$project': {
852 'cov': {'$divide': ['$value', count]},
853 'x': x,
854 'y': y,
855 }
856 }
857 pipeline = [agg, project]
858 outcol = f'{x}_{y}'
859 facets[outcol] = pipeline
860 unwinds.append(
861 {'$unwind': '$' + outcol},
862 )
863 facet = {'$facet': facets}
864 expand = [
865 {
866 '$project': {
867 'value': {'$objectToArray': '$$CURRENT'},
868 }
869 },
870 {'$unwind': '$value'},
871 {
872 '$replaceRoot': {'newRoot': '$value.v'},
873 },
874 ]
875 return [facet, *unwinds, *expand]
876
877 def _pearson(self, ctx):
878 # this works
879 # source http://ilearnasigoalong.blogspot.ch/2017/10/calculating-correlation-inside-mongodb.html
880 facets = {}
881 unwinds = []
882 for x, y in product(ctx.columns, ctx.columns):
883 xcol = '$' + x
884 ycol = '$' + y
885 sumcolumns = {
886 '$group': {
887 '_id': None,
888 'count': {'$sum': 1},
889 'sumx': {'$sum': xcol},
890 'sumy': {'$sum': ycol},
891 'sumxsquared': {'$sum': {'$multiply': [xcol, xcol]}},
892 'sumysquared': {'$sum': {'$multiply': [ycol, ycol]}},
893 'sumxy': {'$sum': {'$multiply': [xcol, ycol]}},
894 }
895 }
896
897 multiply_sumx_sumy = {
898 '$multiply': ["$sumx", "$sumy"],
899 }
900 multiply_sumxy_count = {
901 '$multiply': ["$sumxy", "$count"],
902 }
903 partone = {
904 '$subtract': [multiply_sumxy_count, multiply_sumx_sumy],
905 }
906
907 multiply_sumxsquared_count = {
908 '$multiply': ["$sumxsquared", "$count"],
909 }
910 sumx_squared = {
911 '$multiply': ["$sumx", "$sumx"],
912 }
913 subparttwo = {
914 '$subtract': [multiply_sumxsquared_count, sumx_squared],
915 }
916
917 multiply_sumysquared_count = {
918 '$multiply': ["$sumysquared", "$count"],
919 }
920 sumy_squared = {
921 '$multiply': ["$sumy", "$sumy"],
922 }
923 subpartthree = {
924 '$subtract': [multiply_sumysquared_count, sumy_squared],
925 }
926
927 parttwo = {
928 '$sqrt': {
929 '$multiply': [subparttwo, subpartthree],
930 },
931 }
932
933 rho = {
934 '$project': {
935 'rho': {
936 '$divide': [partone, parttwo],
937 },
938 'x': x,
939 'y': y,
940 }
941 }
942 pipeline = [sumcolumns, rho]
943 outcol = f'{x}_{y}'
944 facets[outcol] = pipeline
945 unwinds.append(
946 {'$unwind': '$' + outcol},
947 )
948 facet = {'$facet': facets}
949 expand = [
950 {
951 '$project': {
952 'value': {'$objectToArray': '$$CURRENT'},
953 }
954 },
955 {'$unwind': '$value'},
956 {
957 '$replaceRoot': {'newRoot': '$value.v'},
958 },
959 ]
960 return [facet, *unwinds, *expand]
961
962 def _percentile(self, pctls=None):
963 """
964 calculate percentiles for all columns
965 """
966 pctls = pctls or [0.25, 0.5, 0.75]
967 if not isinstance(pctls, (list, tuple)):
968 pctls = [pctls]
969
970 def calc(col, p, outcol):
971 # sort values
972 sort = {
973 '$sort': {col: 1},
974 }
975 # group/push to get an array of all values
976 group = {
977 '$group': {
978 '_id': col,
979 'values': {
980 '$push': "$" + col,
981 },
982 },
983 }
984 # find value at requested percentile
985 perc = {
986 '$arrayElemAt': [
987 '$values',
988 {
989 '$floor': {
990 '$multiply': [
991 {
992 '$size': '$values',
993 },
994 p,
995 ]
996 }
997 },
998 ],
999 }
1000 # map percentile value to output column
1001 project = {
1002 '$project': {
1003 'var': col,
1004 'percentile': f'p{p}',
1005 'value': perc,
1006 },
1007 }
1008 return [sort, group, project]
1009
1010 def inner(ctx):
1011 # for each column and requested percentile, build a pipeline
1012 # all pipelines will be combined into a $facet stage to
1013 # calculate every column/percentile tuple in parallel
1014 facets = {}
1015 unwind = []
1016 # for each column build a pipeline to calculate the percentiles
1017 for col in ctx.columns:
1018 for p in pctls:
1019 # e.g. outcol for perc .25 of column abc => abcp25
1020 outcol = f'{col}_p{p}'.replace('0.', '')
1021 facets[outcol] = calc(col, p, outcol)
1022 unwind.append(
1023 {'$unwind': '$' + outcol},
1024 )
1025 # process per-column pipelines in parallel, resulting in one
1026 # document for each variable + percentile combination
1027 facet = {'$facet': facets}
1028 # expand single document into one document per variable + percentile combo
1029 # the resulting set of documents contains var/percentile/value
1030 expand = [
1031 {
1032 '$project': {
1033 'value': {
1034 '$objectToArray': '$$CURRENT',
1035 },
1036 },
1037 },
1038 {'$unwind': '$value'},
1039 {
1040 '$replaceRoot': {
1041 'newRoot': '$value.v',
1042 },
1043 },
1044 ]
1045 pipeline = [facet, *unwind, *expand]
1046 return pipeline
1047
1048 return inner