Apache Beam Python：传递变量或等待前置步骤适配SqlTransform

问题描述 #

我有如下Pipeline代码：

from apache_beam import coders
from apache_beam.coders import RowCoder
from typing import NamedTuple
import apache_beam as beam
from fastavro.schema import load_schema
import typing
from apache_beam.transforms.sql import SqlTransform
from apache_beam.typehints import schemas
import json


def infer_schema(data):
    print(data)
    fields = [(key, str if type(value) == dict else type(value)) for key, value in data.items()]
    print(fields)
    global inferedSchema
    if len(inferedSchema._fields) == 0:
        inferedSchema = typing.NamedTuple('inferedSchema',
                                   fields)
        coders.registry.register_coder(inferedSchema, coders.RowCoder)
    return data

print("HI")
with beam.Pipeline() as pipeline:
    record =( pipeline 
                     | "Read JSON" >> beam.io.ReadFromText("file2.json")
                      |"Infer schema" >> beam.Map(lambda x: infer_schema(json.loads(x)))
                      | "Transform to NamedTuple" >> beam.Map(lambda x: inferedSchema(**x)).with_output_types(inferedSchema) 
              # | "Transform to NamedTuple" >> beam.Map(lambda x: givenSchema(**x)).with_output_types(givenSchema) 
             #          | SqlTransform(
             # """
             #   SELECT
             #   *
             #   FROM PCOLLECTION""")
              | "Print" >> beam.Map(print)
            )
print(inferedSchema._fields)
print("Bye")

不使用SqlTransform时，Pipeline可以正常运行：接收字典（因特定需求必须用字典，无法用ReadFromJson）、推断Schema并转换为NamedTuple。但添加SqlTransform后，会跳过Schema推断返回空结果，除非在Pipeline运行前手动定义Schema，但这样无法适配数据变化。

请问是否存在等待函数或传递结果的方式，让我可以从第一条消息推断Schema并供后续步骤使用？

解决方案 #

核心问题在于SqlTransform需要在Pipeline构建阶段就明确知道输入Schema，而你当前的动态Schema推断是在Pipeline运行阶段（Map转换中）才完成的，这就导致SqlTransform初始化时无法获取到Schema信息，进而无法正确处理数据。

可以通过以下两种方式解决：

1. 提前读取首条数据推断Schema（推荐） #

在Pipeline构建前，先读取JSON文件的首条记录，提前完成Schema推断并定义NamedTuple，这样SqlTransform在初始化时就能获取到明确的Schema：

from apache_beam import coders
from apache_beam.coders import RowCoder
from typing import NamedTuple
import apache_beam as beam
import typing
from apache_beam.transforms.sql import SqlTransform
import json

# 提前读取首条数据推断Schema
def infer_schema_from_first_record(file_path):
    with open(file_path, 'r') as f:
        first_line = f.readline()
        data = json.loads(first_line)
        fields = [(key, str if isinstance(value, dict) else type(value)) for key, value in data.items()]
        inferred_schema = typing.NamedTuple('InferredSchema', fields)
        coders.registry.register_coder(inferred_schema, coders.RowCoder)
        return inferred_schema

# 提前完成Schema推断
inferred_schema = infer_schema_from_first_record("file2.json")

print("HI")
with beam.Pipeline() as pipeline:
    record = (pipeline 
              | "Read JSON" >> beam.io.ReadFromText("file2.json")
              | "Parse JSON" >> beam.Map(json.loads)
              | "Transform to NamedTuple" >> beam.Map(lambda x: inferred_schema(**x)).with_output_types(inferred_schema) 
              | "SQL Transform" >> SqlTransform("SELECT * FROM PCOLLECTION")
              | "Print" >> beam.Map(print)
            )
print(inferred_schema._fields)
print("Bye")

这种方式的优势是Schema在Pipeline构建阶段就已确定，完全适配SqlTransform的要求，同时保留了动态适配数据结构的能力。

2. 使用Beam的Row类型替代NamedTuple #

如果无法提前读取数据，可以使用apache_beam.Row类型结合schemas.add_metadata来动态传递Schema信息：

from apache_beam import coders
import apache_beam as beam
from apache_beam.transforms.sql import SqlTransform
from apache_beam.typehints import schemas
import json
from typing import Dict

def infer_and_convert_to_row(data: Dict):
    # 动态生成Schema字段
    fields = [(key, str if isinstance(value, dict) else type(value)) for key, value in data.items()]
    # 创建Row类型并添加元数据
    row_type = schemas.named_tuple_from_schema(schemas.Schema.from_fields(fields))
    return row_type(**data)

print("HI")
with beam.Pipeline() as pipeline:
    record = (pipeline 
              | "Read JSON" >> beam.io.ReadFromText("file2.json")
              | "Parse and Convert to Row" >> beam.Map(lambda x: infer_and_convert_to_row(json.loads(x)))
              | "SQL Transform" >> SqlTransform("SELECT * FROM PCOLLECTION")
              | "Print" >> beam.Map(print)
            )
print("Bye")

不过这种方式需要确保所有数据的Schema一致，否则在SqlTransform阶段会出现类型不匹配的错误。

关键注意事项 #

• SqlTransform依赖静态类型信息，必须在Pipeline构建阶段明确输入的Schema结构；
• 动态推断Schema时，要确保所有数据的字段和类型一致，否则后续转换会失败；
• 提前读取首条数据的方式兼容性更好，适合大多数动态Schema场景。

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