Skip to content

Commit 6496b15

Browse files
author
Kei
committed
Update according to pr comments
1 parent 3a8597e commit 6496b15

File tree

2 files changed

+52
-301
lines changed

2 files changed

+52
-301
lines changed

pandas/tests/extension/test_arrow.py

+11-18
Original file line numberDiff line numberDiff line change
@@ -225,24 +225,6 @@ def data_for_grouping(dtype):
225225
return pd.array([B, B, None, None, A, A, B, C], dtype=dtype)
226226

227227

228-
def expected_inferred_result_dtype(dtype):
229-
"""
230-
When the data pass through aggregate,
231-
the inferred data type that it will become
232-
233-
"""
234-
235-
pa_dtype = dtype.pyarrow_dtype
236-
if pa.types.is_date(pa_dtype):
237-
return "date32[day][pyarrow]"
238-
elif pa.types.is_time(pa_dtype):
239-
return "time64[us][pyarrow]"
240-
elif pa.types.is_decimal(pa_dtype):
241-
return ArrowDtype(pa.decimal128(4, 3))
242-
else:
243-
return dtype
244-
245-
246228
@pytest.fixture
247229
def data_for_sorting(data_for_grouping):
248230
"""
@@ -1147,6 +1129,17 @@ def test_groupby_agg_extension(self, data_for_grouping):
11471129
# GH#38980 groupby agg on extension type fails for non-numeric types
11481130
df = pd.DataFrame({"A": [1, 1, 2, 2, 3, 3, 1, 4], "B": data_for_grouping})
11491131

1132+
def expected_inferred_result_dtype(dtype):
1133+
pa_dtype = dtype.pyarrow_dtype
1134+
if pa.types.is_date(pa_dtype):
1135+
return "date32[day][pyarrow]"
1136+
elif pa.types.is_time(pa_dtype):
1137+
return "time64[us][pyarrow]"
1138+
elif pa.types.is_decimal(pa_dtype):
1139+
return ArrowDtype(pa.decimal128(4, 3))
1140+
else:
1141+
return dtype
1142+
11501143
expected_df = pd.DataFrame(
11511144
{"A": [1, 1, 2, 2, 3, 3, 1, 4], "B": data_for_grouping}
11521145
)

0 commit comments

Comments
 (0)