如何在GCP（Apache Beam）中同步执行Pipeline的指定作业操作

在Apache Beam（Dataflow Runner）环境中，要实现指定步骤等待所有并行前置作业完成后再执行，核心是捕获前置作业的完成信号，并以此触发后续操作。以下是具体实现方案：

关键思路 #

1. 保留前置写入作业的输出引用（即使是WriteToBigQuery这类无实际数据输出的Transform，其返回的空PCollection可作为作业完成的标记）
2. 合并多个前置作业的完成信号，生成单一触发条件
3. 基于触发条件执行最后的timestamp写入操作

修改后的完整代码 #

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

# 假设以下变量已提前定义：options、jdbc_url、username、password、wisgen_table、recruiter_table、last_sync_tab

p = beam.Pipeline(options=options) 

# 执行前置JDBC读取与BigQuery写入作业，并捕获完成信号
wisgen_data = p | "wisgen job" >> beam.io.ReadFromJdbc (
    jdbc_url=jdbc_url,
    username=username,
    password=password,
    driver_class_name='org.postgresql.Driver',
    query="""SELECT users.id AS user_id, CONCAT(users.first_name,' ', users.last_name) AS full_name""",
    table_name="users"
)
wisgen_write_complete = (wisgen_data 
                         | "Convert TableRow to dict(wisgen data)" >> beam.Map(lambda row: row._asdict())
                         | "Write to BigQuery in wisgen data table" >> beam.io.WriteToBigQuery(
                             wisgen_table,
                             write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND,
                             create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
                             schema='user_id:INTEGER,full_name:STRING'
                         ))

recruiter_data = p | "recruiter job" >> beam.io.ReadFromJdbc(
    jdbc_url=jdbc_url,
    username=username,
    password=password,
    driver_class_name='org.postgresql.Driver',
    query="""SELECT users.id AS user_id, ''""",
    table_name="users"
)
recruiter_write_complete = (recruiter_data 
                            | "Convert TableRow to dict(recruiter data)" >> beam.Map(lambda row: row._asdict())
                            | "Write to BigQuery in recruiter data table" >> beam.io.WriteToBigQuery(
                                recruiter_table,
                                write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND,
                                create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
                                schema='user_id:INTEGER,full_name:STRING'
                            ))

# 合并前置作业完成信号，生成触发条件
trigger_signal = (
    [wisgen_write_complete, recruiter_write_complete]
    | "Flatten completion signals" >> beam.Flatten()
    | "Generate trigger" >> beam.CombineGlobally(lambda _: None).without_defaults()
)

# 等待所有前置作业完成后，执行last_sync写入
_ = (
    trigger_signal
    | "Wait for all jobs to complete" >> beam.pvalue.WaitForSingleton()
    | "Create current timestamp" >> beam.Map(lambda _: {'time': datetime.datetime.utcnow()})
    | "Write to last_sync" >> beam.io.WriteToBigQuery(
        last_sync_tab,
        schema='time:TIMESTAMP',
        write_disposition=beam.io.BigQueryDisposition.WRITE_TRUNCATE,
        create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
    )
)

p.run().wait_until_finish()

代码说明 #

• 捕获完成信号：WriteToBigQuery返回空PCollection，将其作为前置作业完成的标记
• 生成触发信号：通过Flatten合并两个空集合，再用CombineGlobally生成单一信号元素，只有当两个前置作业全部完成时，该信号才会生成
• 等待触发执行：WaitForSingleton会阻塞后续步骤，直到触发信号生成，确保last_sync的写入操作在所有前置作业完成后执行

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