Rework

Author: rhshadrach

pandas/core/groupby/ops.py

@@ -44,7 +44,6 @@
     ensure_uint64,
     is_1d_only_ea_dtype,
 )
-from pandas.core.dtypes.dtypes import ArrowDtype
 from pandas.core.dtypes.missing import (
     isna,
     maybe_fill,
@@ -956,19 +955,22 @@ def agg_series(
         -------
         np.ndarray or ExtensionArray
         """
-
         result = self._aggregate_series_pure_python(obj, func)
         npvalues = lib.maybe_convert_objects(result, try_float=False)
 
         if isinstance(obj._values, ArrowExtensionArray):
-            out = maybe_cast_pointwise_result(
-                npvalues, obj.dtype, numeric_only=True, same_dtype=preserve_dtype
-            )
-            import pyarrow as pa
+            from pandas.core.dtypes.common import is_string_dtype
 
-            if isinstance(out.dtype, ArrowDtype) and pa.types.is_struct(
-                out.dtype.pyarrow_dtype
-            ):
+            if not is_string_dtype(obj.dtype) or is_string_dtype(npvalues):
+                out = maybe_cast_pointwise_result(
+                    npvalues, obj.dtype, numeric_only=True, same_dtype=preserve_dtype
+                )
+
+                # if isinstance(out.dtype, ArrowDtype) and pa.types.is_struct(
+                #     out.dtype.pyarrow_dtype
+                # ):
+                #     out = npvalues
+            else:
                 out = npvalues
 
         elif not isinstance(obj._values, np.ndarray):

pandas/tests/groupby/aggregate/test_aggregate.py

@@ -10,6 +10,7 @@
 import pytest
 
 from pandas.errors import SpecificationError
+import pandas.util._test_decorators as td
 
 from pandas.core.dtypes.common import is_integer_dtype
 
@@ -23,6 +24,7 @@
     to_datetime,
 )
 import pandas._testing as tm
+from pandas.arrays import ArrowExtensionArray
 from pandas.core.groupby.grouper import Grouping
 
 
@@ -1812,16 +1814,18 @@ def test_groupby_aggregation_func_list_multi_index_duplicate_columns():
 @pytest.mark.parametrize(
     "input_dtype, output_dtype",
     [
+        # With NumPy arrays, the results from the UDF would be e.g. np.float32 scalars
+        # which we can therefore preserve. However with PyArrow arrays, the results are
+        # Python scalars so we have no information about size or uint vs int.
         ("float[pyarrow]", "double[pyarrow]"),
         ("int64[pyarrow]", "int64[pyarrow]"),
         ("uint64[pyarrow]", "int64[pyarrow]"),
         ("bool[pyarrow]", "bool[pyarrow]"),
     ],
 )
 def test_agg_lambda_pyarrow_dtype_conversion(input_dtype, output_dtype):
-    # GH#53030
-    # test numpy dtype conversion back to pyarrow dtype
-    # complexes, floats, ints, uints, object
+    # GH#59601
+    # Test PyArrow dtype conversion back to PyArrow dtype
     df = DataFrame(
         {
             "A": ["c1", "c2", "c3", "c1", "c2", "c3"],
@@ -1839,7 +1843,7 @@ def test_agg_lambda_pyarrow_dtype_conversion(input_dtype, output_dtype):
 
 
 def test_agg_lambda_complex128_dtype_conversion():
-    # GH#53030
+    # GH#59601
     df = DataFrame(
         {"A": ["c1", "c2", "c3"], "B": pd.array([100, 200, 255], "int64[pyarrow]")}
     )
@@ -1877,8 +1881,11 @@ def test_agg_lambda_numpy_uint64_to_pyarrow_dtype_conversion():
     tm.assert_frame_equal(result, expected)
 
 
+@td.skip_if_no("pyarrow")
 def test_agg_lambda_pyarrow_struct_to_object_dtype_conversion():
     # GH#53030
+    import pyarrow as pa
+
     df = DataFrame(
         {
             "A": ["c1", "c2", "c3"],
@@ -1888,8 +1895,10 @@ def test_agg_lambda_pyarrow_struct_to_object_dtype_conversion():
     gb = df.groupby("A")
     result = gb.agg(lambda x: {"number": 1})
 
+    arr = pa.array([{"number": 1}, {"number": 1}, {"number": 1}])
     expected = DataFrame(
-        {"B": pd.array([{"number": 1}, {"number": 1}, {"number": 1}], dtype="object")},
+        {"B": ArrowExtensionArray(arr)},
         index=Index(["c1", "c2", "c3"], name="A"),
     )
+
     tm.assert_frame_equal(result, expected)

pandas/tests/groupby/test_groupby.py

@@ -26,7 +26,6 @@
 )
 import pandas._testing as tm
 from pandas.core.arrays import BooleanArray
-from pandas.core.arrays.string_arrow import ArrowStringArrayNumpySemantics
 import pandas.core.common as com
 
 pytestmark = pytest.mark.filterwarnings("ignore:Mean of empty slice:RuntimeWarning")
@@ -2435,30 +2434,28 @@ def test_rolling_wrong_param_min_period():
 
 def test_by_column_values_with_same_starting_value(any_string_dtype):
     # GH29635
+    dtype = any_string_dtype
     df = DataFrame(
         {
             "Name": ["Thomas", "Thomas", "Thomas John"],
             "Credit": [1200, 1300, 900],
-            "Mood": Series(["sad", "happy", "happy"], dtype=any_string_dtype),
+            "Mood": Series(["sad", "happy", "happy"], dtype=dtype),
         }
     )
     aggregate_details = {"Mood": Series.mode, "Credit": "sum"}
 
     result = df.groupby(["Name"]).agg(aggregate_details)
-    expected_result = DataFrame(
+    expected = DataFrame(
         {
             "Mood": [["happy", "sad"], "happy"],
             "Credit": [2500, 900],
             "Name": ["Thomas", "Thomas John"],
         },
     ).set_index("Name")
-    if dtype == "string[pyarrow_numpy]":
-        import pyarrow as pa
-
-        mood_values = ArrowStringArrayNumpySemantics(pa.array(["happy", "sad"]))
-        expected_result["Mood"] = [mood_values, "happy"]
-        expected_result["Mood"] = expected_result["Mood"].astype(dtype)
-    tm.assert_frame_equal(result, expected_result)
+    if getattr(dtype, "storage", None) == "pyarrow":
+        mood_values = pd.array(["happy", "sad"], dtype=dtype)
+        expected["Mood"] = [mood_values, "happy"]
+    tm.assert_frame_equal(result, expected)
 
 
 def test_groupby_none_in_first_mi_level():