ENH: retain masked EA dtypes in groupby with as_index=False

pandas/core/groupby/generic.py

@@ -1046,7 +1046,7 @@ def aggregate(self, func=None, *args, engine=None, engine_kwargs=None, **kwargs)
             self._insert_inaxis_grouper_inplace(result)
             result.index = np.arange(len(result))
 
-        return result._convert(datetime=True)
+        return result
 
     agg = aggregate
 
@@ -1633,7 +1633,7 @@ def _wrap_agged_manager(self, mgr: Manager2D) -> DataFrame:
         if self.axis == 1:
             result = result.T
 
-        return self._reindex_output(result)._convert(datetime=True)
+        return self._reindex_output(result)
 
     def _iterate_column_groupbys(self):
         for i, colname in enumerate(self._selected_obj.columns):

pandas/core/groupby/grouper.py

@@ -10,6 +10,7 @@
 import numpy as np
 
 from pandas._typing import (
+    ArrayLike,
     FrameOrSeries,
     final,
 )
@@ -462,6 +463,8 @@ def __init__(
         self.in_axis = in_axis
         self.dropna = dropna
 
+        self._group_arraylike = None
+
         # right place for this?
         if isinstance(grouper, (Series, Index)) and name is None:
             self.name = grouper.name
@@ -601,6 +604,22 @@ def codes(self) -> np.ndarray:
         # expected "ndarray")
         return self._codes  # type: ignore[return-value]
 
+    @cache_readonly
+    def result_arraylike(self) -> ArrayLike:
+        """
+        Analogous to result_index, but holding an ArrayLike to ensure
+        we can can retain ExtensionDtypes.
+        """
+        ridx = self.result_index  # initialized _group_arraylike
+
+        if self._group_arraylike is None:
+            # This should only occur when ridx is CategoricalIndex
+            self._group_arraylike = ridx._values
+
+        # error: Incompatible return value type (got "None", expected
+        # "Union[ExtensionArray, ndarray]")
+        return self._group_arraylike  # type: ignore[return-value]
+
     @cache_readonly
     def result_index(self) -> Index:
         if self.all_grouper is not None:
@@ -623,7 +642,7 @@ def _make_codes(self) -> None:
         # we have a list of groupers
         if isinstance(self.grouper, ops.BaseGrouper):
             codes = self.grouper.codes_info
-            uniques = self.grouper.result_index
+            uniques = self.grouper.result_arraylike
         else:
             # GH35667, replace dropna=False with na_sentinel=None
             if not self.dropna:
@@ -633,9 +652,9 @@ def _make_codes(self) -> None:
             codes, uniques = algorithms.factorize(
                 self.grouper, sort=self.sort, na_sentinel=na_sentinel
             )
-            uniques = Index(uniques, name=self.name)
         self._codes = codes
-        self._group_index = uniques
+        self._group_arraylike = uniques
+        self._group_index = Index(uniques, name=self.name)
 
     @cache_readonly
     def groups(self) -> dict[Hashable, np.ndarray]:

pandas/core/groupby/ops.py

@@ -908,6 +908,23 @@ def reconstructed_codes(self) -> list[np.ndarray]:
         ids, obs_ids, _ = self.group_info
         return decons_obs_group_ids(ids, obs_ids, self.shape, codes, xnull=True)
 
+    @cache_readonly
+    def result_arraylike(self) -> ArrayLike:
+        """
+        Analogous to result_index, but returning an ndarray/ExtensionArray
+        allowing us to retain ExtensionDtypes not supported by Index.
+        """
+        # TODO: once Index supports arbitrary EAs, this can be removed in favor
+        #  of result_index
+        if len(self.groupings) == 1:
+            return self.groupings[0].result_arraylike
+
+        codes = self.reconstructed_codes
+        levels = [ping.result_arraylike for ping in self.groupings]
+        return MultiIndex(
+            levels=levels, codes=codes, verify_integrity=False, names=self.names
+        )._values
+
     @cache_readonly
     def result_index(self) -> Index:
         if len(self.groupings) == 1:
@@ -924,12 +941,12 @@ def get_group_levels(self) -> list[Index]:
         # Note: only called from _insert_inaxis_grouper_inplace, which
         #  is only called for BaseGrouper, never for BinGrouper
         if len(self.groupings) == 1:
-            return [self.groupings[0].result_index]
+            return [self.groupings[0].result_arraylike]
 
         name_list = []
         for ping, codes in zip(self.groupings, self.reconstructed_codes):
             codes = ensure_platform_int(codes)
-            levels = ping.result_index.take(codes)
+            levels = ping.result_arraylike.take(codes)
 
             name_list.append(levels)
 
@@ -991,7 +1008,10 @@ def agg_series(self, obj: Series, func: F) -> ArrayLike:
             result = self._aggregate_series_fast(obj, func)
             cast_back = False
 
-        npvalues = lib.maybe_convert_objects(result, try_float=False)
+        convert_datetime = obj.dtype.kind == "M"
+        npvalues = lib.maybe_convert_objects(
+            result, try_float=False, convert_datetime=convert_datetime
+        )
         if cast_back:
             # TODO: Is there a documented reason why we dont always cast_back?
             out = maybe_cast_pointwise_result(npvalues, obj.dtype, numeric_only=True)

pandas/tests/extension/base/groupby.py

@@ -22,14 +22,14 @@ def test_grouping_grouper(self, data_for_grouping):
     def test_groupby_extension_agg(self, as_index, data_for_grouping):
         df = pd.DataFrame({"A": [1, 1, 2, 2, 3, 3, 1, 4], "B": data_for_grouping})
         result = df.groupby("B", as_index=as_index).A.mean()
-        _, index = pd.factorize(data_for_grouping, sort=True)
+        _, uniques = pd.factorize(data_for_grouping, sort=True)
 
-        index = pd.Index(index, name="B")
-        expected = pd.Series([3, 1, 4], index=index, name="A")
         if as_index:
+            index = pd.Index(uniques, name="B")
+            expected = pd.Series([3, 1, 4], index=index, name="A")
             self.assert_series_equal(result, expected)
         else:
-            expected = expected.reset_index()
+            expected = pd.DataFrame({"B": uniques, "A": [3, 1, 4]})
             self.assert_frame_equal(result, expected)
 
     def test_groupby_agg_extension(self, data_for_grouping):

pandas/tests/extension/json/test_json.py

@@ -312,10 +312,6 @@ def test_groupby_extension_apply(self):
         we'll be able to dispatch unique.
         """
 
-    @pytest.mark.parametrize("as_index", [True, False])
-    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)

pandas/tests/extension/test_boolean.py

@@ -269,14 +269,14 @@ def test_grouping_grouper(self, data_for_grouping):
     def test_groupby_extension_agg(self, as_index, data_for_grouping):
         df = pd.DataFrame({"A": [1, 1, 2, 2, 3, 3, 1], "B": data_for_grouping})
         result = df.groupby("B", as_index=as_index).A.mean()
-        _, index = pd.factorize(data_for_grouping, sort=True)
+        _, uniques = pd.factorize(data_for_grouping, sort=True)
 
-        index = pd.Index(index, name="B")
-        expected = pd.Series([3, 1], index=index, name="A")
         if as_index:
+            index = pd.Index(uniques, name="B")
+            expected = pd.Series([3, 1], index=index, name="A")
             self.assert_series_equal(result, expected)
         else:
-            expected = expected.reset_index()
+            expected = pd.DataFrame({"B": uniques, "A": [3, 1]})
             self.assert_frame_equal(result, expected)
 
     def test_groupby_agg_extension(self, data_for_grouping):

pandas/tests/groupby/test_groupby.py

@@ -694,6 +694,10 @@ def test_ops_not_as_index(reduction_func):
         expected = expected.rename("size")
     expected = expected.reset_index()
 
+    if reduction_func != "size":
+        # 32 bit compat -> groupby preserves dtype whereas reset_index casts to int64
+        expected["a"] = expected["a"].astype(df["a"].dtype)
+
     g = df.groupby("a", as_index=False)
 
     result = getattr(g, reduction_func)()