如何使用Apache Beam Python SDK合并PCollection中两个Key的列表值

要将PCollection里Key1和Key2对应的列表合并为Key3下的完整列表，可通过以下步骤实现：

核心思路 #

1. 提取所有目标Key对应的列表值；
2. 将所有列表合并为一个全局列表；
3. 为合并后的列表关联新的KeyKey3。

代码实现 #

方法一：自定义合并逻辑（适用于复杂场景）

通过自定义CombineFn实现列表合并，逻辑清晰且便于扩展：

import apache_beam as beam
from apache_beam.transforms.core import CombineFn

class MergeListsFn(CombineFn):
    def create_accumulator(self):
        return []
    
    def add_input(self, accumulator, input_list):
        accumulator.extend(input_list)
        return accumulator
    
    def merge_accumulators(self, accumulators):
        merged = []
        for acc in accumulators:
            merged.extend(acc)
        return merged
    
    def extract_output(self, accumulator):
        return accumulator

with beam.Pipeline() as pipeline:
    # 模拟输入PCollection
    input_data = pipeline | beam.Create([
        ("Key1", ["file1", "file2"]),
        ("Key2", ["file3", "file4"])
    ])
    
    result = (
        input_data
        | beam.Values()  # 提取所有列表值
        | beam.CombineGlobally(MergeListsFn())  # 全局合并列表
        | beam.Map(lambda merged_list: ("Key3", merged_list))  # 绑定新Key
    )
    
    # 打印结果验证
    result | beam.Map(print)

方法二：简洁写法（适用于简单场景）

用lambda表达式替代自定义CombineFn，代码更紧凑：

import apache_beam as beam

with beam.Pipeline() as pipeline:
    input_data = pipeline | beam.Create([
        ("Key1", ["file1", "file2"]),
        ("Key2", ["file3", "file4"])
    ])
    
    result = (
        input_data
        | beam.Values()
        | beam.CombineGlobally(lambda lists: sum(lists, []))  # 直接合并所有列表
        | beam.Map(lambda lst: ("Key3", lst))
    )
    
    result | beam.Map(print)

过滤特定Key（可选） #

如果只需合并指定Key（如仅Key1和Key2），可先添加过滤步骤：

result = (
    input_data
    | beam.Filter(lambda kv: kv[0] in ["Key1", "Key2"])  # 仅保留目标Key的元素
    | beam.Values()
    | beam.CombineGlobally(lambda lists: sum(lists, []))
    | beam.Map(lambda lst: ("Key3", lst))
)

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