From 04e8db370a412db07809650cf5eec79d063c373b Mon Sep 17 00:00:00 2001 From: Vedant Madane <6527493+VedantMadane@users.noreply.github.com> Date: Tue, 10 Feb 2026 22:26:47 +0530 Subject: [PATCH 1/2] TST: Add test for groupby.var() arrow dtype retention --- pandas/tests/groupby/aggregate/test_cython.py | 38 +++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/pandas/tests/groupby/aggregate/test_cython.py b/pandas/tests/groupby/aggregate/test_cython.py index a706ea795a0e2..f2471ab662091 100644 --- a/pandas/tests/groupby/aggregate/test_cython.py +++ b/pandas/tests/groupby/aggregate/test_cython.py @@ -413,3 +413,41 @@ def test_cython_agg_EA_known_dtypes(data, op_name, action, with_na): result = grouped["col"].aggregate(op_name) assert result.dtype == expected_dtype + + +@pytest.mark.parametrize( + "op_name", + [ + "var", + "std", + "sem", + "mean", + ], +) +@pytest.mark.parametrize( + "dtype", + [ + "int64[pyarrow]", + "float64[pyarrow]", + ], +) +def test_cython_agg_pyarrow_dtype_retention(op_name, dtype): + # GH#54627 - groupby.var() should return arrow types with arrow backed input + pa = pytest.importorskip("pyarrow") + + arr = pd.array([1, 2, 3, 4], dtype=dtype) + df = DataFrame({"key": ["a", "a", "b", "b"], "col": arr}) + grouped = df.groupby("key") + + result = getattr(grouped, op_name)() + # var/std/sem/mean of numeric pyarrow types should return double[pyarrow] + assert result["col"].dtype == pd.ArrowDtype(pa.float64()) + + result = grouped.aggregate(op_name) + assert result["col"].dtype == pd.ArrowDtype(pa.float64()) + + result = getattr(grouped["col"], op_name)() + assert result.dtype == pd.ArrowDtype(pa.float64()) + + result = grouped["col"].aggregate(op_name) + assert result.dtype == pd.ArrowDtype(pa.float64()) From d4fe63089ce637610377e7c4c2c3adb3886a52f9 Mon Sep 17 00:00:00 2001 From: Vedant Madane <6527493+VedantMadane@users.noreply.github.com> Date: Fri, 6 Mar 2026 09:19:29 +0530 Subject: [PATCH 2/2] TST: Move pyarrow groupby agg test to tests/extension/test_arrow.py Made-with: Cursor --- pandas/tests/extension/test_arrow.py | 22 +++++++++++ pandas/tests/groupby/aggregate/test_cython.py | 38 ------------------- 2 files changed, 22 insertions(+), 38 deletions(-) diff --git a/pandas/tests/extension/test_arrow.py b/pandas/tests/extension/test_arrow.py index 17e90f29eddc8..0a61c2db2eae4 100644 --- a/pandas/tests/extension/test_arrow.py +++ b/pandas/tests/extension/test_arrow.py @@ -3363,6 +3363,28 @@ def test_groupby_count_return_arrow_dtype(data_missing): tm.assert_frame_equal(result, expected) +@pytest.mark.parametrize("op_name", ["var", "std", "sem", "mean"]) +@pytest.mark.parametrize("dtype", ["int64[pyarrow]", "float64[pyarrow]"]) +def test_groupby_cython_agg_pyarrow_dtype_retention(op_name, dtype): + # GH#54627 + arr = pd.array([1, 2, 3, 4], dtype=dtype) + df = pd.DataFrame({"key": ["a", "a", "b", "b"], "col": arr}) + grouped = df.groupby("key") + expected_dtype = ArrowDtype(pa.float64()) + + result = getattr(grouped, op_name)() + assert result["col"].dtype == expected_dtype + + result = grouped.aggregate(op_name) + assert result["col"].dtype == expected_dtype + + result = getattr(grouped["col"], op_name)() + assert result.dtype == expected_dtype + + result = grouped["col"].aggregate(op_name) + assert result.dtype == expected_dtype + + def test_fixed_size_list(): # GH#55000 ser = pd.Series( diff --git a/pandas/tests/groupby/aggregate/test_cython.py b/pandas/tests/groupby/aggregate/test_cython.py index f2471ab662091..a706ea795a0e2 100644 --- a/pandas/tests/groupby/aggregate/test_cython.py +++ b/pandas/tests/groupby/aggregate/test_cython.py @@ -413,41 +413,3 @@ def test_cython_agg_EA_known_dtypes(data, op_name, action, with_na): result = grouped["col"].aggregate(op_name) assert result.dtype == expected_dtype - - -@pytest.mark.parametrize( - "op_name", - [ - "var", - "std", - "sem", - "mean", - ], -) -@pytest.mark.parametrize( - "dtype", - [ - "int64[pyarrow]", - "float64[pyarrow]", - ], -) -def test_cython_agg_pyarrow_dtype_retention(op_name, dtype): - # GH#54627 - groupby.var() should return arrow types with arrow backed input - pa = pytest.importorskip("pyarrow") - - arr = pd.array([1, 2, 3, 4], dtype=dtype) - df = DataFrame({"key": ["a", "a", "b", "b"], "col": arr}) - grouped = df.groupby("key") - - result = getattr(grouped, op_name)() - # var/std/sem/mean of numeric pyarrow types should return double[pyarrow] - assert result["col"].dtype == pd.ArrowDtype(pa.float64()) - - result = grouped.aggregate(op_name) - assert result["col"].dtype == pd.ArrowDtype(pa.float64()) - - result = getattr(grouped["col"], op_name)() - assert result.dtype == pd.ArrowDtype(pa.float64()) - - result = grouped["col"].aggregate(op_name) - assert result.dtype == pd.ArrowDtype(pa.float64())