ENH: Support mask in GroupBy.cumsum

doc/source/whatsnew/v1.5.0.rst

@@ -1070,8 +1070,9 @@ Groupby/resample/rolling
 - Bug when using ``engine="numba"`` would return the same jitted function when modifying ``engine_kwargs`` (:issue:`46086`)
 - Bug in :meth:`.DataFrameGroupBy.transform` fails when ``axis=1`` and ``func`` is ``"first"`` or ``"last"`` (:issue:`45986`)
 - Bug in :meth:`DataFrameGroupBy.cumsum` with ``skipna=False`` giving incorrect results (:issue:`46216`)
-- Bug in :meth:`GroupBy.sum` with integer dtypes losing precision (:issue:`37493`)
+- Bug in :meth:`GroupBy.sum` and :meth:`GroupBy.cumsum` with integer dtypes losing precision (:issue:`37493`)
 - Bug in :meth:`.GroupBy.cumsum` with ``timedelta64[ns]`` dtype failing to recognize ``NaT`` as a null value (:issue:`46216`)
+- Bug in :meth:`GroupBy.cumsum` with integer dtypes causing overflows when sum was bigger than maximum of dtype (:issue:`37493`)
 - Bug in :meth:`.GroupBy.cummin` and :meth:`.GroupBy.cummax` with nullable dtypes incorrectly altering the original data in place (:issue:`46220`)
 - Bug in :meth:`DataFrame.groupby` raising error when ``None`` is in first level of :class:`MultiIndex` (:issue:`47348`)
 - Bug in :meth:`.GroupBy.cummax` with ``int64`` dtype with leading value being the smallest possible int64 (:issue:`46382`)

pandas/_libs/groupby.pyi

@@ -20,12 +20,14 @@ def group_cumprod_float64(
     skipna: bool = ...,
 ) -> None: ...
 def group_cumsum(
-    out: np.ndarray,  # numeric[:, ::1]
-    values: np.ndarray,  # ndarray[numeric, ndim=2]
+    out: np.ndarray,  # cumsum_t[:, ::1]
+    values: np.ndarray,  # ndarray[cumsum_t, ndim=2]
     labels: np.ndarray,  # const int64_t[:]
     ngroups: int,
     is_datetimelike: bool,
     skipna: bool = ...,
+    mask: np.ndarray | None = ...,
+    result_mask: np.ndarray | None = ...,
 ) -> None: ...
 def group_shift_indexer(
     out: np.ndarray,  # int64_t[::1]

pandas/_libs/groupby.pyx

@@ -207,15 +207,24 @@ def group_cumprod_float64(
                         break
 
 
+ctypedef fused cumsum_t:
+    int64_t
+    uint64_t
+    float32_t
+    float64_t
+
+
 @cython.boundscheck(False)
 @cython.wraparound(False)
 def group_cumsum(
-    numeric_t[:, ::1] out,
-    ndarray[numeric_t, ndim=2] values,
+    cumsum_t[:, ::1] out,
+    ndarray[cumsum_t, ndim=2] values,
     const intp_t[::1] labels,
     int ngroups,
     bint is_datetimelike,
     bint skipna=True,
+    const uint8_t[:, :] mask=None,
+    uint8_t[:, ::1] result_mask=None,
 ) -> None:
     """
     Cumulative sum of columns of `values`, in row groups `labels`.
@@ -234,23 +243,33 @@ def group_cumsum(
         True if `values` contains datetime-like entries.
     skipna : bool
         If true, ignore nans in `values`.
+    mask: np.ndarray[uint8], optional
+        Mask of values
+    result_mask: np.ndarray[int8], optional
+        Mask of out array
 
     Notes
     -----
     This method modifies the `out` parameter, rather than returning an object.
     """
     cdef:
         Py_ssize_t i, j, N, K, size
-        numeric_t val, y, t, na_val
-        numeric_t[:, ::1] accum, compensation
+        cumsum_t val, y, t, na_val
+        cumsum_t[:, ::1] accum, compensation
+        uint8_t[:, ::1] accum_mask
         intp_t lab
         bint isna_entry, isna_prev = False
+        bint uses_mask = mask is not None
 
     N, K = (<object>values).shape
+
+    if uses_mask:
+        accum_mask = np.zeros((ngroups, K), dtype="uint8")
+
     accum = np.zeros((ngroups, K), dtype=np.asarray(values).dtype)
     compensation = np.zeros((ngroups, K), dtype=np.asarray(values).dtype)
 
-    na_val = _get_na_val(<numeric_t>0, is_datetimelike)
+    na_val = _get_na_val(<cumsum_t>0, is_datetimelike)
 
     with nogil:
         for i in range(N):
@@ -261,23 +280,41 @@ def group_cumsum(
             for j in range(K):
                 val = values[i, j]
 
-                isna_entry = _treat_as_na(val, is_datetimelike)
+                if uses_mask:
+                    isna_entry = mask[i, j]
+                else:
+                    isna_entry = _treat_as_na(val, is_datetimelike)
 
                 if not skipna:
-                    isna_prev = _treat_as_na(accum[lab, j], is_datetimelike)
+                    if uses_mask:
+                        isna_prev = accum_mask[lab, j]
+                    else:
+                        isna_prev = _treat_as_na(accum[lab, j], is_datetimelike)
+
                     if isna_prev:
-                        out[i, j] = na_val
+                        if uses_mask:
+                            result_mask[i, j] = True
+                        else:
+                            out[i, j] = na_val
                         continue
 
                 if isna_entry:
-                    out[i, j] = na_val
+
+                    if uses_mask:
+                        result_mask[i, j] = True
+                    else:
+                        out[i, j] = na_val
+
                     if not skipna:
-                        accum[lab, j] = na_val
+                        if uses_mask:
+                            accum_mask[lab, j] = True
+                        else:
+                            accum[lab, j] = na_val
 
                 else:
                     # For floats, use Kahan summation to reduce floating-point
                     # error (https://en.wikipedia.org/wiki/Kahan_summation_algorithm)
-                    if numeric_t == float32_t or numeric_t == float64_t:
+                    if cumsum_t == float32_t or cumsum_t == float64_t:
                         y = val - compensation[lab, j]
                         t = accum[lab, j] + y
                         compensation[lab, j] = t - accum[lab, j] - y
@@ -1041,7 +1078,7 @@ cdef numeric_t _get_na_val(numeric_t val, bint is_datetimelike):
     elif numeric_t is int64_t and is_datetimelike:
         na_val = NPY_NAT
     else:
-        # Will not be used, but define to avoid uninitialized warning.
+        # Used in case of masks
         na_val = 0
     return na_val
 

pandas/core/groupby/ops.py

@@ -157,6 +157,7 @@ def __init__(self, kind: str, how: str, has_dropped_na: bool) -> None:
         "first",
         "rank",
         "sum",
+        "cumsum",
     }
 
     _cython_arity = {"ohlc": 4}  # OHLC
@@ -225,7 +226,7 @@ def _get_cython_vals(self, values: np.ndarray) -> np.ndarray:
                 # result may still include NaN, so we have to cast
                 values = ensure_float64(values)
 
-            elif how == "sum":
+            elif how in ["sum", "cumsum"]:
                 # Avoid overflow during group op
                 if values.dtype.kind == "i":
                     values = ensure_int64(values)

pandas/tests/groupby/test_groupby.py

@@ -2810,12 +2810,15 @@ def test_single_element_list_grouping():
         values, _ = next(iter(df.groupby(["a"])))
 
 
-def test_groupby_sum_avoid_casting_to_float():
+@pytest.mark.parametrize("func", ["sum", "cumsum"])
+def test_groupby_sum_avoid_casting_to_float(func):
     # GH#37493
     val = 922337203685477580
     df = DataFrame({"a": 1, "b": [val]})
-    result = df.groupby("a").sum() - val
+    result = getattr(df.groupby("a"), func)() - val
     expected = DataFrame({"b": [0]}, index=Index([1], name="a"))
+    if func == "cumsum":
+        expected = expected.reset_index(drop=True)
     tm.assert_frame_equal(result, expected)
 
 
@@ -2832,7 +2835,7 @@ def test_groupby_sum_support_mask(any_numeric_ea_dtype):
 
 
 @pytest.mark.parametrize("val, dtype", [(111, "int"), (222, "uint")])
-def test_groupby_sum_overflow(val, dtype):
+def test_groupby_overflow(val, dtype):
     # GH#37493
     df = DataFrame({"a": 1, "b": [val, val]}, dtype=f"{dtype}8")
     result = df.groupby("a").sum()
@@ -2842,3 +2845,19 @@ def test_groupby_sum_overflow(val, dtype):
         dtype=f"{dtype}64",
     )
     tm.assert_frame_equal(result, expected)
+
+    result = df.groupby("a").cumsum()
+    expected = DataFrame({"b": [val, val * 2]}, dtype=f"{dtype}64")
+    tm.assert_frame_equal(result, expected)
+
+
+@pytest.mark.parametrize("skipna, val", [(True, 3), (False, pd.NA)])
+def test_groupby_cumsum_mask(any_numeric_ea_dtype, skipna, val):
+    # GH#37493
+    df = DataFrame({"a": 1, "b": [1, pd.NA, 2]}, dtype=any_numeric_ea_dtype)
+    result = df.groupby("a").cumsum(skipna=skipna)
+    expected = DataFrame(
+        {"b": [1, pd.NA, val]},
+        dtype=any_numeric_ea_dtype,
+    )
+    tm.assert_frame_equal(result, expected)

pandas/tests/groupby/test_libgroupby.py

@@ -183,9 +183,10 @@ def _check_cython_group_transform_cumulative(pd_op, np_op, dtype):
     tm.assert_numpy_array_equal(np_op(data), answer[:, 0], check_dtype=False)
 
 
-def test_cython_group_transform_cumsum(any_real_numpy_dtype):
+@pytest.mark.parametrize("np_dtype", ["int64", "uint64", "float32", "float64"])
+def test_cython_group_transform_cumsum(np_dtype):
     # see gh-4095
-    dtype = np.dtype(any_real_numpy_dtype).type
+    dtype = np.dtype(np_dtype).type
     pd_op, np_op = group_cumsum, np.cumsum
     _check_cython_group_transform_cumulative(pd_op, np_op, dtype)