在Apache Beam Dataflow中比较单列数据集及CoGroupByKey报错解决

问题背景 #

需要对比两个仅含单列的数据集，输出三类结果：两者共有的数据、仅存在于数据集1的数据、仅存在于数据集2的数据。尝试用CoGroupByKey实现时出现报错，代码及错误信息如下：

错误代码 #

import apache_beam as beam

with beam.Pipeline() as p:
  dataset1 = ['data1a', 'data1b', 'data1c', 'data1d']
  dataset2 = ['data2a', 'data2b', 'data1b', 'data2d']
  dataset1_pcoll = p | 'Read Dataset 1' >> beam.Create(dataset1)
  dataset2_pcoll = p | 'Read Dataset 2' >> beam.Create(dataset2)
  
  combined_data = (
            {
                'file1': dataset1_pcoll,
                'file2': dataset2_pcoll,
            }
            | 'CoGroup Files' >> beam.CoGroupByKey()          
        )

报错信息 #

wrapper = lambda x: [fn(*x)]
TypeError: <lambda>() takes 2 positional arguments but 6 were given [while running 'CoGroup Files/CoGroupByKeyImpl/Tag[file2]']

报错原因 #

CoGroupByKey要求输入的PCollection必须是**键值对（Key-Value Tuple）**格式，它基于Key对多个PCollection的元素进行分组合并。上述代码直接传入了一维字符串列表，没有指定Key，Beam无法识别分组规则，因此抛出参数不匹配的错误。

解决方案 #

核心思路 #

1. 给两个数据集的每个元素添加统一的Key（比如固定值'compare_key'），将一维列表转换为键值对形式的PCollection。
2. 通过CoGroupByKey合并两个PCollection，得到对应Key下的两个数据集元素列表。
3. 利用集合的交集、差集运算，提取三类目标结果。

完整实现代码 #

import apache_beam as beam

def compare_datasets(element):
    _, grouped_data = element
    dataset1_items = set(grouped_data['file1'])
    dataset2_items = set(grouped_data['file2'])
    
    # 计算三类结果
    common_items = dataset1_items & dataset2_items
    only_dataset1 = dataset1_items - dataset2_items
    only_dataset2 = dataset2_items - dataset1_items
    
    return {
        'common_data': list(common_items),
        'only_in_dataset1': list(only_dataset1),
        'only_in_dataset2': list(only_dataset2)
    }

with beam.Pipeline() as p:
    dataset1 = ['data1a', 'data1b', 'data1c', 'data1d']
    dataset2 = ['data2a', 'data2b', 'data1b', 'data2d']
    
    # 转换为键值对PCollection
    dataset1_kv = p | 'Read Dataset1' >> beam.Create(dataset1) | 'Add Key to Dataset1' >> beam.Map(lambda x: ('compare_key', x))
    dataset2_kv = p | 'Read Dataset2' >> beam.Create(dataset2) | 'Add Key to Dataset2' >> beam.Map(lambda x: ('compare_key', x))
    
    # CoGroupByKey合并
    combined = (
        {
            'file1': dataset1_kv,
            'file2': dataset2_kv
        }
        | 'CoGroup Datasets' >> beam.CoGroupByKey()
    )
    
    # 处理合并结果，提取三类数据
    comparison_result = combined | 'Compare Datasets' >> beam.Map(compare_datasets)
    
    # 输出结果
    comparison_result | 'Print Result' >> beam.Map(print)

代码说明 #

• 添加Key：通过beam.Map(lambda x: ('compare_key', x))将每个元素转换为(Key, Value)格式，确保所有元素都被分到同一个Key下，方便合并。
• CoGroupByKey合并：此时传入的是键值对PCollection，CoGroupByKey会将相同Key对应的两个数据集元素列表整合到一起。
• 结果计算：在compare_datasets函数中，将两个列表转为集合，利用集合运算快速得到交集（共有数据）和差集（仅某数据集存在的数据）。
• 输出：最后将结果打印，也可根据需求写入存储系统（如GCS、BigQuery等）。

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