REGR: Bug fix for ExtensionArray groupby aggregation on non-numeric types (#38982)

Author: BryanCutler

doc/source/whatsnew/v1.2.1.rst

@@ -28,6 +28,7 @@ Fixed regressions
 - Fixed regression in :meth:`DataFrame.replace` raising ``ValueError`` when :class:`DataFrame` has dtype ``bytes`` (:issue:`38900`)
 - Fixed regression in :meth:`DataFrameGroupBy.diff` raising for ``int8`` and ``int16`` columns (:issue:`39050`)
 - Fixed regression that raised ``AttributeError`` with PyArrow versions [0.16.0, 1.0.0) (:issue:`38801`)
+- Fixed regression in :meth:`DataFrame.groupby` when aggregating an :class:`ExtensionDType` that could fail for non-numeric values (:issue:`38980`)
 -
 -
 

pandas/core/groupby/ops.py

@@ -540,7 +540,9 @@ def _ea_wrap_cython_operation(
             result = type(orig_values)._from_sequence(res_values)
             return result
 
-        raise NotImplementedError(values.dtype)
+        raise NotImplementedError(
+            f"function is not implemented for this dtype: {values.dtype}"
+        )
 
     @final
     def _cython_operation(

pandas/tests/extension/base/groupby.py

@@ -33,6 +33,22 @@ def test_groupby_extension_agg(self, as_index, data_for_grouping):
             expected = expected.reset_index()
             self.assert_frame_equal(result, expected)
 
+    def test_groupby_agg_extension(self, data_for_grouping):
+        # GH#38980 groupby agg on extension type fails for non-numeric types
+        df = pd.DataFrame({"A": [1, 1, 2, 2, 3, 3, 1, 4], "B": data_for_grouping})
+
+        expected = df.iloc[[0, 2, 4, 7]]
+        expected = expected.set_index("A")
+
+        result = df.groupby("A").agg({"B": "first"})
+        self.assert_frame_equal(result, expected)
+
+        result = df.groupby("A").agg("first")
+        self.assert_frame_equal(result, expected)
+
+        result = df.groupby("A").first()
+        self.assert_frame_equal(result, expected)
+
     def test_groupby_extension_no_sort(self, data_for_grouping):
         df = pd.DataFrame({"A": [1, 1, 2, 2, 3, 3, 1, 4], "B": data_for_grouping})
         result = df.groupby("B", sort=False).A.mean()

pandas/tests/extension/decimal/test_decimal.py

@@ -197,6 +197,10 @@ class TestGroupby(BaseDecimal, base.BaseGroupbyTests):
     def test_groupby_apply_identity(self, data_for_grouping):
         super().test_groupby_apply_identity(data_for_grouping)
 
+    @pytest.mark.xfail(reason="GH#39098: Converts agg result to object")
+    def test_groupby_agg_extension(self, data_for_grouping):
+        super().test_groupby_agg_extension(data_for_grouping)
+
 
 class TestSetitem(BaseDecimal, base.BaseSetitemTests):
     pass

pandas/tests/extension/json/test_json.py

@@ -313,6 +313,10 @@ def test_groupby_extension_apply(self):
     def test_groupby_extension_agg(self, as_index, data_for_grouping):
         super().test_groupby_extension_agg(as_index, data_for_grouping)
 
+    @pytest.mark.xfail(reason="GH#39098: Converts agg result to object")
+    def test_groupby_agg_extension(self, data_for_grouping):
+        super().test_groupby_agg_extension(data_for_grouping)
+
 
 class TestArithmeticOps(BaseJSON, base.BaseArithmeticOpsTests):
     def test_error(self, data, all_arithmetic_operators):

pandas/tests/extension/test_boolean.py

@@ -291,6 +291,22 @@ def test_groupby_extension_agg(self, as_index, data_for_grouping):
             expected = expected.reset_index()
             self.assert_frame_equal(result, expected)
 
+    def test_groupby_agg_extension(self, data_for_grouping):
+        # GH#38980 groupby agg on extension type fails for non-numeric types
+        df = pd.DataFrame({"A": [1, 1, 2, 2, 3, 3, 1], "B": data_for_grouping})
+
+        expected = df.iloc[[0, 2, 4]]
+        expected = expected.set_index("A")
+
+        result = df.groupby("A").agg({"B": "first"})
+        self.assert_frame_equal(result, expected)
+
+        result = df.groupby("A").agg("first")
+        self.assert_frame_equal(result, expected)
+
+        result = df.groupby("A").first()
+        self.assert_frame_equal(result, expected)
+
     def test_groupby_extension_no_sort(self, data_for_grouping):
         df = pd.DataFrame({"A": [1, 1, 2, 2, 3, 3, 1], "B": data_for_grouping})
         result = df.groupby("B", sort=False).A.mean()