Convert result of group by agg to pyarrow if input is pyarrow

pandas/core/dtypes/cast.py

@@ -478,6 +478,42 @@ def maybe_cast_pointwise_result(
     return result
 
 
+def maybe_cast_to_pyarrow_dtype(
+    result: ArrayLike, converted_result: ArrayLike
+) -> ArrayLike:
+    """
+    Try casting result of a pointwise operation to its pyarrow dtype if
+    appropriate.
+
+    Parameters
+    ----------
+    result : array-like
+        Result to cast.
+
+    Returns
+    -------
+    result : array-like
+        result maybe casted to the dtype.
+    """
+    try:
+        import pyarrow as pa
+        from pyarrow import (
+            ArrowInvalid,
+            ArrowNotImplementedError,
+        )
+
+        from pandas.core.construction import array as pd_array
+
+        result[isna(result)] = np.nan
+        pyarrow_result = pa.array(result)
+        pandas_pyarrow_dtype = ArrowDtype(pyarrow_result.type)
+        result = pd_array(result, dtype=pandas_pyarrow_dtype)
+    except (ArrowNotImplementedError, ArrowInvalid):
+        return converted_result
+
+    return result
+
+
 def _maybe_cast_to_extension_array(
     cls: type[ExtensionArray], obj: ArrayLike, dtype: ExtensionDtype | None = None
 ) -> ArrayLike:

pandas/core/groupby/ops.py

@@ -36,6 +36,7 @@
 
 from pandas.core.dtypes.cast import (
     maybe_cast_pointwise_result,
+    maybe_cast_to_pyarrow_dtype,
     maybe_downcast_to_dtype,
 )
 from pandas.core.dtypes.common import (
@@ -51,6 +52,7 @@
 )
 
 from pandas.core.arrays import Categorical
+from pandas.core.arrays.arrow.array import ArrowExtensionArray
 from pandas.core.frame import DataFrame
 from pandas.core.groupby import grouper
 from pandas.core.indexes.api import (
@@ -914,18 +916,22 @@ def agg_series(
         np.ndarray or ExtensionArray
         """
 
-        if not isinstance(obj._values, np.ndarray):
+        if not isinstance(obj._values, np.ndarray) and not isinstance(
+            obj._values, ArrowExtensionArray
+        ):
             # we can preserve a little bit more aggressively with EA dtype
             #  because maybe_cast_pointwise_result will do a try/except
             #  with _from_sequence.  NB we are assuming here that _from_sequence
             #  is sufficiently strict that it casts appropriately.
             preserve_dtype = True
 
         result = self._aggregate_series_pure_python(obj, func)
-
         npvalues = lib.maybe_convert_objects(result, try_float=False)
+
         if preserve_dtype:
             out = maybe_cast_pointwise_result(npvalues, obj.dtype, numeric_only=True)
+        elif isinstance(obj._values, ArrowExtensionArray):
+            out = maybe_cast_to_pyarrow_dtype(result, npvalues)
         else:
             out = npvalues
         return out