BUG: errors and segfaults in groupby cython transforms (#16771)

Author: abull

doc/source/whatsnew/v0.25.0.rst

@@ -387,7 +387,8 @@ Groupby/Resample/Rolling
 - Ensured that ordering of outputs in ``groupby`` aggregation functions is consistent across all versions of Python (:issue:`25692`)
 - Ensured that result group order is correct when grouping on an ordered ``Categorical`` and specifying ``observed=True`` (:issue:`25871`, :issue:`25167`)
 - Bug in :meth:`pandas.core.window.Rolling.min` and :meth:`pandas.core.window.Rolling.max` that caused a memory leak (:issue:`25893`)
-- Bug in :func:`idxmax` and :func:`idxmin` on :meth:`DataFrame.groupby` with datetime column would return incorrect dtype (:issue:`25444`, :issue:`15306`)
+- Bug in :meth:`pandas.core.groupby.GroupBy.idxmax` and :meth:`pandas.core.groupby.GroupBy.idxmin` with datetime column would return incorrect dtype (:issue:`25444`, :issue:`15306`)
+- Bug in :meth:`pandas.core.groupby.GroupBy.cumsum`, :meth:`pandas.core.groupby.GroupBy.cumprod`, :meth:`pandas.core.groupby.GroupBy.cummin` and :meth:`pandas.core.groupby.GroupBy.cummax` with categorical column having absent categories, would return incorrect result or segfault (:issue:`16771`)
 
 Reshaping
 ^^^^^^^^^

pandas/_libs/groupby.pyx

@@ -142,19 +142,29 @@ def group_median_float64(ndarray[float64_t, ndim=2] out,
 def group_cumprod_float64(float64_t[:, :] out,
                           const float64_t[:, :] values,
                           const int64_t[:] labels,
+                          int ngroups,
                           bint is_datetimelike,
                           bint skipna=True):
+    """Cumulative product of columns of `values`, in row groups `labels`.
+
+Arguments:
+    `out`: array to store cumprod in
+    `values`: values to take cumprod of
+    `labels`: labels to group by
+    `ngroups`: number of groups, larger than all entries of `labels`
+    `is_datetimelike`: always false, `values` is never datetime-like
+    `skipna`: if true, ignore nans in `values`
+
     """
-    Only transforms on axis=0
-    """
+
     cdef:
         Py_ssize_t i, j, N, K, size
         float64_t val
         float64_t[:, :] accum
         int64_t lab
 
     N, K = (<object>values).shape
-    accum = np.ones_like(values)
+    accum = np.ones((ngroups, K), np.float64)
 
     with nogil:
         for i in range(N):
@@ -179,19 +189,29 @@ def group_cumprod_float64(float64_t[:, :] out,
 def group_cumsum(numeric[:, :] out,
                  numeric[:, :] values,
                  const int64_t[:] labels,
+                 int ngroups,
                  is_datetimelike,
                  bint skipna=True):
+    """Cumulative sum of columns of `values`, in row groups `labels`.
+
+Arguments:
+    `out`: array to store cumsum in
+    `values`: values to take cumsum of
+    `labels`: labels to group by
+    `ngroups`: number of groups, larger than all entries of `labels`
+    `is_datetimelike`: true if `values` contains datetime-like entries
+    `skipna`: if true, ignore nans in `values`
+
     """
-    Only transforms on axis=0
-    """
+
     cdef:
         Py_ssize_t i, j, N, K, size
         numeric val
         numeric[:, :] accum
         int64_t lab
 
     N, K = (<object>values).shape
-    accum = np.zeros_like(values)
+    accum = np.zeros((ngroups, K), np.asarray(values).dtype)
 
     with nogil:
         for i in range(N):

pandas/_libs/groupby_helper.pxi.in

@@ -474,18 +474,27 @@ def group_min(groupby_t[:, :] out,
 def group_cummin(groupby_t[:, :] out,
                  groupby_t[:, :] values,
                  const int64_t[:] labels,
+                 int ngroups,
                  bint is_datetimelike):
+    """Cumulative minimum of columns of `values`, in row groups `labels`.
+
+Arguments:
+    `out`: array to store cummin in
+    `values`: values to take cummin of
+    `labels`: labels to group by
+    `ngroups`: number of groups, larger than all entries of `labels`
+    `is_datetimelike`: true if `values` contains datetime-like entries
+
     """
-    Only transforms on axis=0
-    """
+
     cdef:
         Py_ssize_t i, j, N, K, size
         groupby_t val, mval
         ndarray[groupby_t, ndim=2] accum
         int64_t lab
 
     N, K = (<object>values).shape
-    accum = np.empty_like(values)
+    accum = np.empty((ngroups, K), np.asarray(values).dtype)
     if groupby_t is int64_t:
         accum[:] = _int64_max
     else:
@@ -522,18 +531,27 @@ def group_cummin(groupby_t[:, :] out,
 def group_cummax(groupby_t[:, :] out,
                  groupby_t[:, :] values,
                  const int64_t[:] labels,
+		 int ngroups,
                  bint is_datetimelike):
+    """Cumulative maximum of columns of `values`, in row groups `labels`.
+
+Arguments:
+    `out`: array to store cummax in
+    `values`: values to take cummax of
+    `labels`: labels to group by
+    `ngroups`: number of groups, larger than all entries of `labels`
+    `is_datetimelike`: true if `values` contains datetime-like entries
+
     """
-    Only transforms on axis=0
-    """
+
     cdef:
         Py_ssize_t i, j, N, K, size
         groupby_t val, mval
         ndarray[groupby_t, ndim=2] accum
         int64_t lab
 
     N, K = (<object>values).shape
-    accum = np.empty_like(values)
+    accum = np.empty((ngroups, K), np.asarray(values).dtype)
     if groupby_t is int64_t:
         accum[:] = -_int64_max
     else:

pandas/core/groupby/ops.py

@@ -361,8 +361,8 @@ def get_group_levels(self):
             'cummax': 'group_cummax',
             'rank': {
                 'name': 'group_rank',
-                'f': lambda func, a, b, c, d, **kwargs: func(
-                    a, b, c, d,
+                'f': lambda func, a, b, c, d, e, **kwargs: func(
+                    a, b, c, e,
                     kwargs.get('ties_method', 'average'),
                     kwargs.get('ascending', True),
                     kwargs.get('pct', False),
@@ -600,9 +600,10 @@ def _transform(self, result, values, comp_ids, transform_func,
             for i, chunk in enumerate(values.transpose(2, 0, 1)):
 
                 transform_func(result[:, :, i], values,
-                               comp_ids, is_datetimelike, **kwargs)
+                               comp_ids, ngroups, is_datetimelike, **kwargs)
         else:
-            transform_func(result, values, comp_ids, is_datetimelike, **kwargs)
+            transform_func(result, values, comp_ids, ngroups, is_datetimelike,
+                           **kwargs)
 
         return result
 

pandas/tests/groupby/test_transform.py

@@ -9,7 +9,8 @@
 from pandas.core.dtypes.common import ensure_platform_int, is_timedelta64_dtype
 
 import pandas as pd
-from pandas import DataFrame, MultiIndex, Series, Timestamp, concat, date_range
+from pandas import (
+    Categorical, DataFrame, MultiIndex, Series, Timestamp, concat, date_range)
 from pandas.core.groupby.groupby import DataError
 from pandas.util import testing as tm
 from pandas.util.testing import assert_frame_equal, assert_series_equal
@@ -470,7 +471,8 @@ def _check_cython_group_transform_cumulative(pd_op, np_op, dtype):
     ans = np.zeros_like(data)
 
     labels = np.array([0, 0, 0, 0], dtype=np.int64)
-    pd_op(ans, data, labels, is_datetimelike)
+    ngroups = 1
+    pd_op(ans, data, labels, ngroups, is_datetimelike)
 
     tm.assert_numpy_array_equal(np_op(data), ans[:, 0],
                                 check_dtype=False)
@@ -496,17 +498,19 @@ def test_cython_group_transform_algos():
 
     # with nans
     labels = np.array([0, 0, 0, 0, 0], dtype=np.int64)
+    ngroups = 1
 
     data = np.array([[1], [2], [3], [np.nan], [4]], dtype='float64')
     actual = np.zeros_like(data)
     actual.fill(np.nan)
-    groupby.group_cumprod_float64(actual, data, labels, is_datetimelike)
+    groupby.group_cumprod_float64(actual, data, labels, ngroups,
+                                  is_datetimelike)
     expected = np.array([1, 2, 6, np.nan, 24], dtype='float64')
     tm.assert_numpy_array_equal(actual[:, 0], expected)
 
     actual = np.zeros_like(data)
     actual.fill(np.nan)
-    groupby.group_cumsum(actual, data, labels, is_datetimelike)
+    groupby.group_cumsum(actual, data, labels, ngroups, is_datetimelike)
     expected = np.array([1, 3, 6, np.nan, 10], dtype='float64')
     tm.assert_numpy_array_equal(actual[:, 0], expected)
 
@@ -515,7 +519,7 @@ def test_cython_group_transform_algos():
     data = np.array([np.timedelta64(1, 'ns')] * 5, dtype='m8[ns]')[:, None]
     actual = np.zeros_like(data, dtype='int64')
     groupby.group_cumsum(actual, data.view('int64'), labels,
-                         is_datetimelike)
+                         ngroups, is_datetimelike)
     expected = np.array([np.timedelta64(1, 'ns'), np.timedelta64(
         2, 'ns'), np.timedelta64(3, 'ns'), np.timedelta64(4, 'ns'),
         np.timedelta64(5, 'ns')])
@@ -863,3 +867,15 @@ def test_groupby_transform_with_datetimes(func, values):
                          index=dates, name="price")
 
     tm.assert_series_equal(result, expected)
+
+
+def test_transform_absent_categories():
+    # GH 16771
+    # cython transforms with more groups than rows
+    x_vals = [1]
+    x_cats = range(2)
+    y = [1]
+    df = DataFrame(dict(x=Categorical(x_vals, x_cats), y=y))
+    result = df.y.groupby(df.x).cumsum()
+    expected = df.y
+    assert_series_equal(result, expected)