如何在Python Apache Beam中自动获取Parquet Schema用于写入？

现有基于Python的Apache Beam管道流程：读取Parquet文件→转为Pandas DataFrame清洗→转回Parquet格式写入，但执行时因WriteToParquet缺少必填的schema参数报错。

管道代码 #

with beam.Pipeline(options=pipeline_options) as p:
    dataframes = p \
      | 'Read' >> beam.io.ReadFromParquetBatched(known_args.input) \
      | 'Convert to pandas' >> beam.Map(lambda table: table.to_pandas()) \
      | 'Process df' >> beam.ParDo(ProcessDataFrame()) \
      | 'Convert to parquet' >> beam.Map(lambda table: table.to_parquet()) \
      | 'Write to parquet' >> beam.io.WriteToParquet(known_args.output)

错误信息 #

Traceback (most recent call last):
  File "/Users/kgallatin/dataflow/example.py", line 75, in <module>
    main()
  File "/Users/kgallatin/dataflow/example.py", line 70, in main
    | 'Write to parquet' >> beam.io.WriteToParquet(known_args.output)
TypeError: __init__() missing 1 required positional argument: 'schema'

因数据列数多且结构可能变化，无需手动编写PyArrow Schema，可通过以下两种方式自动提取：

方案1：从原始Parquet文件提前读取Schema（处理后Schema无变化时适用） #

直接用PyArrow读取输入Parquet文件的Schema，传递给WriteToParquet，简单高效。

修改后代码：

import pyarrow.parquet as pq

# 提前读取输入Parquet的Schema
schema = pq.read_schema(known_args.input)

with beam.Pipeline(options=pipeline_options) as p:
    dataframes = p \
      | 'Read' >> beam.io.ReadFromParquetBatched(known_args.input) \
      | 'Convert to pandas' >> beam.Map(lambda table: table.to_pandas()) \
      | 'Process df' >> beam.ParDo(ProcessDataFrame()) \
      | 'Convert to Arrow Table' >> beam.Map(lambda df: pyarrow.Table.from_pandas(df)) \
      | 'Write to parquet' >> beam.io.WriteToParquet(known_args.output, schema=schema)

方案2：从处理后的DataFrame动态提取Schema（处理后Schema可能变化时适用） #

通过Beam的分布式操作提取样本数据的Schema，再广播给所有写入任务，适配Schema变更场景。

修改后代码：

import pyarrow as pa

def extract_schema(df):
    # 将Pandas DataFrame转为Arrow Table并提取Schema
    return pa.Table.from_pandas(df).schema

with beam.Pipeline(options=pipeline_options) as p:
    processed_dfs = p \
      | 'Read' >> beam.io.ReadFromParquetBatched(known_args.input) \
      | 'Convert to pandas' >> beam.Map(lambda table: table.to_pandas()) \
      | 'Process df' >> beam.ParDo(ProcessDataFrame())
    
    # 从处理后的样本数据中提取Schema（取第一个有效DataFrame的Schema）
    schema = processed_dfs | 'Extract schema' >> beam.CombineGlobally(
        lambda dfs: extract_schema(next(iter(dfs))) if dfs else None
    ).without_defaults()
    
    # 将Schema广播到所有处理后的DataFrame，转换为Arrow Table后写入
    processed_dfs | 'Pair with schema' >> beam.Map(lambda df, schema: (df, schema), beam.pvalue.AsSingleton(schema)) \
                  | 'Convert to Arrow Table' >> beam.Map(lambda x: pa.Table.from_pandas(x[0], schema=x[1])) \
                  | 'Write to parquet' >> beam.io.WriteToParquet(known_args.output)

注意事项 #

• 若处理后各批次数据Schema不一致，需先确保Schema统一，否则写入会失败。
• 方案2中CombineGlobally取第一个样本的Schema，需保证样本能代表所有数据的Schema结构。

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