Moved function groupby_result_dtype to _GroupBy._result_dtype

 - Reverted maybe_downcast_numeric to its original state
 - Parameterized tests

Author: rhshadrach

pandas/core/dtypes/cast.py

@@ -30,7 +30,6 @@
     ensure_int64,
     ensure_object,
     ensure_str,
-    groupby_result_dtype,
     is_bool,
     is_bool_dtype,
     is_complex,
@@ -173,9 +172,7 @@ def maybe_downcast_to_dtype(result, dtype):
     return result
 
 
-def maybe_downcast_numeric(
-    result, dtype, do_round: bool = False, how: str = "",
-):
+def maybe_downcast_numeric(result, dtype, do_round: bool = False):
     """
     Subset of maybe_downcast_to_dtype restricted to numeric dtypes.
 
@@ -184,7 +181,6 @@ def maybe_downcast_numeric(
     result : ndarray or ExtensionArray
     dtype : np.dtype or ExtensionDtype
     do_round : bool
-    how : str
 
     Returns
     -------
@@ -199,8 +195,6 @@ def maybe_downcast_numeric(
         #  earlier
         result = np.array(result)
 
-    dtype = groupby_result_dtype(dtype, how)
-
     def trans(x):
         if do_round:
             return x.round()

pandas/core/dtypes/common.py

@@ -1786,27 +1786,3 @@ def pandas_dtype(dtype) -> DtypeObj:
         raise TypeError(f"dtype '{dtype}' not understood")
 
     return npdtype
-
-
-def groupby_result_dtype(dtype, how) -> DtypeObj:
-    """
-    Get the desired dtype of an aggregation result based on the
-    input dtype and how the aggregation is done.
-
-    Parameters
-    ----------
-    dtype : dtype, type
-        The input dtype for the groupby.
-    how : str
-        How the aggregation is performed.
-
-    Returns
-    -------
-    The desired dtype of the aggregation result.
-    """
-    d = {
-        (np.dtype(np.bool), "add"): np.dtype(np.int64),
-        (np.dtype(np.bool), "cumsum"): np.dtype(np.int64),
-        (np.dtype(np.bool), "sum"): np.dtype(np.int64),
-    }
-    return d.get((dtype, how), dtype)

pandas/core/groupby/generic.py

@@ -1072,8 +1072,10 @@ def _cython_agg_blocks(
             assert not isinstance(result, DataFrame)
 
             if result is not no_result:
-                # see if we can cast the block back to the original dtype
-                result = maybe_downcast_numeric(result, block.dtype, how=how)
+                # see if we can cast the block to the desired dtype
+                # this may not be the original dtype
+                dtype = self._result_dtype(block.dtype, how)
+                result = maybe_downcast_numeric(result, dtype)
 
                 if block.is_extension and isinstance(result, np.ndarray):
                     # e.g. block.values was an IntegerArray

pandas/core/groupby/groupby.py

@@ -33,7 +33,7 @@ class providing the base-class of operations.
 
 from pandas._libs import Timestamp
 import pandas._libs.groupby as libgroupby
-from pandas._typing import FrameOrSeries, Scalar
+from pandas._typing import DtypeObj, FrameOrSeries, Scalar
 from pandas.compat import set_function_name
 from pandas.compat.numpy import function as nv
 from pandas.errors import AbstractMethodError
@@ -42,7 +42,6 @@ class providing the base-class of operations.
 from pandas.core.dtypes.cast import maybe_downcast_to_dtype
 from pandas.core.dtypes.common import (
     ensure_float,
-    groupby_result_dtype,
     is_datetime64_dtype,
     is_extension_array_dtype,
     is_integer_dtype,
@@ -795,7 +794,7 @@ def _cumcount_array(self, ascending: bool = True):
 
     def _try_cast(self, result, obj, numeric_only: bool = False, how: str = ""):
         """
-        Try to cast the result to our obj original type,
+        Try to cast the result to the desired type,
         we may have roundtripped through object in the mean-time.
 
         If numeric_only is True, then only try to cast numerics
@@ -806,8 +805,7 @@ def _try_cast(self, result, obj, numeric_only: bool = False, how: str = ""):
             dtype = obj._values.dtype
         else:
             dtype = obj.dtype
-
-        dtype = groupby_result_dtype(dtype, how)
+        dtype = self._result_dtype(dtype, how)
 
         if not is_scalar(result):
             if is_extension_array_dtype(dtype) and dtype.kind != "M":
@@ -1028,6 +1026,30 @@ def _apply_filter(self, indices, dropna):
             filtered = self._selected_obj.where(mask)  # Fill with NaNs.
         return filtered
 
+    @staticmethod
+    def _result_dtype(dtype, how) -> DtypeObj:
+        """
+        Get the desired dtype of a groupby result based on the
+        input dtype and how the aggregation is done.
+
+        Parameters
+        ----------
+        dtype : dtype, type
+            The input dtype of the groupby.
+        how : str
+            How the aggregation is performed.
+
+        Returns
+        -------
+        The desired dtype of the aggregation result.
+        """
+        d = {
+            (np.dtype(np.bool), "add"): np.dtype(np.int64),
+            (np.dtype(np.bool), "cumsum"): np.dtype(np.int64),
+            (np.dtype(np.bool), "sum"): np.dtype(np.int64),
+        }
+        return d.get((dtype, how), dtype)
+
 
 class GroupBy(_GroupBy):
     """

pandas/tests/groupby/test_groupby.py

@@ -2061,44 +2061,25 @@ def test_groups_repr_truncates(max_seq_items, expected):
         assert result == expected
 
 
-def test_bool_agg_dtype():
+@pytest.mark.parametrize(
+    "op",
+    [
+        lambda x: x.sum(),
+        lambda x: x.cumsum(),
+        lambda x: x.transform("sum"),
+        lambda x: x.transform("cumsum"),
+        lambda x: x.agg("sum"),
+        lambda x: x.agg("cumsum"),
+    ],
+)
+def test_bool_agg_dtype(op):
     # GH 7001
-    # Bool aggregation results in int
+    # Bool sum aggregations result in int
     df = pd.DataFrame({"a": [1, 1], "b": [False, True]})
     s = df.set_index("a")["b"]
 
-    result = df.groupby("a").sum()["b"].dtype
-    assert is_integer_dtype(result)
-
-    result = s.groupby("a").sum().dtype
-    assert is_integer_dtype(result)
-
-    result = df.groupby("a").cumsum()["b"].dtype
-    assert is_integer_dtype(result)
-
-    result = s.groupby("a").cumsum().dtype
-    assert is_integer_dtype(result)
-
-    result = df.groupby("a").agg("sum")["b"].dtype
-    assert is_integer_dtype(result)
-
-    result = s.groupby("a").agg("sum").dtype
-    assert is_integer_dtype(result)
-
-    result = df.groupby("a").agg("cumsum")["b"].dtype
-    assert is_integer_dtype(result)
-
-    result = s.groupby("a").agg("cumsum").dtype
-    assert is_integer_dtype(result)
-
-    result = df.groupby("a").transform("sum")["b"].dtype
-    assert is_integer_dtype(result)
-
-    result = s.groupby("a").transform("sum").dtype
-    assert is_integer_dtype(result)
-
-    result = df.groupby("a").transform("cumsum")["b"].dtype
+    result = op(df.groupby("a"))["b"].dtype
     assert is_integer_dtype(result)
 
-    result = s.groupby("a").transform("cumsum").dtype
+    result = op(s.groupby("a")).dtype
     assert is_integer_dtype(result)