如何确保执行CoGroupByKey前已接收两个PubSub流的文件？

要解决这个问题，你需要从键提取、窗口策略、结果过滤三个层面调整现有代码，同时修正原代码中的语法错误，具体如下：

核心问题分析 #

你的现有代码存在几个关键问题：

• 没有将数据转换成(共同键, 内容)的键值对结构，CoGroupByKey无法识别关联依据
• 固定窗口可能因通知到达间隔超过窗口时长，导致同键数据分属不同窗口
• 未定义file_1_collected/file_2_collected变量，且多余的Flatten操作会破坏关联所需的结构
• 没有过滤仅单端存在数据的结果，无法保证关联时两端文件都已到达

具体实现步骤 #

1. 解析JSON并提取共同键 #

首先需要将读取到的JSON文本解析为对象，提取用于关联的共同键（比如假设键为id字段），将数据转换成(键, 内容)的结构，这是CoGroupByKey的必要输入格式。

2. 优化窗口策略 #

如果两个文件的通知到达时间不确定，推荐使用会话窗口（SessionWindows）：设置一个超时间隔（比如5分钟），同一个键的后续数据只要在超时内到达，就会合并到同一个窗口，避免因时间差导致的窗口拆分问题。如果确定通知间隔固定，也可以保留固定窗口，但需配合允许迟到时间。

3. 过滤单端数据 #

CoGroupByKey会保留所有键的关联结果，包括仅一端有数据的情况，因此需要添加过滤步骤，只保留两端都有数据的结果。

修正后的完整代码 #

import json
import apache_beam as beam
from apache_beam.transforms import window

def parse_json_and_extract_key(element, key_field='id'):
    # 解析JSON文本，提取共同键，返回(键, 解析后的JSON数据)
    json_data = json.loads(element)
    key = json_data.get(key_field)
    if not key:
        raise ValueError(f"数据中缺少关联键字段 '{key_field}'：{element}")
    return key, json_data

with beam.Pipeline(options=pipeline_options) as p:
    # 处理第一个PubSub数据流
    file_1 = (
        p
        | "file_1: 读取PubSub通知" >> beam.io.ReadFromPubSub(subscription="sub1").with_output_types(bytes)
        | "file_1: 字节转字符串" >> beam.Map(lambda x: x.decode('utf-8'))
        | "file_1: 读取文件内容" >> beam.io.ReadAllFromText()
        | "file_1: 解析JSON并提取键" >> beam.Map(parse_json_and_extract_key, key_field='id')
        | "file_1: 会话窗口分组" >> beam.WindowInto(window.Sessions(gap_duration=300))  # 5分钟会话超时
    )
    
    # 处理第二个PubSub数据流
    file_2 = (
        p
        | "file_2: 读取PubSub通知" >> beam.io.ReadFromPubSub(subscription="sub2").with_output_types(bytes)
        | "file_2: 字节转字符串" >> beam.Map(lambda x: x.decode('utf-8'))
        | "file_2: 读取文件内容" >> beam.io.ReadAllFromText()
        | "file_2: 解析JSON并提取键" >> beam.Map(parse_json_and_extract_key, key_field='id')
        | "file_2: 会话窗口分组" >> beam.WindowInto(window.Sessions(gap_duration=300))
    )
    
    # 执行CoGroupByKey关联操作
    joined_data = (
        {'file_1': file_1, 'file_2': file_2}
        | "按共同键关联" >> beam.CoGroupByKey()
    )
    
    # 过滤仅单端存在的数据，确保两端文件都已到达
    filtered_result = (
        joined_data
        | "过滤单端无效数据" >> beam.Filter(lambda x: len(x[1]['file_1']) > 0 and len(x[1]['file_2']) > 0)
        | "整理关联结果" >> beam.Map(lambda x: (x[0], {'file_1': x[1]['file_1'][0], 'file_2': x[1]['file_2'][0]}))
        # 注：如果同一个键对应多个文件，可根据实际需求调整结果整理逻辑
    )
    
    # 后续可将结果写入存储（如BigQuery、GCS等）
    # filtered_result | "写入结果存储" >> ...

额外优化建议 #

• 迟到数据处理：如果存在通知延迟的情况，可以给窗口添加允许迟到时间，比如beam.WindowInto(window.Sessions(300), allowed_lateness=window.Duration(3600))，允许1小时内的迟到数据纳入关联。
• 数据去重：如果PubSub可能重复推送通知，可以在读取文件后添加去重逻辑（比如基于文件路径或键+内容哈希），避免重复处理。
• 触发策略：可自定义窗口触发条件，比如trigger=AfterWatermark(early=AfterCount(1), late=AfterCount(1))，实现数据到达后的及时处理，同时兼顾迟到数据。

内容的提问来源于stack exchange，提问作者Grégoire Borel