[Python] Consistent handling of categoricals

[asfimport]

What is the current state of categoricals with pyarrow? The categories parameter mentioned in this GitHub issue does not seem to be accepted in pd.read_parquet anymore. I see that read/write of int categoricals does not work, though str do – except if the file is written by fastparquet.

Using pandas 1.1.5, pyarrow 2.0.0, and fastparquet 0.4.1, I see the following handling of categoricals:

 

import os
import pandas as pd


fname = '/tmp/tst'


data = {
    'int': pd.Series([0, 1] * 1000, dtype=pd.CategoricalDtype([0,1])),
    'str': pd.Series(['foo', 'bar'] * 1000, dtype=pd.CategoricalDtype(['foo', 'bar'])),
}
df = pd.DataFrame(data)


for write in ['fastparquet', 'pyarrow']:
    for read in ['fastparquet', 'pyarrow']:
        if os.path.exists(fname):
            os.remove(fname)
        df.to_parquet(fname, engine=write, compression=None)
        df_read = pd.read_parquet(fname, engine=read)


        print()
        print('write:', write, 'read:', read)
        for t in data.keys():
            print(t, df[t].dtype == df_read[t].dtype)

 

 

write: fastparquet read: fastparquet
int True
str True
write: fastparquet read: pyarrow
int False
str False
write: pyarrow read: fastparquet
int True
str True
write: pyarrow read: pyarrow
int False
str True

Reporter: Chris Roat

Related issues:

• [Python] Categorical boolean column saved as regular boolean in parquet (relates to)

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

[asfimport]

Wes McKinney / @wesm:
This looks buggy to me. Less consideration has been given to non-string categorical data since it appears less frequently in practice.

cc @jorisvandenbossche

[asfimport]

Joris Van den Bossche / @jorisvandenbossche:
Looking a bit into this. Converting the pandas DataFrame to pyarrow Table still preserves the categorical (dictionary) data type:

In [55]: table = pa.table(df)

In [56]: table
Out[56]: 
pyarrow.Table
int: dictionary<values=int64, indices=int8, ordered=0>
str: dictionary<values=string, indices=int8, ordered=0>

but when reading, we restore the dictionary type for the string column but not for the int column (as you already observed):

In [58]: pq.write_table(table, "test_categoricals.parquet")

In [59]: pq.write_table?

In [60]: pq.read_table("test_categoricals.parquet")
Out[60]: 
pyarrow.Table
int: int64
str: dictionary<values=string, indices=int32, ordered=0>

We store the original Arrow schema metadata, so we can restore the original types (since parquet doesn't have the concept of a separate dictionary/categorical type, only uses that as compression). The code where this is done lives here: https://github.com/apache/arrow/blob/b3b62412140182a0ac7b4f585fac27c8c57f1662/cpp/src/parquet/arrow/schema.cc#L874-L884
The IsDictionaryReadSupported in the if-check is basically checking if the data type is binary or string, so we only restore the dictionary type for those data types and not for integer (and thus when finally converting to pandas, you also don't see a categorical data type for the integer column).

I am not familiar enough with this part of the codebase to really know the reason for this, but I suppose it is related to ARROW-6140. For string/binary values, we can read the parquet data directly into an Arrow DictionaryArray, while this is not yet supported for other types. So if we would currently want to restore the integer columns also as dictionary type, that would basically mean to do a dictionary encoding after the fact of the materialized integer column (as long as ARROW-6140 is not implemented).

It is also documented that the read_dictionary keyword is (for this reason) only valid for string/bytes columns: https://arrow.apache.org/docs/python/parquet.html#reading-types-as-dictionaryarray

Now, the above is the actual parquet reading. If we know, based on the pandas dtypes information we store in the metadata, that the original column was a categorical column, we could also still preserve the categorical dtype on conversion from pyarrow -> pandas (but that's something we currently also don't do for eg string columns).

[asfimport]

Joris Van den Bossche / @jorisvandenbossche:
Additional answer specifically about this:

The categories parameter mentioned in this GitHub issue does not seem to be accepted in pd.read_parquet anymore

It certainly still works in the to_pandas method:

In [89]: pa.table({'int': [1, 2]}).to_pandas().dtypes
Out[89]: 
int    int64
dtype: object

In [90]: pa.table({'int': [1, 2]}).to_pandas(categories=['int']).dtypes
Out[90]: 
int    category
dtype: object

It might be that this worked at some point directly in the pd.read_parquet function, but indeed not at the moment. The main problem here is that additional keywords (kwargs) can be passed either to parquet.read_table or to Table.to_pandas. And right now they are passed to read_table. We should probably have a way to specify keywords for to_pandas as well.

As a workaround, you can read with pyarrow and do the conversion to pandas manually. So basically instead of pd.parquet(..) you can do pyarrow.parquet.read_table(..).to_pandas(..).

[asfimport]

A:
I need the following to work: 

Given a parquet file, automatically load it and convert to Pandas DataFrame with correct categorical columns (which are required by XGBoost for example).

The suggested workarounds do not work since this has to be done programmatically without prior knowledge of any schema - based on the parquet file alone.

[github-actions]

This issue has been marked as stale because it has had no activity in the past 365 days. Please remove the stale label or comment below, or this issue will be closed in 14 days. If this improvement is still desired but has no current owner, please add the 'Status: needs champion' label.