如何在Python版Beam中实现类似ARRAY_AGG(STRUCT(table))的操作？

核心思路 #

通过分布式分组聚合将明细数据按ID打包为数组（对应ARRAY_AGG），再与主表关联，最终构造符合BigQuery嵌套结构的输出行。你之前用原生groupby失败是因为它是单机操作，无法适配Beam的分布式管道，必须使用Beam提供的分布式分组转换。

完整实现代码 #

import apache_beam as beam
from apache_beam.io.gcp.bigquery import ReadFromBigQuery, WriteToBigQuery

# 处理主表数据：转换为键值对（ID, 主数据字典）
def format_main_row(row):
    return (row['ID'], {
        'id': row['ID'],
        'total': row['total']
    })

# 关联主表与聚合后的明细，构造最终输出结构
def combine_data(grouped):
    table_id, (main_rows, detail_groups) = grouped
    # 主表每个ID仅一行，取第一个元素
    main_data = main_rows[0]
    # 若当前ID无明细，赋值空列表
    main_data['line_items'] = detail_groups[0] if detail_groups else []
    return main_data

with beam.Pipeline() as p:
    # 1. 读取并格式化主表数据
    main_data = (
        p
        | "读取主表" >> ReadFromBigQuery(
            query="SELECT ID, total FROM `your-project.your-dataset.main_table`",
            use_standard_sql=True
        )
        | "格式化主表" >> beam.Map(format_main_row)
    )

    # 2. 读取明细数据，按ID分组聚合为结构化数组
    aggregated_details = (
        p
        | "读取明细表" >> ReadFromBigQuery(
            query="SELECT table1_id, item_name, item_price FROM `your-project.your-dataset.detail_table`",
            use_standard_sql=True
        )
        # 将每条明细转换为STRUCT对应的字典
        | "转换明细为结构" >> beam.Map(lambda row: (row['table1_id'], {
            'item_name': row['item_name'],
            'item_price': row['item_price']
        }))
        # 分布式分组：按table1_id聚合所有明细
        | "按ID分组明细" >> beam.GroupByKey()
        # 将分组后的明细转换为列表（对应ARRAY_AGG）
        | "聚合明细为数组" >> beam.Map(lambda x: (x[0], list(x[1])))
    )

    # 3. 关联主表与聚合后的明细
    final_output = (
        {'main': main_data, 'details': aggregated_details}
        | "关联主表与明细" >> beam.CoGroupByKey()
        | "构造最终输出行" >> beam.Map(combine_data)
    )

    # 4. 写入BigQuery，指定嵌套Schema
    final_output | "写入BigQuery" >> WriteToBigQuery(
        table="your-project.your-dataset.output_table",
        schema="""
            id:STRING,
            total:INTEGER,
            line_items:RECORD REPEATED,
                item_name:STRING,
                item_price:INTEGER
        """,
        write_disposition=beam.io.BigQueryDisposition.WRITE_TRUNCATE,
        create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED
    )

简化写法（使用beam.GroupBy） #

Beam 2.20+支持更简洁的GroupBy转换，无需手动处理键值对：

aggregated_details = (
    p
    | "读取明细表" >> ReadFromBigQuery(...)
    | "分组聚合明细" >> beam.GroupBy('table1_id').aggregate(
        lambda rows: [{'item_name': r['item_name'], 'item_price': r['item_price']} for r in rows],
        'line_items'
    )
    | "转换为键值对" >> beam.Map(lambda row: (row['table1_id'], row['line_items']))
)

关键注意事项 #

• Schema匹配：BigQuery输出表的line_items必须定义为RECORD REPEATED类型，与Python中的列表嵌套字典结构对应。
• 空明细处理：代码中已处理无明细的ID，确保输出line_items为空列表而非None。
• 分布式分组：必须使用Beam的GroupByKey或GroupBy，不能用Python原生groupby，后者无法在分布式环境中运行。

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