如何使用 Pandas 编写分区的 Parquet 文件

Question

我正在尝试将 Pandas 数据帧写入分区文件：

df.to_parquet('output.parquet', engine='pyarrow', partition_cols = ['partone', 'partwo'])

TypeError: __cinit__() got an unexpected keyword argument 'partition_cols'

从文档中我预计partition_cols将作为 kwargs 传递给 pyarrow 库。如何使用 Pandas 将分区文件写入本地磁盘？

Answer 1

大熊猫DataFrame.to_parquet是薄的包装上table = pa.Table.from_pandas(...)和pq.write_table(table, ...)（见pandas.parquet.py#L120），并且pq.write_table不支持写入分区的数据集。你应该pq.write_to_dataset改用。

import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq

df = pd.DataFrame(yourData)
table = pa.Table.from_pandas(df)

pq.write_to_dataset(
    table,
    root_path='output.parquet',
    partition_cols=['partone', 'parttwo'],
)

有关更多信息，请参阅pyarrow 文档。

通常，在读取/写入镶木地板文件时，我总是直接使用 PyArrow API，因为 Pandas 包装器的功能相当有限。

Answer 2

您需要更新到 Pandas 0.24 或更高版本。partition_cols 的功能是从该版本开始添加的。

Answer 3

首先确保你有一个相当新的 pandas 和 pyarrow 版本：

pyenv shell 3.8.2
python -m venv venv
source venv/bin/activate
pip install pandas pyarrow
pip freeze | grep pandas # pandas==1.2.3
pip freeze | grep pyarrow # pyarrow==3.0.0

然后您可以使用partition_cols生成分区的镶木地板文件：

import pandas as pd

# example dataframe with 3 rows and columns year,month,day,value
df = pd.DataFrame(data={'year':  [2020, 2020, 2021],
                        'month': [1,12,2], 
                        'day':   [1,31,28], 
                        'value': [1000,2000,3000]})

df.to_parquet('./mydf', partition_cols=['year', 'month', 'day'])

这产生：

mydf/year=2020/month=1/day=1/6f0258e6c48a48dbb56cae0494adf659.parquet
mydf/year=2020/month=12/day=31/cf8a45116d8441668c3a397b816cd5f3.parquet
mydf/year=2021/month=2/day=28/7f9ba3f37cb9417a8689290d3f5f9e6e.parquet