为分区Parquet数据集写入pyarrow元数据遇Schema不匹配错误的解决方法

当你指定partition_cols=["n_legs"]时，PyArrow会把n_legs列从每个输出的Parquet文件中移除，仅将其作为分区目录的标识（比如生成n_legs=2/、n_legs=4/这类目录）。因此每个分区文件的schema只有animal这1列，但你写入元数据时用的是原table的schema（包含n_legs和animal共2列），两者不匹配就触发了AppendRowGroups requires equal schemas错误。

有两种可靠的解决方法：

方法1：手动生成移除分区列后的schema #

直接从原schema中剔除分区列，用这个精简后的schema来写元数据：

import pyarrow as pa
import pyarrow.parquet as pq
from pathlib import Path

table = pa.table({'n_legs': [2, 2, 4, 4, 5, 100],
                  'animal': ["Flamingo", "Parrot", "Dog", "Horse",
                             "Brittle stars", "Centipede"]})

metadata_collector = []
root_path = Path("animals.parquet")
pq.write_to_dataset(
    table,
    root_path,
    partition_cols=["n_legs"],
    metadata_collector=metadata_collector,
)

# 生成移除分区列后的schema
partitioned_schema = table.schema.remove(["n_legs"])

# 写入_common_metadata
pq.write_metadata(partitioned_schema, root_path / "_common_metadata")

# 写入带统计信息的_metadata
pq.write_metadata(
    partitioned_schema, root_path / "_metadata", metadata_collector=metadata_collector
)

方法2：从metadata_collector中提取schema（更推荐） #

metadata_collector里已经存储了每个分区文件的schema信息，直接取第一个元素的schema即可（所有分区文件schema一致），无需手动修改：

import pyarrow as pa
import pyarrow.parquet as pq
from pathlib import Path

table = pa.table({'n_legs': [2, 2, 4, 4, 5, 100],
                  'animal': ["Flamingo", "Parrot", "Dog", "Horse",
                             "Brittle stars", "Centipede"]})

metadata_collector = []
root_path = Path("animals.parquet")
pq.write_to_dataset(
    table,
    root_path,
    partition_cols=["n_legs"],
    metadata_collector=metadata_collector,
)

# 从metadata_collector中获取正确的schema
partitioned_schema = metadata_collector[0].schema

# 写入_common_metadata
pq.write_metadata(partitioned_schema, root_path / "_common_metadata")

# 写入带统计信息的_metadata
pq.write_metadata(
    partitioned_schema, root_path / "_metadata", metadata_collector=metadata_collector
)

两种方法都能解决schema不匹配的问题，方法2更灵活，尤其是当分区列较多或schema动态变化时，不需要手动调整。

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