BUG: MultiIndex loses category after .stack()

[maroth96]

• I have checked that this issue has not already been reported.

• I have confirmed this bug exists on the latest version of pandas.

• (optional) I have confirmed this bug exists on the master branch of pandas.

---

Code Sample, a copy-pastable example

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

Problem description

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

Expected Output

Output of pd.show_versions()

INSTALLED VERSIONS

commit : db08276
python : 3.7.3.final.0
python-bits : 64
OS : Darwin
OS-release : 19.6.0
Version : Darwin Kernel Version 19.6.0: Mon Aug 31 22:12:52 PDT 2020; root:xnu-6153.141.2~1/RELEASE_X86_64
machine : x86_64
processor : i386
byteorder : little
LC_ALL : None
LANG : en_US.UTF-8
LOCALE : en_US.UTF-8

pandas : 1.1.3
numpy : 1.18.4
pytz : 2020.1
dateutil : 2.8.1
pip : 20.2.3
setuptools : 40.8.0
Cython : None
pytest : None
hypothesis : None
sphinx : None
blosc : None
feather : None
xlsxwriter : None
lxml.etree : None
html5lib : None
pymysql : None
psycopg2 : None
jinja2 : 2.11.2
IPython : 7.16.1
pandas_datareader: None
bs4 : None
bottleneck : None
fsspec : 0.7.4
fastparquet : None
gcsfs : None
matplotlib : 3.3.0
numexpr : None
odfpy : None
openpyxl : None
pandas_gbq : None
pyarrow : None
pytables : None
pyxlsb : None
s3fs : None
scipy : 1.5.2
sqlalchemy : None
tables : None
tabulate : None
xarray : None
xlrd : None
xlwt : None
numba : None

[rhshadrach]

Thanks for the report! Confirmed on master, further investigations and PRs are welcome.

[maroth96]

The problem is in _stack_multi_columns (core/reshape/reshape.py)

The following line causes the dtypes to be lost:

new_columns = MultiIndex.from_tuples(unique_groups, names=new_names)

The most straightforward solution is simply to add:

new_columns = new_columns.set_levels([
    new_columns.levels[i].astype(this.columns.levels[i].dtype)
    for i in range(0, new_columns.size)
])

This causes my assertions to succeed. The additional conversions are not so elegant, however.

@rhshadrach Shall I create such a change, or you see a better solution?

[rhshadrach]

@maroth96 - It's not clear to me if there is a better way to preserve dtype; I recommend opening a PR with this and seeing if others have suggestions.