Skip to content

Commit bb6343b

Browse files
author
Kei
committed
Override test case for test_arrow.py
1 parent bebc442 commit bb6343b

File tree

2 files changed

+23
-9
lines changed

2 files changed

+23
-9
lines changed

pandas/tests/extension/base/groupby.py

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -58,18 +58,12 @@ def test_groupby_extension_agg(self, as_index, data_for_grouping):
5858
expected = pd.DataFrame({"B": uniques, "A": exp_vals})
5959
tm.assert_frame_equal(result, expected)
6060

61-
def test_groupby_agg_extension(
62-
self, data_for_grouping, expected_inferred_result_dtype
63-
):
61+
def test_groupby_agg_extension(self, data_for_grouping):
6462
# GH#38980 groupby agg on extension type fails for non-numeric types
6563
df = pd.DataFrame({"A": [1, 1, 2, 2, 3, 3, 1, 4], "B": data_for_grouping})
6664

67-
expected_df = pd.DataFrame(
68-
{"A": [1, 1, 2, 2, 3, 3, 1, 4], "B": data_for_grouping}
69-
)
70-
expected = expected_df.iloc[[0, 2, 4, 7]]
65+
expected = df.iloc[[0, 2, 4, 7]]
7166
expected = expected.set_index("A")
72-
expected["B"] = expected["B"].astype(expected_inferred_result_dtype)
7367

7468
result = df.groupby("A").agg({"B": "first"})
7569
tm.assert_frame_equal(result, expected)

pandas/tests/extension/test_arrow.py

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -225,7 +225,6 @@ def data_for_grouping(dtype):
225225
return pd.array([B, B, None, None, A, A, B, C], dtype=dtype)
226226

227227

228-
@pytest.fixture
229228
def expected_inferred_result_dtype(dtype):
230229
"""
231230
When the data pass through aggregate,
@@ -1144,6 +1143,27 @@ def test_comp_masked_numpy(self, masked_dtype, comparison_op):
11441143
expected = pd.Series(exp, dtype=ArrowDtype(pa.bool_()))
11451144
tm.assert_series_equal(result, expected)
11461145

1146+
def test_groupby_agg_extension(self, data_for_grouping):
1147+
# GH#38980 groupby agg on extension type fails for non-numeric types
1148+
df = pd.DataFrame({"A": [1, 1, 2, 2, 3, 3, 1, 4], "B": data_for_grouping})
1149+
1150+
expected_df = pd.DataFrame(
1151+
{"A": [1, 1, 2, 2, 3, 3, 1, 4], "B": data_for_grouping}
1152+
)
1153+
expected = expected_df.iloc[[0, 2, 4, 7]]
1154+
expected = expected.set_index("A")
1155+
expected_dtype = expected_inferred_result_dtype(expected["B"].dtype)
1156+
expected["B"] = expected["B"].astype(expected_dtype)
1157+
1158+
result = df.groupby("A").agg({"B": "first"})
1159+
tm.assert_frame_equal(result, expected)
1160+
1161+
result = df.groupby("A").agg("first")
1162+
tm.assert_frame_equal(result, expected)
1163+
1164+
result = df.groupby("A").first()
1165+
tm.assert_frame_equal(result, expected)
1166+
11471167

11481168
class TestLogicalOps:
11491169
"""Various Series and DataFrame logical ops methods."""

0 commit comments

Comments
 (0)