Skip to content

Commit 096f6eb

Browse files
CLN: Refactor tests to adress PR feedback
1 parent 1600a1d commit 096f6eb

File tree

2 files changed

+27
-24
lines changed

2 files changed

+27
-24
lines changed

pandas/tests/reshape/test_pivot.py

+14-12
Original file line numberDiff line numberDiff line change
@@ -2858,22 +2858,24 @@ def test_pivot_with_pyarrow_categorical(self):
28582858
# GH#53051
28592859
pa = pytest.importorskip("pyarrow")
28602860

2861-
data = {"string_column": ["A", "B", "C"], "number_column": [1, 2, 3]}
2862-
df = (
2863-
DataFrame(data)
2864-
.astype({"string_column": "category", "number_column": "float32"})
2865-
.astype(
2866-
{
2867-
"string_column": ArrowDtype(pa.dictionary(pa.int32(), pa.string())),
2868-
"number_column": "float[pyarrow]",
2869-
}
2870-
)
2861+
df = DataFrame(
2862+
{"string_column": ["A", "B", "C"], "number_column": [1, 2, 3]}
2863+
).astype(
2864+
{
2865+
"string_column": ArrowDtype(pa.dictionary(pa.int32(), pa.string())),
2866+
"number_column": "float[pyarrow]",
2867+
}
28712868
)
28722869

28732870
df = df.pivot(columns=["string_column"], values=["number_column"])
28742871

2875-
df_expected = DataFrame(data).pivot(
2876-
columns=["string_column"], values=["number_column"]
2872+
multi_index = MultiIndex.from_arrays(
2873+
[["number_column", "number_column", "number_column"], ["A", "B", "C"]],
2874+
names=(None, "string_column"),
2875+
)
2876+
df_expected = DataFrame(
2877+
[[1.0, np.nan, np.nan], [np.nan, 2.0, np.nan], [np.nan, np.nan, 3.0]],
2878+
columns=multi_index,
28772879
)
28782880
tm.assert_frame_equal(
28792881
df, df_expected, check_dtype=False, check_column_type=False

pandas/tests/test_multilevel.py

+13-12
Original file line numberDiff line numberDiff line change
@@ -324,26 +324,27 @@ def test_multiindex_with_pyarrow_categorical(self):
324324
# GH#53051
325325
pa = pytest.importorskip("pyarrow")
326326

327-
data = {"string_column": ["A", "B", "C"], "number_column": [1, 2, 3]}
328-
df = (
329-
DataFrame(data)
330-
.astype({"string_column": "category", "number_column": "float32"})
331-
.astype(
332-
{
333-
"string_column": ArrowDtype(pa.dictionary(pa.int32(), pa.string())),
334-
"number_column": "float[pyarrow]",
335-
}
336-
)
327+
df = DataFrame(
328+
{"string_column": ["A", "B", "C"], "number_column": [1, 2, 3]}
329+
).astype(
330+
{
331+
"string_column": ArrowDtype(pa.dictionary(pa.int32(), pa.string())),
332+
"number_column": "float[pyarrow]",
333+
}
337334
)
338335

339336
df = df.set_index(["string_column", "number_column"])
340337

341-
df_expected = DataFrame(data).set_index(["string_column", "number_column"])
338+
df_expected = DataFrame(
339+
index=MultiIndex.from_arrays(
340+
[["A", "B", "C"], [1, 2, 3]], names=["string_column", "number_column"]
341+
)
342+
)
342343
tm.assert_frame_equal(
343344
df,
344345
df_expected,
345-
check_dtype=False,
346346
check_index_type=False,
347+
check_column_type=False,
347348
)
348349

349350

0 commit comments

Comments
 (0)