如何使用pyarrow为ParquetDataset存储自定义元数据

你设想的操作流程是完全可行的，ParquetDataset的全局元数据默认存储在数据集根目录下的_common_metadata文件中，只需修改该文件的元数据即可实现需求，完整实现代码如下：

import pyarrow.parquet as pq

# 读取现有公共元数据
meta = pq.read_metadata('temp.parq/_common_metadata')

# 定义自定义元数据，注意键值必须为bytes类型
custom_metadata = {b'type': b'mydataset'}

# 合并元数据，此处将自定义元数据放在后位，保证同名字段自定义值优先覆盖原有值
merged_metadata = {**meta.metadata, **custom_metadata}

# 生成新的FileMetaData对象，替换元数据
new_meta = meta.replace_metadata(merged_metadata)

# 写入新的公共元数据文件
new_meta.write_metadata_file('temp.parq/_common_metadata')

验证元数据是否写入成功 #

读取数据集后直接打印schema的元数据即可验证：

dataset = pq.ParquetDataset('temp.parq')
print(dataset.schema.metadata.get(b'type'))  # 输出应为b'mydataset'

特殊场景处理：数据集无默认_common_metadata文件 #

部分生成Parquet数据集的工具不会默认生成公共元数据文件，可通过数据集schema直接生成：

import pyarrow.parquet as pq

dataset = pq.ParquetDataset('temp.parq')
custom_metadata = {b'type': b'mydataset'}

# 给schema添加自定义元数据
schema_with_meta = dataset.schema.with_metadata(custom_metadata)

# 直接写入公共元数据文件
pq.write_metadata(schema_with_meta, 'temp.parq/_common_metadata')

额外优化建议 #

如果是使用Dask生成Parquet数据集，可在写入时显式指定write_metadata_file=True，会自动生成_common_metadata文件，无需后续手动创建：

dask.datasets.timeseries().to_parquet('temp.parq', write_metadata_file=True)

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