如何利用PCollection列表长度与值对PCollection进行分区？

从BigQuery动态获取站点ID列表并转为Beam列表，计算其长度后，希望用该列表和长度对大型PCollection做分区，再将每个分区写入GCS对应CSV文件。但直接将sites_list和site_list_len传给beam.Partition时遇到阻碍——Beam仅允许通过ParDo的Side Inputs（基于pvalue模块）传递PCollections，需从Side Inputs、Singletons、自定义PTransforms/ParDo方向寻求解决方案。

原代码示例：

with beam.Pipeline(options=pipeline_options) as pipeline:
        sites_list = (
            pipeline
            | "Get site_id from D_Sites Table"
            >> beam.io.ReadFromBigQuery(
                query="SELECT DISTINCT(site_id) FROM `dataset.D_Sites`;",
                use_standard_sql=True,
                project="myproject",
            )
            | "Get Keys" >> beam.Map(lambda x: x["site_id"])
            | "Combine to list" >> beam.combiners.ToList()
        )

        site_list_len = sites_list | "Calculate Site List length" >> beam.Map(len)

        edge_data = (
            pipeline
            | "Query and Filter Edge data"
            >> beam.io.ReadFromBigQuery(
                query=special_query,
                use_standard_sql=True,
                project="myproject",
            )
            | "Partition"
            >> beam.Partition(
                lambda data, num_partitions: sites_list.index(data["site_id"]),
                site_list_len,
            )
            | "Print" >> beam.Map(print)
        )

sites_list示例内容：

["S10001","S10002","S10023","S10003","A10012"]

预期后续添加的写入代码：

for i, x in enumerate(sites_list):
    d[i] | f"Map to CSV strings {x}" >> beam.Map(
        lambda x: ",".join(x.values())
    ) | f"Write to CSV {x}" >> beam.io.WriteToText(
        f"gs://test-bucket/partition-test/{x}",
        file_name_suffix=".csv",
        header="col1,name,duration,...",
    )

核心思路是将sites_list转为Singleton Side Input传给自定义分区逻辑，替代原生beam.Partition的直接参数传递，再通过动态分支处理每个分区的写入任务。

完整修正代码 #

import apache_beam as beam
from apache_beam.pvalue import AsSingleton

def get_partition_index(element, site_index_map):
    # 通过预构建的字典快速获取site_id对应的分区索引
    return site_index_map[element["site_id"]]

def create_site_index_map(sites_list):
    # 将站点列表转为{site_id: 分区索引}的字典，提升查找效率
    return {site: idx for idx, site in enumerate(sites_list)}

with beam.Pipeline(options=pipeline_options) as pipeline:
    # 1. 获取站点列表并转为索引映射的Singleton Side Input
    sites_list_pcoll = (
        pipeline
        | "Fetch distinct site_ids"
        >> beam.io.ReadFromBigQuery(
            query="SELECT DISTINCT(site_id) FROM `dataset.D_Sites`;",
            use_standard_sql=True,
            project="myproject",
        )
        | "Extract site_id values" >> beam.Map(lambda x: x["site_id"])
        | "Combine to single list" >> beam.combiners.ToList()
    )
    # 生成站点-索引映射的Side Input
    site_index_map_pcoll = sites_list_pcoll | "Create site-index map" >> beam.Map(create_site_index_map)
    site_index_side_input = AsSingleton(site_index_map_pcoll)

    # 2. 提前获取站点列表的实际值，用于后续动态分支处理
    result = pipeline.run()
    result.wait_until_finish()
    sites_list = result.get(sites_list_pcoll)[0]
    num_partitions = len(sites_list)

    # 3. 处理边缘数据并按site_id分区
    edge_partitions = (
        pipeline
        | "Load filtered edge data"
        >> beam.io.ReadFromBigQuery(
            query=special_query,
            use_standard_sql=True,
            project="myproject",
        )
        | "Partition by site_id"
        >> beam.Partition(
            get_partition_index,
            num_partitions,
            site_index_map=site_index_side_input  # 传入Side Input
        )
    )

    # 4. 遍历每个分区，转换为CSV并写入GCS
    for idx, site_id in enumerate(sites_list):
        (
            edge_partitions[idx]
            | f"Convert to CSV - {site_id}" >> beam.Map(lambda x: ",".join(x.values()))
            | f"Write CSV to GCS - {site_id}" >> beam.io.WriteToText(
                f"gs://test-bucket/partition-test/{site_id}",
                file_name_suffix=".csv",
                header="col1,name,duration,...",
                shard_name_template="",  # 可选：关闭分片，生成单个文件
                append_trailing_newlines=True
            )
        )

关键细节说明 #

• Singleton Side Input：将站点索引映射作为单个元素传递给每个Worker，避免重复传输大列表，提升效率。
• 字典索引映射：将站点列表转为字典后，查找分区索引的时间复杂度从O(n)降到O(1)，适合站点数量较多的场景。
• 动态分支处理：通过result.get(sites_list_pcoll)获取实际的站点列表，再遍历生成对应分区的写入任务，适配动态获取的站点列表。

注意事项 #

• 确保special_query返回的数据包含site_id字段，否则会触发KeyError。
• 如果站点列表极大，需考虑Side Input的传输开销，可结合Beam的分布式缓存优化。

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