diff --git a/doc/source/whatsnew/v3.0.0.rst b/doc/source/whatsnew/v3.0.0.rst index f91d40c4d9ea9..5341f338e6249 100644 --- a/doc/source/whatsnew/v3.0.0.rst +++ b/doc/source/whatsnew/v3.0.0.rst @@ -1064,6 +1064,7 @@ Groupby/resample/rolling - Bug in :meth:`DataFrame.resample` changing index type to :class:`MultiIndex` when the dataframe is empty and using an upsample method (:issue:`55572`) - Bug in :meth:`DataFrameGroupBy.agg` and :meth:`SeriesGroupBy.agg` that was returning numpy dtype values when input values are pyarrow dtype values, instead of returning pyarrow dtype values. (:issue:`53030`) - Bug in :meth:`DataFrameGroupBy.agg` that raises ``AttributeError`` when there is dictionary input and duplicated columns, instead of returning a DataFrame with the aggregation of all duplicate columns. (:issue:`55041`) +- Bug in :meth:`DataFrameGroupBy.agg` when applied to columns with :class:`ArrowDtype`, where pandas attempted to cast the result back to the original dtype (:issue:`61636`) - Bug in :meth:`DataFrameGroupBy.agg` where applying a user-defined function to an empty DataFrame returned a Series instead of an empty DataFrame. (:issue:`61503`) - Bug in :meth:`DataFrameGroupBy.apply` and :meth:`SeriesGroupBy.apply` for empty data frame with ``group_keys=False`` still creating output index using group keys. (:issue:`60471`) - Bug in :meth:`DataFrameGroupBy.apply` that was returning a completely empty DataFrame when all return values of ``func`` were ``None`` instead of returning an empty DataFrame with the original columns and dtypes. (:issue:`57775`) diff --git a/pandas/core/arrays/arrow/array.py b/pandas/core/arrays/arrow/array.py index 653a900fbfe45..f341d13903613 100644 --- a/pandas/core/arrays/arrow/array.py +++ b/pandas/core/arrays/arrow/array.py @@ -189,6 +189,7 @@ def floordiv_compat( ArrayLike, AxisInt, Dtype, + DtypeObj, FillnaOptions, InterpolateOptions, Iterator, @@ -308,6 +309,20 @@ def __init__(self, values: pa.Array | pa.ChunkedArray) -> None: ) self._dtype = ArrowDtype(self._pa_array.type) + @classmethod + def _from_scalars(cls, scalars, dtype: DtypeObj) -> Self: + inferred_dtype = lib.infer_dtype(scalars, skipna=True) + try: + pa_array = cls._from_sequence(scalars, dtype=dtype) + except pa.ArrowNotImplementedError as err: + # _from_scalars should only raise ValueError or TypeError. + raise ValueError from err + + same_dtype = lib.infer_dtype(pa_array, skipna=True) == inferred_dtype + if not same_dtype: + raise ValueError + return pa_array + @classmethod def _from_sequence( cls, scalars, *, dtype: Dtype | None = None, copy: bool = False diff --git a/pandas/core/arrays/string_arrow.py b/pandas/core/arrays/string_arrow.py index 9046d83dcc09f..a91b16482963a 100644 --- a/pandas/core/arrays/string_arrow.py +++ b/pandas/core/arrays/string_arrow.py @@ -53,6 +53,7 @@ from pandas._typing import ( ArrayLike, Dtype, + DtypeObj, NpDtype, Scalar, npt, @@ -190,6 +191,12 @@ def __len__(self) -> int: """ return len(self._pa_array) + @classmethod + def _from_scalars(cls, scalars, dtype: DtypeObj) -> Self: + if lib.infer_dtype(scalars, skipna=True) not in ["string", "empty"]: + raise ValueError + return cls._from_sequence(scalars, dtype=dtype) + @classmethod def _from_sequence( cls, scalars, *, dtype: Dtype | None = None, copy: bool = False diff --git a/pandas/tests/extension/test_arrow.py b/pandas/tests/extension/test_arrow.py index e9d014a0eb29d..b56962dc1a28d 100644 --- a/pandas/tests/extension/test_arrow.py +++ b/pandas/tests/extension/test_arrow.py @@ -3251,6 +3251,27 @@ def test_groupby_count_return_arrow_dtype(data_missing): tm.assert_frame_equal(result, expected) +@pytest.mark.parametrize( + "func, func_dtype", + [ + [lambda x: x.to_dict(), "object"], + [lambda x: 1, "int64"], + [lambda x: "s", ArrowDtype(pa.string())], + ], +) +def test_groupby_aggregate_coersion(func, func_dtype): + # GH 61636 + df = pd.DataFrame( + { + "b": pd.array([0, 1]), + "c": pd.array(["X", "Y"], dtype=ArrowDtype(pa.string())), + }, + index=pd.Index(["A", "B"], name="a"), + ) + result = df.groupby("b").agg(func) + assert result["c"].dtype == func_dtype + + def test_fixed_size_list(): # GH#55000 ser = pd.Series(