如何在Python版Apache Beam中将单条Pub/Sub消息写入三个BigQuery表

原代码的关键问题 #

• DoFn方法命名错误：Beam的DoFn只会执行名为process的方法，自定义的process1/2/3不会被框架调用。
• 字典字段访问错误：json.loads返回字典对象，必须用parsed["key"]而非parsed.key的方式访问字段。
• 不必要的元组类型：赋值语句末尾的逗号会将字段值转为元组（如parsed["et"] = parsed["et"],），不符合BigQuery的字段类型要求。
• 字段映射不匹配：部分处理逻辑的输出字段名和目标表Schema不对应（比如第二个流程写parsed["name"]但目标表Schema是Place:string）。

修正后的完整代码 #

import json
import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions

# 配置参数，替换为你的实际信息
subscription = "pub/sub/writeToBigquery"
project_id = "your-gcp-project-id"
dataset_id = "your-bigquery-dataset-id"
dataset_table1 = f"{project_id}.{dataset_id}.A"
dataset_table2 = f"{project_id}.{dataset_id}.B"
dataset_table3 = f"{project_id}.{dataset_id}.c"

# 修正后的Schema，确保和输出字段完全匹配
schema1 = "et:timestamp,name:string,Day:integer"
schema2 = "et:timestamp,Place:string,Month:integer"
schema3 = "et:timestamp,Location:string,Year:integer"

# 拆分三个独立的DoFn，对应不同的表处理逻辑
class ParseForTableA(beam.DoFn):
    def process(self, element, timestamp=beam.DoFn.TimestampParam):
        parsed = json.loads(element.decode("utf-8"))
        yield {
            "et": parsed["et"],
            "name": parsed["name"],
            "Day": parsed["day"]
        }

class ParseForTableB(beam.DoFn):
    def process(self, element, timestamp=beam.DoFn.TimestampParam):
        parsed = json.loads(element.decode("utf-8"))
        yield {
            "et": parsed["et"],
            "Place": parsed["place"],
            "Month": parsed["month"]
        }

class ParseForTableC(beam.DoFn):
    def process(self, element, timestamp=beam.DoFn.TimestampParam):
        parsed = json.loads(element.decode("utf-8"))
        yield {
            "et": parsed["et"],
            "Location": parsed["location"],
            "Year": parsed["year"]
        }

def run():
    pipeline_options = PipelineOptions()
    with beam.Pipeline(options=pipeline_options) as p:
        # 从Pub/Sub订阅读取数据
        messages = p | "Read from Pub/Sub" >> beam.io.ReadFromPubSub(subscription=subscription)

        # 分支1：处理并写入表A
        (messages
         | "Parse for Table A" >> beam.ParDo(ParseForTableA())
         | "Write to Table A" >> beam.io.WriteToBigQuery(
             dataset_table1,
             schema=schema1,
             write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND,
             create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED
         ))

        # 分支2：处理并写入表B
        (messages
         | "Parse for Table B" >> beam.ParDo(ParseForTableB())
         | "Write to Table B" >> beam.io.WriteToBigQuery(
             dataset_table2,
             schema=schema2,
             write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND,
             create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED
         ))

        # 分支3：处理并写入表C
        (messages
         | "Parse for Table C" >> beam.ParDo(ParseForTableC())
         | "Write to Table C" >> beam.io.WriteToBigQuery(
             dataset_table3,
             schema=schema3,
             write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND,
             create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED
         ))

if __name__ == "__main__":
    run()

关键说明 #

• 参数替换：将project_id和dataset_id替换为你的GCP项目ID和BigQuery数据集ID。
• 写入模式：WRITE_APPEND表示追加数据到现有表，CREATE_IF_NEEDED会自动创建不存在的表（需确保运行账号有BigQuery创建权限）。
• 字段一致性：处理后输出的字典键必须和BigQuery Schema的字段名完全一致，数据类型也要匹配（比如et字段需为可解析的时间戳格式）。
• 权限配置：运行Pipeline的服务账号需要具备Pub/Sub订阅读取权限和BigQuery写入权限。

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