Skip to content

Commit c2a216d

Browse files
committed
Add test code
1 parent 3c9cf09 commit c2a216d

File tree

1 file changed

+21
-0
lines changed

1 file changed

+21
-0
lines changed

pandas/tests/extension/test_arrow.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3251,6 +3251,27 @@ def test_groupby_count_return_arrow_dtype(data_missing):
32513251
tm.assert_frame_equal(result, expected)
32523252

32533253

3254+
@pytest.mark.parametrize(
3255+
"func, func_dtype",
3256+
[
3257+
[lambda x: x.to_dict(), "object"],
3258+
[lambda x: 1, "int64"],
3259+
[lambda x: "s", ArrowDtype(pa.string())],
3260+
],
3261+
)
3262+
def test_groupby_aggregate_coersion(func, func_dtype):
3263+
# GH 61636
3264+
df = pd.DataFrame(
3265+
{
3266+
"b": pd.array([0, 1]),
3267+
"c": pd.array(["X", "Y"], dtype=ArrowDtype(pa.string())),
3268+
},
3269+
index=pd.Index(["A", "B"], name="a"),
3270+
)
3271+
result = df.groupby("b").agg(func)
3272+
assert result["c"].dtype == func_dtype
3273+
3274+
32543275
def test_fixed_size_list():
32553276
# GH#55000
32563277
ser = pd.Series(

0 commit comments

Comments
 (0)