diff --git a/pandas/tests/reshape/test_pivot.py b/pandas/tests/reshape/test_pivot.py index 582084e3bfb5a..a8386d21ba27f 100644 --- a/pandas/tests/reshape/test_pivot.py +++ b/pandas/tests/reshape/test_pivot.py @@ -1656,6 +1656,51 @@ def test_categorical_margins_category(self, observed): table = df.pivot_table("x", "y", "z", dropna=observed, margins=True) tm.assert_frame_equal(table, expected) + def test_pivot_with_categorical(self, observed, ordered_fixture): + # gh-21370 + idx = [np.nan, "low", "high", "low", np.nan] + col = [np.nan, "A", "B", np.nan, "A"] + df = pd.DataFrame( + { + "In": pd.Categorical( + idx, categories=["low", "high"], ordered=ordered_fixture + ), + "Col": pd.Categorical( + col, categories=["A", "B"], ordered=ordered_fixture + ), + "Val": range(1, 6), + } + ) + # case with index/columns/value + result = df.pivot_table( + index="In", columns="Col", values="Val", observed=observed + ) + + expected_cols = pd.CategoricalIndex( + ["A", "B"], ordered=ordered_fixture, name="Col" + ) + + expected = pd.DataFrame( + data=[[2.0, np.nan], [np.nan, 3.0]], columns=expected_cols + ) + expected.index = Index( + pd.Categorical( + ["low", "high"], categories=["low", "high"], ordered=ordered_fixture + ), + name="In", + ) + + tm.assert_frame_equal(result, expected) + + # case with columns/value + result = df.pivot_table(columns="Col", values="Val", observed=observed) + + expected = pd.DataFrame( + data=[[3.5, 3.0]], columns=expected_cols, index=Index(["Val"]) + ) + + tm.assert_frame_equal(result, expected) + def test_categorical_aggfunc(self, observed): # GH 9534 df = pd.DataFrame(