Possible to read categoricals back into Pandas from Parquet using Pyarrow?

[kylebarron]

Apologies if this is solved somewhere in the documentation; I've tried to look through it and through the issues here and on JIRA. If so just close this issue.

Is it possible to have a Pandas dataset with categorical variables, save it to Parquet, and read in those variables as categorical again?

Setup:

In [1]: import fastparquet as fp
In [2]: import pyarrow as pa
In [3]: import pyarrow.parquet as pq
In [4]: import pandas as pd
In [5]: df = pd.DataFrame({'A':['a','b','c','a']})
In [6]: df['B'] = df['A'].astype('category')
In [7]: df.dtypes
Out[7]: 
A      object
B    category
dtype: object

Using Pandas to read/write only reads B as categorical using fastparquet, but interestingly, reads that even if it was written by pyarrow:

In [11]: df.to_parquet('pa.parquet', engine='pyarrow')
    ...: df.to_parquet('fp.parquet', engine='fastparquet')

In [12]: pd.read_parquet('pa.parquet', engine='pyarrow').dtypes
Out[12]: 
A    object
B    object
dtype: object

In [13]: pd.read_parquet('pa.parquet', engine='fastparquet').dtypes
Out[13]: 
A      object
B    category
dtype: object

In [14]: pd.read_parquet('fp.parquet', engine='pyarrow').dtypes
Out[14]: 
A    object
B    object
dtype: object

In [15]: pd.read_parquet('fp.parquet', engine='fastparquet').dtypes
Out[15]: 
A      object
B    category
dtype: object

I'm guessing that fastparquet reads the pandas_type option in the metadata and sees categorical.

In [16]: pq.ParquetFile('pa.parquet').read(use_pandas_metadata=True)
Out[16]: 
pyarrow.Table
A: string
B: 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": "A", "field_name": "A", "pandas_type": "unicode", "nu'
            b'mpy_type": "object", "metadata": null}, {"name": "B", "field_nam'
            b'e": "B", "pandas_type": "categorical", "numpy_type": "int8", "me'
            b'tadata": {"num_categories": 3, "ordered": false}}, {"name": null'
            b', "field_name": "__index_level_0__", "pandas_type": "int64", "nu'
            b'mpy_type": "int64", "metadata": null}], "pandas_version": "0.22.'
            b'0"}'}

It seems that the strings_to_categorical option of pa.Table.to_pandas() doesn't work in this situation either (maybe I'm using it wrong; also I'd prefer to only read as categorical the columns that were originally categorical in Pandas):

In [17]: table = pq.ParquetFile('pa.parquet').read(use_pandas_metadata=True)
    ...: table.to_pandas(strings_to_categorical=True)
    ...: 
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-17-42387e24e864> in <module>()
      1 table = pq.ParquetFile('pa.parquet').read(use_pandas_metadata=True)
----> 2 table.to_pandas(strings_to_categorical=True)

~/local/anaconda3/lib/python3.6/site-packages/pyarrow/table.pxi in pyarrow.lib.Table.to_pandas (/arrow/python/build/temp.linux-x86_64-3.6/lib.cxx:46331)()

~/local/anaconda3/lib/python3.6/site-packages/pyarrow/pandas_compat.py in table_to_blockmanager(options, table, memory_pool, nthreads, categoricals)
    526             )
    527 
--> 528     blocks = _table_to_blocks(options, block_table, nthreads, memory_pool)
    529 
    530     # Construct the row index

~/local/anaconda3/lib/python3.6/site-packages/pyarrow/pandas_compat.py in _table_to_blocks(options, block_table, nthreads, memory_pool)
    620 
    621     # Defined above
--> 622     return [_reconstruct_block(item) for item in result]
    623 
    624 

~/local/anaconda3/lib/python3.6/site-packages/pyarrow/pandas_compat.py in <listcomp>(.0)
    620 
    621     # Defined above
--> 622     return [_reconstruct_block(item) for item in result]
    623 
    624 

~/local/anaconda3/lib/python3.6/site-packages/pyarrow/pandas_compat.py in _reconstruct_block(item)
    434         cat = pd.Categorical.from_codes(block_arr,
    435                                         categories=item['dictionary'],
--> 436                                         ordered=item['ordered'])
    437         block = _int.make_block(cat, placement=placement,
    438                                 klass=_int.CategoricalBlock,

~/local/anaconda3/lib/python3.6/site-packages/pandas/core/categorical.py in from_codes(cls, codes, categories, ordered)
    619 
    620         if len(codes) and (codes.max() >= len(categories) or codes.min() < -1):
--> 621             raise ValueError("codes need to be between -1 and "
    622                              "len(categories)-1")
    623 

ValueError: codes need to be between -1 and len(categories)-1

It seems that the pd.DataFrame > pa.Table > pd.DataFrame conversion keeps categoricals as a dictionary Arrow type, but that the pa.Table > .parquet > pa.Table process loses the categoricals.

In [22]: table = pa.Table.from_pandas(df)

In [23]: table.schema
Out[23]: 
A: string
B: dictionary<values=string, indices=int8, ordered=0>
  dictionary: ["a", "b", "c"]
__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": "A", "field_name": "A", "pandas_type": "unicode", "nu'
            b'mpy_type": "object", "metadata": null}, {"name": "B", "field_nam'
            b'e": "B", "pandas_type": "categorical", "numpy_type": "int8", "me'
            b'tadata": {"num_categories": 3, "ordered": false}}, {"name": null'
            b', "field_name": "__index_level_0__", "pandas_type": "int64", "nu'
            b'mpy_type": "int64", "metadata": null}], "pandas_version": "0.22.'
            b'0"}'}

In [24]: table.to_pandas().dtypes
Out[24]: 
A      object
B    category
dtype: object

In [25]: pq.write_table(table, 'table.parquet')

In [26]: pq.ParquetFile('table.parquet').read()
Out[26]: 
pyarrow.Table
A: string
B: 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": "A", "field_name": "A", "pandas_type": "unicode", "nu'
            b'mpy_type": "object", "metadata": null}, {"name": "B", "field_nam'
            b'e": "B", "pandas_type": "categorical", "numpy_type": "int8", "me'
            b'tadata": {"num_categories": 3, "ordered": false}}, {"name": null'
            b', "field_name": "__index_level_0__", "pandas_type": "int64", "nu'
            b'mpy_type": "int64", "metadata": null}], "pandas_version": "0.22.'
            b'0"}'}

Using fastparquet works well at the moment, but

1. Some files written by fastparquet are unable to be written by pyarrow. I get the trace

   In [3]: df = pf.read()
   ---------------------------------------------------------------------------
   ArrowIOError                              Traceback (most recent call last)
   <ipython-input-3-910db268c9c9> in <module>()
   ----> 1 df = pf.read()
   
   ~/local/anaconda3/lib/python3.6/site-packages/pyarrow/parquet.py in read(self, columns, nthreads, use_pandas_metadata)
       140             columns, use_pandas_metadata=use_pandas_metadata)
       141         return self.reader.read_all(column_indices=column_indices,
   --> 142                                     nthreads=nthreads)
       143 
       144     def scan_contents(self, columns=None, batch_size=65536):
   
   ~/local/anaconda3/lib/python3.6/site-packages/pyarrow/_parquet.pyx in pyarrow._parquet.ParquetReader.read_all (/arrow/python/build/temp.linux-x86_64-3.6/_parquet.cxx:12865)()
   
   ~/local/anaconda3/lib/python3.6/site-packages/pyarrow/error.pxi in pyarrow.lib.check_status (/arrow/python/build/temp.linux-x86_64-3.6/lib.cxx:8345)()
   
   ArrowIOError: Unexpected end of stream.

   I'm working with restricted data, but I could try to reproduce this error with sample data.
2. I'm trying to write a wrapper for the Parquet C++ library that would read/write Parquet files to Stata, and so I'd need all my files to be readable by pyarrow.
3. Fastparquet is single threaded so I suppose pyarrow could be faster when using multiple cores

This was done using versions:

In [18]: pa.__version__
Out[18]: '0.8.0'

In [20]: fp.__version__
Out[20]: '0.1.4'

In [21]: pd.__version__
Out[21]: '0.22.0'

[wesm]

@kylebarron the Parquet format has dictionary encoding as a compression strategy, but does not support categoricals in the way that pandas does.

The fastparquet library does support returning the dictionary pages as though they were categorical, but there are ways that things can fail:

• Dictionaries may differ across row groups, files
• A column chunk may switch from dictionary to plain encoding mid-stream if the dictionary got too big

We have planned to implement a best-efforts direct-to-categorical read like fastparquet with the caveat that it may not always work. This hasn't been implemented yet, though.

It looks like there's some other bugs happening here that should be investigated. It looks like you may have hit dask/fastparquet#305

[xhochy]

@kylebarron The best way to deal with this at the moment is to use on to_pandas()/read_parquet() the option strings_to_categorical=True* (available in the released pyarrow version) that will turn all string columns to categoricals or the option categories=[…] that will convert all listed columns into categoricals. The latter is only available in git but will be part of the Arrow 0.9 release.

These are not the perfect options but the most pragmatic ones currently.

• As @wesm already pointed out, the error you are seeing is ColumnMetaData.total_(un)compressed_size does not take DictionaryPage into account dask/fastparquet#305, thus if that file was written using pyarrow or the latest fastparquet version it should vanish and strings_to_categorical should work.

[kylebarron]

Thanks for your replies.

@xhochy the error I got from strings_to_categoricals=True came from the dataset saved by pyarrow 0.8, not the one saved by fastparquet. Concisely:

In [1]: import pyarrow as pa

In [2]: import pyarrow.parquet as pq

In [3]: import pandas as pd

In [4]: df = pd.DataFrame({'A':['a','b','c','a']})

In [5]: df['B'] = df['A'].astype('category')

In [6]: df.dtypes
Out[6]: 
A      object
B    category
dtype: object

In [7]: df.to_parquet('pa.parquet', engine='pyarrow')

In [8]: table = pq.ParquetFile('pa.parquet').read(use_pandas_metadata=True)

In [9]: table.to_pandas(strings_to_categorical=True)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-9-fff9603d5df6> in <module>()
----> 1 table.to_pandas(strings_to_categorical=True)

~barron/anaconda/lib/python3.6/site-packages/pyarrow/table.pxi in pyarrow.lib.Table.to_pandas (/Users/travis/build/BryanCutler/arrow-dist/arrow/python/build/temp.macosx-10.6-intel-3.6/lib.cxx:46331)()

~barron/anaconda/lib/python3.6/site-packages/pyarrow/pandas_compat.py in table_to_blockmanager(options, table, memory_pool, nthreads, categoricals)
    526             )
    527 
--> 528     blocks = _table_to_blocks(options, block_table, nthreads, memory_pool)
    529 
    530     # Construct the row index

~barron/anaconda/lib/python3.6/site-packages/pyarrow/pandas_compat.py in _table_to_blocks(options, block_table, nthreads, memory_pool)
    620 
    621     # Defined above
--> 622     return [_reconstruct_block(item) for item in result]
    623 
    624 

~barron/anaconda/lib/python3.6/site-packages/pyarrow/pandas_compat.py in <listcomp>(.0)
    620 
    621     # Defined above
--> 622     return [_reconstruct_block(item) for item in result]
    623 
    624 

~barron/anaconda/lib/python3.6/site-packages/pyarrow/pandas_compat.py in _reconstruct_block(item)
    434         cat = pd.Categorical.from_codes(block_arr,
    435                                         categories=item['dictionary'],
--> 436                                         ordered=item['ordered'])
    437         block = _int.make_block(cat, placement=placement,
    438                                 klass=_int.CategoricalBlock,

~barron/anaconda/lib/python3.6/site-packages/pandas/core/categorical.py in from_codes(cls, codes, categories, ordered)
    619 
    620         if len(codes) and (codes.max() >= len(categories) or codes.min() < -1):
--> 621             raise ValueError("codes need to be between -1 and "
    622                              "len(categories)-1")
    623 

ValueError: codes need to be between -1 and len(categories)-1

In [9]: pa.__version__
Out[9]: '0.8.0'

[xhochy]

This is fixed on master and will then also be part of the shortly upcoming 0.9.0 release.

[Demetrio92]

Stumbled across this bug again.
pyarrow preserves category as dtype, fastparquet does not.

Docs don't mention it. They even kinda mislead the users:
https://pandas.pydata.org/pandas-docs/stable/user_guide/io.html#io-parquet

^ this works, because the second option has engine='pyarrow'!!!
So, the default auto, which uses fastparquet would actually result in those being object.

This took me a bit, before I noticed where and why do I have memory leaks despite optimizing my dataframes...

---

pd.__version__
Out[3]: '1.0.4'

[wesm]

Just to confirm, this project isn't doing anything wrong?

[Demetrio92]

@wesm yeah, sorry, guys, you're awesome. I thought this was pandas repo...

[chrisroat]

What is the current state of categoricals with pyarrow? The categories parameter mentioned above 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

[superChing]

@chrisroat
FYI, this is documented in Pandas

In the pyarrow engine, categorical dtypes for non-string types can be serialized to parquet, but will de-serialize as their primitive dtype.

[Demetrio92]

Seems like the issue is back. But the guys are working on it.

https://issues.apache.org/jira/browse/ARROW-11157

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(..).

---

My 2c: last summer I had a dataset which wouldn't fit into RAM with categories stored as strings in pandas. This pretty much prevented me from writing and reading it as parquet using fastparquet. PyArrow literally saved the day by properly handling categories.

Would be really nice to see this working again.