BUG: groupby().agg fails on categorical column

pandas/core/groupby/base.py

@@ -92,6 +92,11 @@ def _gotitem(self, key, ndim, subset=None):
 
 cython_cast_blacklist = frozenset(["rank", "count", "size", "idxmin", "idxmax"])
 
+cython_cast_cat_type_list = frozenset(["first", "last"])
+cython_cast_keep_type_list = cython_cast_cat_type_list | frozenset(
+    ["sum", "min", "max"]
+)
+
 # List of aggregation/reduction functions.
 # These map each group to a single numeric value
 reduction_kernels = frozenset(

pandas/core/groupby/groupby.py

@@ -59,6 +59,7 @@ class providing the base-class of operations.
 from pandas.core.frame import DataFrame
 from pandas.core.generic import NDFrame
 from pandas.core.groupby import base, ops
+from pandas.core.groupby.base import cython_cast_keep_type_list
 from pandas.core.indexes.api import CategoricalIndex, Index, MultiIndex
 from pandas.core.series import Series
 from pandas.core.sorting import get_group_index_sorter
@@ -792,7 +793,7 @@ def _cumcount_array(self, ascending: bool = True):
         rev[sorter] = np.arange(count, dtype=np.intp)
         return out[rev].astype(np.int64, copy=False)
 
-    def _try_cast(self, result, obj, numeric_only: bool = False):
+    def _try_cast(self, result, obj, numeric_only: bool = False, how=None):
         """
         Try to cast the result to our obj original type,
         we may have roundtripped through object in the mean-time.
@@ -807,13 +808,19 @@ def _try_cast(self, result, obj, numeric_only: bool = False):
             dtype = obj.dtype
 
         if not is_scalar(result):
-            if is_extension_array_dtype(dtype) and dtype.kind != "M":
-                # The function can return something of any type, so check
-                #  if the type is compatible with the calling EA.
-                # datetime64tz is handled correctly in agg_series,
-                #  so is excluded here.
 
-                if len(result) and isinstance(result[0], dtype.type):
+            # The function can return something of any type, so check
+            #  if the type is compatible with the calling EA.
+            # datetime64tz is handled correctly in agg_series,
+            #  so is excluded here.
+            if is_extension_array_dtype(dtype) and dtype.kind != "M":
+                # if how is in cython_cast_keep_type_list, which means it
+                # should be cast back to return the same type as obj
+                if (
+                    len(result)
+                    and isinstance(result[0], dtype.type)
+                    or how in cython_cast_keep_type_list
+                ):
                     cls = dtype.construct_array_type()
                     result = try_cast_to_ea(cls, result, dtype=dtype)
 
@@ -900,7 +907,7 @@ def _cython_agg_general(
             else:
                 assert result.ndim == 1
                 key = base.OutputKey(label=name, position=idx)
-                output[key] = self._try_cast(result, obj)
+                output[key] = self._try_cast(result, obj, how=how)
                 idx += 1
 
         if len(output) == 0:

pandas/core/groupby/ops.py

@@ -43,6 +43,7 @@
 from pandas.core.frame import DataFrame
 from pandas.core.generic import NDFrame
 from pandas.core.groupby import base, grouper
+from pandas.core.groupby.base import cython_cast_cat_type_list
 from pandas.core.indexes.api import Index, MultiIndex, ensure_index
 from pandas.core.series import Series
 from pandas.core.sorting import (
@@ -451,7 +452,12 @@ def _cython_operation(
 
         # categoricals are only 1d, so we
         # are not setup for dim transforming
-        if is_categorical_dtype(values) or is_sparse(values):
+        # those four cython agg that should work with categoricals
+        if (
+            is_categorical_dtype(values)
+            and how not in cython_cast_cat_type_list
+            or is_sparse(values)
+        ):
             raise NotImplementedError(f"{values.dtype} dtype not supported")
         elif is_datetime64_any_dtype(values):
             if how in ["add", "prod", "cumsum", "cumprod"]:

pandas/tests/groupby/test_categorical.py

@@ -1376,3 +1376,14 @@ def test_groupby_agg_non_numeric():
 
     result = df.groupby([1, 2, 1]).nunique()
     tm.assert_frame_equal(result, expected)
+
+
+@pytest.mark.parametrize("func", ["first", "last"])
+def test_groupby_agg_categorical_first_last(func):
+    # GH 31450
+    df = pd.DataFrame({"col_num": [1, 1, 2, 3]})
+    df["col_cat"] = df["col_num"].astype("category")
+
+    grouped = df.groupby("col_num").agg({"col_cat": func})
+    expected = df.groupby("col_num").agg(func)
+    tm.assert_frame_equal(grouped, expected)