使用PyArrow写入Parquet分区数据集时分区列丢失的解决办法

问题原因 #

当使用pyarrow的write_to_dataset按year分区写入Parquet文件时，分区列的值会被存储在文件路径的目录结构中（例如dataset_name_3/year=2019/xxx.parquet），而不会写入单个Parquet文件的内容里。直接读取单个文件时，自然无法获取到分区列数据。

解决办法 #

方法1：读取整个分区数据集（推荐） #

通过ParquetDataset读取整个数据集，PyArrow会自动从目录路径中解析出分区列并添加到结果中：

import pyarrow.parquet as pq
import pandas as pd

# 读取整个分区数据集
dataset = pq.ParquetDataset('dataset_name_3', use_legacy_dataset=False)
table = dataset.read()
df = table.to_pandas()

# 或者直接用pandas读取整个数据集路径
df = pd.read_parquet('dataset_name_3', engine='pyarrow')

这样得到的DataFrame会包含year列，且对应每个分区的正确值。

方法2：手动从文件路径提取分区列（仅适用于读取单个文件） #

如果必须读取单个文件，可以解析文件路径中的分区信息，手动添加到DataFrame：

import pandas as pd
import os

file_path = p_files[0]
# 从路径中提取year值（按目录结构拆分）
parts = file_path.split(os.sep)
year_part = [p for p in parts if p.startswith('year=')][0]
year_value = int(year_part.split('=')[1])

# 读取文件并添加year列
df = pd.read_parquet(file_path)
df['year'] = year_value

执行后得到的DataFrame就会包含year列，值为对应分区的年份。

内容的提问来源于stack exchange，提问作者Swavig