Pandas DataFrame 使用 read_parquet 来自 Parquet 文件的分类列？

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 进行进一步分析时，我似乎无法恢复类别数据类型。以下

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

生成 DataFrame，用 object dtypes 代替所需的 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

我们遇到了类似的问题。使用多文件镶木地板时，解决方法如下：使用 Table.to_pandas() documentation 以下代码可能是相关的：

import pyarrow.parquet as pq
dft = pq.read_table('path/to/data_parquet/', use_pandas_metadata=True)
df = dft.to_pandas(categories=['column_2'] )

use_panadas_metadata 适用于 dtype datetime64[ns]

Answer 2

这已在 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

Pandas DataFrame 使用 read_parquet 来自 Parquet 文件的分类列？

Pandas DataFrame with categorical columns from a Parquet file using read_parquet?

python-3.x

pandas

parquet

pyarrow