[Python] CategoricalIndex is lost after reading back

[asfimport]

When a CategoricalIndex is written and read back the resulting index is not more categorical.

df = pd.DataFrame([['a', 'b'], ['c', 'd']], columns=['c1', 'c2'])
df['c1'] = df['c1'].astype('category')
df = df.set_index(['c1'])

table = pa.Table.from_pandas(df)
pq.write_table(table, 'test.parquet')

ref_df = pq.read_pandas('test.parquet').to_pandas()

print(df.index)
# CategoricalIndex(['a', 'c'], categories=['a', 'c'], ordered=False, name='c1', dtype='category')

print(ref_df.index)
# Index(['a', 'c'], dtype='object', name='c1')

In the metadata the information is correctly contained:

{"name": "c1", "field_name": "c1", "p'
 b'andas_type": "categorical", "numpy_type": "int8", "metadata": {"'
 b'num_categories": 2, "ordered": false}

 

Reporter: Armin Berres
Assignee: Wes McKinney / @wesm

Related issues:

• [Python] Support reading Parquet binary/string columns directly as DictionaryArray (relates to)
• [C++] Read Parquet dictionary encoded ColumnChunks directly into an Arrow DictionaryArray (relates to)
• [Python][Parquet] direct reading/writing of pandas categoricals in parquet (relates to)
• [Python] Pandas categorical type doesn't survive a round-trip through parquet (relates to)

PRs and other links:

• GitHub Pull Request #5117

_{Note: This issue was originally created as ARROW-3652. Please see the migration documentation for further details.}

[asfimport]

Matthew Rocklin / @mrocklin:
This is lost for categorical columns generally, not only the index

import pandas as pd
df = pd.DataFrame({'x': pd.Categorical(['a'])})
df.to_parquet('df.parquet')

import pyarrow.parquet as pq
table = pq.read_table('df.parquet', use_pandas_metadata=True)

>>> df.dtypes
x category
dtype: object

>>> table.to_pandas().dtypes
x object
dtype: object

>>> table
pyarrow.Table
x: string
__index_level_0__: int64
metadata
--------
{b'pandas': b'{"index_columns": ["__index_level_0__"], "column_indexes": [{"na'
 b'me": null, "field_name": null, "pandas_type": "unicode", "numpy_'
 b'type": "object", "metadata": {"encoding": "UTF-8"}}], "columns":'
 b' [{"name": "x", "field_name": "x", "pandas_type": "categorical",'
 b' "numpy_type": "int8", "metadata": {"num_categories": 1, "ordere'
 b'd": false}}, {"name": null, "field_name": "__index_level_0__", "'
 b'pandas_type": "int64", "numpy_type": "int64", "metadata": null}]'
 b', "pandas_version": "0.23.4"}'}

[asfimport]

Wes McKinney / @wesm:
As has been discussed in other issues, this is a somewhat complex topic:

• Parquet does not have a categorical type

• Dictionary encoding is used in general for compression only

• A single RowGroup may stop using dictionary encoding mid-stream if the dictionary grows large

• The dictionary will be different from row group to roup group in general

  In the case where both the writer and the reader is pandas, we can add pandas-specific metadata to indicate that the same dictionary has been written in each row group, and that all the data pages are dictionary encoded, to avoid the issues above. It's a doable project but with a steep learning curve, requiring comfort level with multiple layers of the C++ library and the Arrow conversion layer.

[asfimport]

Matthew Rocklin / @mrocklin:
Right, I've just done a search on this topic and found a few other issues. Happy to move conversation to one of them if desired.

In this particular case we already have pandas-specific metadata telling us that we have a categorical column (this is in the pandas_type entry for the column). As you say we don't know the category-mapping, but we could still create a Categorical column and just let Pandas decide the mapping. I suspect that it sorts the entries as a sane default.. This wouldn't necessarily roundtrip the right category-to-code mapping, but we can at least roundtrip the fact that these should be represented as categoricals.

Speaking from a Dask Dataframe perspective this would be desirable. We're accustomed to having to remap mismatched categorical columns to make things match up. Pandas generally handles this case well today. However, I can understand if Arrow wants to play things safe here and resist the temptation to let Pandas be magical.

 

So, the concrete ask here is that when converting a table to a Pandas dataframe and using the use_pandas_metadata= keyword that we convert columns that Pandas marked with pandas_type='category' to be categorical columns.

[asfimport]

Wes McKinney / @wesm:
Makes sense. We've recently done a bunch of work to improve the hashing machinery on the C++ side (for dictionary encoding, or whatever purpose). This has been a common enough issue / ask that I think it would be a good idea to tackle this in the next couple of months, so for the Arrow 0.13 release (0.12 is right around the corner)

[asfimport]

Wes McKinney / @wesm:
I'm pretty close to being able to take care of this one finally. There are some Parquet-related oddities such as dealing with "large" categories if we want to accurately preserve categories in all cases.

[asfimport]

Wes McKinney / @wesm:
This is working in master. I'm going to add a unit test

[asfimport]

Wes McKinney / @wesm:
Issue resolved by pull request 5117
#5117