BUG: Aggregation on arrow array return same type.

Signed-off-by: Liang Yan <[email protected]>

Author: liang3zy22

pandas/core/groupby/ops.py

@@ -73,6 +73,8 @@
 if TYPE_CHECKING:
     from pandas.core.generic import NDFrame
 
+from pandas.arrays import ArrowExtensionArray
+
 
 def check_result_array(obj, dtype):
     # Our operation is supposed to be an aggregation/reduction. If
@@ -837,7 +839,11 @@ def agg_series(
         # test_groupby_empty_with_category gets here with self.ngroups == 0
         #  and len(obj) > 0
 
-        if len(obj) > 0 and not isinstance(obj._values, np.ndarray):
+        if (
+            len(obj) > 0
+            and 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

pandas/tests/groupby/aggregate/test_aggregate.py

@@ -5,6 +5,7 @@
 import functools
 from functools import partial
 import re
+import typing
 
 import numpy as np
 import pytest
@@ -1603,3 +1604,23 @@ def test_agg_with_as_index_false_with_list():
         columns=MultiIndex.from_tuples([("a1", ""), ("a2", ""), ("b", "sum")]),
     )
     tm.assert_frame_equal(result, expected)
+
+
+@pytest.mark.skipif(
+    not typing.TYPE_CHECKING, reason="TYPE_CHECKING must be True to import pyarrow"
+)
+def test_agg_arrow_type():
+    df = DataFrame.from_dict(
+        {
+            "category": ["A"] * 10 + ["B"] * 10,
+            "bool_numpy": [True] * 5 + [False] * 5 + [True] * 5 + [False] * 5,
+        }
+    )
+    df["bool_arrow"] = df["bool_numpy"].astype("bool[pyarrow]")
+    result = df.groupby("category").agg(lambda x: x.sum() / x.count())
+    expected = DataFrame(
+        [[0.5, 0.5], [0.5, 0.5]],
+        columns=["bool_numpy", "bool_arrow"],
+        index=Index(["A", "B"], name="category"),
+    )
+    tm.assert_frame_equal(result, expected)