使用 read_parquet 的 Parquet 文件中带有分类列的 Pandas DataFrame？

Question

我正在将大型 CSV 文件转换为 Parquet 文件以供进一步分析。我将 CSV 数据读入 Pandas 并按dtypes如下方式指定列

_dtype = {"column_1": "float64",
          "column_2": "category",
          "column_3": "int64",
          "column_4": "int64"}

df = pd.read_csv("data.csv", dtype=_dtype)

然后我做一些更多的数据清理并将数据写入 Parquet 以供下游使用。

_parquet_kwargs = {"engine": "pyarrow",
                   "compression": "snappy",
                   "index": False}

df.to_parquet("data.parquet", **_parquet_kwargs)

但是当我使用 Pandas 将数据读入进行进一步分析时，from_parquet我似乎无法恢复类别 dtypes。下列

df = pd.read_parquet("data.parquet")

结果DataFrame用objectdtypes 代替了所需的category.

以下似乎按预期工作

import pyarrow.parquet as pq

_table = (pq.ParquetFile("data.parquet")
            .read(use_pandas_metadata=True))

df = _table.to_pandas(strings_to_categorical=True)

但是我想知道如何使用pd.read_parquet.

Answer 1

这是在 中修复的Arrow 0.15，现在下一个代码将列保留为类别（并且性能明显更快）：

import pandas

df = pandas.DataFrame({'foo': list('aabbcc'),
                       'bar': list('xxxyyy')}).astype('category')

df.to_parquet('my_file.parquet')
df = pandas.read_parquet('my_file.parquet')
df.dtypes