BUG: incorrect EA casting in groubpy.agg

doc/source/whatsnew/v1.3.0.rst

@@ -163,7 +163,7 @@ Plotting
 
 Groupby/resample/rolling
 ^^^^^^^^^^^^^^^^^^^^^^^^
-
+- Bug in :meth:`DataFrameGroupBy.agg` and :meth:`SeriesGroupBy.agg` with ``ExtensionDtype`` columns incorrectly casting results too aggressively (:issue:`38254`)
 -
 -
 

pandas/core/groupby/ops.py

@@ -28,11 +28,7 @@
 from pandas.errors import AbstractMethodError
 from pandas.util._decorators import cache_readonly
 
-from pandas.core.dtypes.cast import (
-    maybe_cast_result,
-    maybe_cast_result_dtype,
-    maybe_downcast_to_dtype,
-)
+from pandas.core.dtypes.cast import maybe_cast_result_dtype, maybe_downcast_to_dtype
 from pandas.core.dtypes.common import (
     ensure_float,
     ensure_float64,
@@ -750,7 +746,7 @@ def _aggregate_series_pure_python(self, obj: Series, func: F):
             result[label] = res
 
         result = lib.maybe_convert_objects(result, try_float=0)
-        result = maybe_cast_result(result, obj, numeric_only=True)
+        # TODO: cast to EA once _from_sequence is reliably strict GH#38254
 
         return result, counts
 

pandas/tests/extension/decimal/test_decimal.py

@@ -445,7 +445,8 @@ def test_groupby_agg():
     )
 
     # single key, selected column
-    expected = pd.Series(to_decimal([data[0], data[3]]))
+    # GH#38254 until _from_sequence is reliably strict, we cant retain dtype
+    expected = pd.Series(to_decimal([data[0], data[3]])).astype(object)
     result = df.groupby("id1")["decimals"].agg(lambda x: x.iloc[0])
     tm.assert_series_equal(result, expected, check_names=False)
     result = df["decimals"].groupby(df["id1"]).agg(lambda x: x.iloc[0])
@@ -455,14 +456,16 @@ def test_groupby_agg():
     expected = pd.Series(
         to_decimal([data[0], data[1], data[3]]),
         index=pd.MultiIndex.from_tuples([(0, 0), (0, 1), (1, 1)]),
-    )
+    ).astype(object)
     result = df.groupby(["id1", "id2"])["decimals"].agg(lambda x: x.iloc[0])
     tm.assert_series_equal(result, expected, check_names=False)
     result = df["decimals"].groupby([df["id1"], df["id2"]]).agg(lambda x: x.iloc[0])
     tm.assert_series_equal(result, expected, check_names=False)
 
     # multiple columns
-    expected = pd.DataFrame({"id2": [0, 1], "decimals": to_decimal([data[0], data[3]])})
+    expected = pd.DataFrame(
+        {"id2": [0, 1], "decimals": to_decimal([data[0], data[3]]).astype(object)}
+    )
     result = df.groupby("id1").agg(lambda x: x.iloc[0])
     tm.assert_frame_equal(result, expected, check_names=False)
 
@@ -478,7 +481,11 @@ def DecimalArray__my_sum(self):
 
     data = make_data()[:5]
     df = pd.DataFrame({"id": [0, 0, 0, 1, 1], "decimals": DecimalArray(data)})
-    expected = pd.Series(to_decimal([data[0] + data[1] + data[2], data[3] + data[4]]))
+
+    # GH#38254 until _from_sequence is reliably strict, we cant retain dtype
+    expected = pd.Series(
+        to_decimal([data[0] + data[1] + data[2], data[3] + data[4]])
+    ).astype(object)
 
     result = df.groupby("id")["decimals"].agg(lambda x: x.values.my_sum())
     tm.assert_series_equal(result, expected, check_names=False)

pandas/tests/groupby/aggregate/test_other.py

@@ -432,10 +432,13 @@ def test_agg_over_numpy_arrays():
     tm.assert_frame_equal(result, expected)
 
 
-def test_agg_tzaware_non_datetime_result():
+@pytest.mark.parametrize("as_period", [True, False])
+def test_agg_tzaware_non_datetime_result(as_period):
     # discussed in GH#29589, fixed in GH#29641, operating on tzaware values
     #  with function that is not dtype-preserving
     dti = pd.date_range("2012-01-01", periods=4, tz="UTC")
+    if as_period:
+        dti = dti.tz_localize(None).to_period("D")
     df = DataFrame({"a": [0, 0, 1, 1], "b": dti})
     gb = df.groupby("a")
 
@@ -454,6 +457,8 @@ def test_agg_tzaware_non_datetime_result():
     result = gb["b"].agg(lambda x: x.iloc[-1] - x.iloc[0])
     expected = Series([pd.Timedelta(days=1), pd.Timedelta(days=1)], name="b")
     expected.index.name = "a"
+    if as_period:
+        expected = expected.astype(object)
     tm.assert_series_equal(result, expected)
 
 
@@ -627,7 +632,8 @@ def test_groupby_agg_err_catching(err_cls):
         {"id1": [0, 0, 0, 1, 1], "id2": [0, 1, 0, 1, 1], "decimals": DecimalArray(data)}
     )
 
-    expected = Series(to_decimal([data[0], data[3]]))
+    # GH#38254 until _from_sequence is strict, we cannot reliably cast agg results
+    expected = Series(to_decimal([data[0], data[3]])).astype(object)
 
     def weird_func(x):
         # weird function that raise something other than TypeError or IndexError