Update implementation to use pyarrow array method

Author: Kei

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 (
@@ -45,15 +46,13 @@
     ensure_uint64,
     is_1d_only_ea_dtype,
 )
-from pandas.core.dtypes.dtypes import ArrowDtype
 from pandas.core.dtypes.missing import (
     isna,
     maybe_fill,
 )
 
 from pandas.core.arrays import Categorical
 from pandas.core.arrays.arrow.array import ArrowExtensionArray
-from pandas.core.construction import array as pd_array
 from pandas.core.frame import DataFrame
 from pandas.core.groupby import grouper
 from pandas.core.indexes.api import (
@@ -927,21 +926,12 @@ def agg_series(
             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)
-            and npvalues.dtype != np.dtype("object")
-            and npvalues.dtype != np.dtype("complex128")
-        ):
-            import pyarrow as pa
-
-            pyarrow_dtype = pa.from_numpy_dtype(npvalues.dtype)
-            pandas_pyarrow_dtype = ArrowDtype(pyarrow_dtype)
-            out = pd_array(npvalues, dtype=pandas_pyarrow_dtype)
-
+        elif isinstance(obj._values, ArrowExtensionArray):
+            out = maybe_cast_to_pyarrow_dtype(result, npvalues)
         else:
             out = npvalues
         return out