BUG: MultiIndex loses category after .stack() · Issue #36991 · pandas-dev/pandas (original) (raw)

import pandas as pd


columns = pd.MultiIndex.from_tuples(
    [('x', 1, 3), ('x', 2, 3), ('y', 1, 3), ('y', 2, 3)], names=['a', 'b', 'c'])
columns = columns.set_levels(
    [columns.levels[i].astype('category') for i in range(0, 2)], level=[0, 1])


df = pd.DataFrame([[1, 2, 3, 4]], columns=columns)
dtypes = df.columns.to_frame().dtypes

assert isinstance(dtypes.a, pd.CategoricalDtype)
assert isinstance(dtypes.b, pd.CategoricalDtype)


df2 = df.stack(['a', 'b'])
dtypes2 = df2.index.to_frame().dtypes

assert isinstance(dtypes2.a, pd.CategoricalDtype)
assert isinstance(dtypes2.b, pd.CategoricalDtype)  # broken

Categorical types within a MultiIndex should be preserved after calling .stack().