Apache Beam中基于窗口分组的Pub/Sub消息写入BigQuery失败问题求助

咱们一步步来看，你的流水线没法写入BigQuery的核心原因是窗口分组阶段意外修改了数据结构，但后续转换步骤没跟上这个变化，直接导致PrepareRow没法正确解析数据，最终失败。

具体来说，在AddTimestamps这个DoFn里，你把原始的字典对象转换成了(cluster_time_str, payload_str)这样的元组，但后面的PrepareRow还在傻乎乎地用element['payload']去访问字典字段——这肯定会抛出KeyError啊！元素处理失败了，自然就写不进BigQuery了。

另外，你原来的pipeline.run().wait_until_finish()放在了with beam.Pipeline()的块外面，这也是个小问题：with块结束后流水线资源已经被释放了，再调用run()容易出问题。

一步步修复 #

1. 修正AddTimestamps的输出结构 #

咱们得保留原始的字典结构，只给元素打上时间戳就行，这样后续步骤能正常访问字段：

class AddTimestamps(beam.DoFn, ABC):
    def process(self, element, **kwargs):
        unix_timestamp = int(element['meta']['cluster_time'].split("-")[0])
        # 保留原始字典，只添加时间戳标记
        yield TimestampedValue(element, unix_timestamp)

2. 调整窗口分组逻辑（保持你的自定义Transform） #

你的GroupWindowsIntoBatches逻辑本身没问题，只要数据结构正确就可以正常工作，修改后如下（其实就是不用改结构，只是保留原来的字典）：

class GroupWindowsIntoBatches(beam.PTransform):
    """
    按发布时间分组Pub/Sub消息，输出一批批原始消息字典的列表
    """
    def __init__(self, window_size):
        super().__init__()
        self.window_size = window_size

    def expand(self, pcoll):
        return (
            pcoll
            | "按固定间隔划分窗口" >> beam.WindowInto(FixedWindows(self.window_size))
            | "为消息添加时间戳" >> beam.ParDo(AddTimestamps())
            | "添加虚拟Key" >> beam.Map(lambda elem: (None, elem))
            | "按Key分组" >> beam.GroupByKey()
            | "丢弃虚拟Key" >> beam.MapTuple(lambda _, val: val)
        )

3. 修正流水线的执行逻辑 #

把pipeline.run().wait_until_finish()移到with beam.Pipeline(...) as pipeline:的块内部，这样就能保证在流水线资源释放前完成执行：

4. 验证PrepareRow的正确性 #

现在元素还是原始的字典结构，PrepareRow里的element['payload']就能正常访问了，如果你想更稳妥，可以加个日志看看元素内容：

class PrepareRow(beam.DoFn, ABC):
    def process(self, element, **kwargs):
        logging.info(f"当前处理的元素: {element}")
        migration_datetime = datetime.datetime.now().strftime("%Y-%m-%dT%H:%M:%S")
        migration_id = "windowed migration_id"
        return [
            {"raw_contents": element['payload'], "migration_datetime": migration_datetime, "migration_id": migration_id}
        ]

修改后的完整代码 #

import argparse
import json
import logging
import os
from abc import ABC
import datetime
import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.transforms.window import TimestampedValue
from apache_beam.transforms.window import FixedWindows

os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = "./credentials/dataflow_service_account.json"
SUBSCRIPTION = "projects/project-data-engineering-ci/subscriptions/beam-local-sub"

class GroupWindowsIntoBatches(beam.PTransform):
    """
    A composite transform that groups Pub/Sub messages based on publish time and outputs a list of dictionaries, where each contains one message and its publish timestamp.
    """
    def __init__(self, window_size):
        super().__init__()
        self.window_size = window_size

    def expand(self, pcoll):
        return (
            pcoll
            # Assigns window info to each Pub/Sub message based on its cluster time.
            | "Window into Fixed Intervals" >> beam.WindowInto(FixedWindows(self.window_size))
            | "Add timestamps to messages" >> beam.ParDo(AddTimestamps())
            | "Add Dummy Key" >> beam.Map(lambda elem: (None, elem))
            | "Groupby" >> beam.GroupByKey()
            | "Abandon Dummy Key" >> beam.MapTuple(lambda _, val: val)
        )

class AddTimestamps(beam.DoFn, ABC):
    def process(self, element, **kwargs):
        unix_timestamp = int(element['meta']['cluster_time'].split("-")[0])
        # 保留原始字典结构，只添加时间戳
        yield TimestampedValue(element, unix_timestamp)

class PrintValue(beam.DoFn, ABC):
    def process(self, element, **kwargs):
        print(element)
        return [element]

class PrepareRow(beam.DoFn, ABC):
    def process(self, element, **kwargs):
        logging.info(f"Processing element: {element}")
        migration_datetime = datetime.datetime.now().strftime("%Y-%m-%dT%H:%M:%S")
        migration_id = "windowed migration_id"
        return [
            {"raw_contents": element['payload'], "migration_datetime": migration_datetime, "migration_id": migration_id}
        ]

TABLE_SCHEMA = {
    "fields": [
        {"name": "raw_contents", 'type': 'STRING', 'mode': 'NULLABLE'},
        {"name": "migration_datetime", 'type': 'DATETIME', 'mode': 'NULLABLE'},
        {"name": "migration_id", 'type': 'STRING', 'mode': 'NULLABLE'}
    ]
}

def run(input_subscription, output_table_bq, output_table_bq_schema, window_size, pipeline_args=None):
    pipeline_options = PipelineOptions(
        pipeline_args,
        streaming=True,
        save_main_session=True,
        direct_running_mode='in_memory',
        direct_num_workers=2
    )

    with beam.Pipeline(options=pipeline_options) as pipeline:
        events = (
            pipeline
            | beam.io.ReadFromPubSub(subscription=input_subscription)
            | "Convert bytes to a dictionary" >> beam.Map(lambda e: json.loads(e.decode('utf-8')))
        )

        _ = (
            events
            | "Group Windows Into Batches" >> GroupWindowsIntoBatches(window_size)
            | "FlatMap" >> beam.FlatMap(lambda elements: elements)
            | "Prepare rows before inserting" >> beam.ParDo(PrepareRow())
            | "Write to BQ" >> beam.io.WriteToBigQuery(
                table=output_table_bq,
                schema=output_table_bq_schema,
                write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND,
                create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED
            )
        )

        # 将run().wait_until_finish()移到with块内部
        pipeline.run().wait_until_finish()

if __name__ == "__main__":
    logging.getLogger().setLevel(logging.INFO)
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input_subscription",
        dest='input_subscription',
        help="The Cloud Pub/Sub subscription to read from.\n"
        '"projects/<PROJECT_NAME>/subscriptions/<SUBSCRIPTION_NAME>".',
    )
    parser.add_argument(
        "--window_size",
        dest='window_size',
        type=int,
        default=10,
        help="Output file's window size in number of seconds.",
    )
    parser.add_argument(
        "--output_table_bq",
        dest='output_table_bq',
        required=True,
        help="BQ Table for output. Format: <project_id:dataset.table>",
    )

    known_args, pipeline_args = parser.parse_known_args()

    run(
        input_subscription=known_args.input_subscription or SUBSCRIPTION,
        output_table_bq=known_args.output_table_bq,
        window_size=known_args.window_size,
        output_table_bq_schema=TABLE_SCHEMA,
        pipeline_args=pipeline_args,
    )

额外优化小建议 #

1. 用Beam内置的GroupIntoBatches替代自定义Transform：如果你的需求只是按窗口批量处理数据，直接用Beam提供的GroupIntoBatches会更简洁，比如：

   events | "Group into batches" >> beam.GroupIntoBatches(batch_size=10)
   

2. 用窗口时间生成批次标识：如果想让migration_datetime和migration_id基于窗口的时间（而不是当前系统时间），可以在PrepareRow中获取窗口参数：

   class PrepareRow(beam.DoFn, ABC):
       def process(self, element, window=beam.DoFn.WindowParam, **kwargs):
           window_end = window.end.to_utc_datetime().strftime("%Y-%m-%dT%H:%M:%S")
           migration_datetime = window_end
           migration_id = f"window_{window.start.to_utc_datetime()}_{window_end}"
           return [
               {"raw_contents": element['payload'], "migration_datetime": migration_datetime, "migration_id": migration_id}
           ]
   

3. 添加错误处理：可以用beam.Map或者ParDo的with_outputs来捕获处理失败的元素，方便排查问题：

   prepared_rows, failed_rows = (
       elements
       | "Prepare rows" >> beam.ParDo(PrepareRow()).with_outputs('failed', main='success')
   )
   failed_rows | "Log failed elements" >> beam.Map(lambda e: logging.error(f"Failed to process: {e}"))
   

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