PERF: Cythonize Groupby.cummin/cummax (pandas-dev#15048)

Author: mroeschke

asv_bench/benchmarks/groupby.py

@@ -344,6 +344,26 @@ def time_groupby_dt_timegrouper_size(self):
         self.df.groupby(TimeGrouper(key='dates', freq='M')).size()
 
 
+#----------------------------------------------------------------------
+# groupby.cummin and groupby.cummax # GH 15048
+
+class groupby_cummax_cummin(object):
+    goal_time = 0.2
+
+    def setup(self):
+        np.random.seed(1234)
+        self.G = 1000
+        self.N = 10000
+        self.df = pd.DataFrame({'A': np.random.randint(0, self.G, size=self.N),
+                                'B': np.random.randn(self.N)})
+
+    def time_cummin(self):
+        self.df.groupby('A').cummin()
+
+    def time_cummax(self):
+        self.df.groupby('A').cummax()
+
+
 #----------------------------------------------------------------------
 # groupby with a variable value for ngroups
 
@@ -690,6 +710,3 @@ def time_shift(self):
     def time_transform_dataframe(self):
         # GH 12737
         self.df_nans.groupby('key').transform('first')
-
-
-

pandas/core/groupby.py

@@ -75,7 +75,7 @@
     'last', 'first',
     'head', 'tail', 'median',
     'mean', 'sum', 'min', 'max',
-    'cumsum', 'cumprod', 'cummin', 'cummax', 'cumcount',
+    'cumsum', 'cumprod', 'cumcount',
     'resample',
     'describe',
     'rank', 'quantile',
@@ -97,7 +97,8 @@
 _dataframe_apply_whitelist = \
     _common_apply_whitelist | frozenset(['dtypes', 'corrwith'])
 
-_cython_transforms = frozenset(['cumprod', 'cumsum', 'shift'])
+_cython_transforms = frozenset(['cumprod', 'cumsum', 'shift',
+                                'cummin', 'cummax'])
 
 
 def _groupby_function(name, alias, npfunc, numeric_only=True,
@@ -1415,6 +1416,26 @@ def cumsum(self, axis=0, *args, **kwargs):
 
         return self._cython_transform('cumsum')
 
+    @Substitution(name='groupby')
+    @Appender(_doc_template)
+    def cummin(self, axis=0, *args, **kwargs):
+        """Cumulative min for each group"""
+        nv.validate_groupby_func('cummin', args, kwargs)
+        if axis != 0:
+            return self.apply(lambda x: np.minimum.accumulate(x, axis))
+
+        return self._cython_transform('cummin')
+
+    @Substitution(name='groupby')
+    @Appender(_doc_template)
+    def cummax(self, axis=0, *args, **kwargs):
+        """Cumulative max for each group"""
+        nv.validate_groupby_func('cummax', args, kwargs)
+        if axis != 0:
+            return self.apply(lambda x: np.maximum.accumulate(x, axis))
+
+        return self._cython_transform('cummax')
+
     @Substitution(name='groupby')
     @Appender(_doc_template)
     def shift(self, periods=1, freq=None, axis=0):
@@ -1752,6 +1773,8 @@ def get_group_levels(self):
         'transform': {
             'cumprod': 'group_cumprod',
             'cumsum': 'group_cumsum',
+            'cummin': 'group_cummin',
+            'cummax': 'group_cummax',
         }
     }
 

pandas/src/algos_groupby_helper.pxi.in

@@ -667,6 +667,73 @@ def group_cumsum(numeric[:, :] out,
                     out[i, j] = accum[lab, j]
 
 
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def group_cummin(numeric[:, :] out,
+                 numeric[:, :] values,
+                 int64_t[:] labels,
+                 numeric[:, :] accum):
+    """
+    Only transforms on axis=0
+    """
+    cdef:
+        Py_ssize_t i, j, N, K, size
+        numeric val, min_val
+        int64_t lab
+
+    N, K = (<object> values).shape
+    accum = np.zeros_like(accum)
+
+    with nogil:
+        for i in range(N):
+            lab = labels[i]
+
+            if lab < 0:
+                continue
+            for j in range(K):
+                val = values[i, j]
+                if val == val:
+                    if j == 0:
+                        min_val = val
+                    elif val < min_val:
+                        min_val = val
+                    accum[lab, j] = min_val
+                    out[i, j] = accum[lab, j]
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def group_cummax(numeric[:, :] out,
+                 numeric[:, :] values,
+                 int64_t[:] labels,
+                 numeric[:, :] accum):
+    """
+    Only transforms on axis=0
+    """
+    cdef:
+        Py_ssize_t i, j, N, K, size
+        numeric val, max_val
+        int64_t lab
+
+    N, K = (<object> values).shape
+    accum = np.zeros_like(accum)
+
+    with nogil:
+        for i in range(N):
+            lab = labels[i]
+
+            if lab < 0:
+                continue
+            for j in range(K):
+                val = values[i, j]
+                if val == val:
+                    if j == 0:
+                        max_val = val
+                    elif val > max_val:
+                        max_val = val
+                    accum[lab, j] = max_val
+                    out[i, j] = accum[lab, j]
+
+
 @cython.boundscheck(False)
 @cython.wraparound(False)
 def group_shift_indexer(int64_t[:] out, int64_t[:] labels,

pandas/tests/groupby/test_groupby.py

@@ -4979,8 +4979,6 @@ def test_groupby_whitelist(self):
             'tail',
             'cumsum',
             'cumprod',
-            'cummin',
-            'cummax',
             'cumcount',
             'resample',
             'describe',
@@ -5020,8 +5018,6 @@ def test_groupby_whitelist(self):
             'tail',
             'cumsum',
             'cumprod',
-            'cummin',
-            'cummax',
             'cumcount',
             'resample',
             'describe',