Apache Beam Python：BatchElements逆操作及批量集合拆分方法

BatchElements的逆操作是将PCollection[List[T]]还原为单个元素的PCollection[T]，常用两种方式实现：

1. 使用beam.FlatMap（最简方案） #

FlatMap会自动遍历每个列表元素，将其展开为独立的PCollection元素，无需额外定义DoFn，直接用lambda或简单函数即可：

# 直接用lambda展开列表
output_regrouped = output | '重新分组' >> beam.FlatMap(lambda batch: batch)

如果需要处理空列表、过滤无效元素等场景，也可以定义一个展开函数：

def expand_batch(batch):
    for item in batch:
        # 可添加过滤、转换等逻辑，比如跳过None值
        if item is not None:
            yield item

output_regrouped = output | '重新分组' >> beam.FlatMap(expand_batch)

2. 使用自定义ParDo（灵活扩展方案） #

如果展开过程中需要复杂的业务逻辑（比如日志记录、异常处理），可以自定义DoFn逐个输出元素：

class ExpandBatchDoFn(beam.DoFn):
    def process(self, batch):
        try:
            for item in batch:
                # 这里可加入自定义处理逻辑
                yield item
        except Exception as e:
            # 异常处理逻辑，比如记录错误日志
            logging.error(f"处理批量元素失败: {str(e)}")

output_regrouped = output | '重新分组' >> beam.ParDo(ExpandBatchDoFn())

完整修改后的流水线代码 #

import logging
import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions

# 假设CastFields、my_costly_function已提前定义
with beam.Pipeline(options=options) as pipeline:
    # 准备themis清理输出
    data = pipeline | beam.io.ReadFromParquet(file_pattern=input_location)
    output = data | '转换字段类型' >> beam.ParDo(CastFields(country, slug_name))
    group_batch = output | "分组为批量" >> beam.BatchElements(min_batch_size=500, max_batch_size=10000)
    output = group_batch | '执行耗时函数，批量处理更高效' >> beam.ParDo(my_costly_function(country, shared_handle)) # 生成List[T]类型的PCollection
    # 执行逆操作还原为单个元素的PCollection
    output_regrouped = output | '重新分组' >> beam.FlatMap(lambda x: x)
    output_regrouped | '写入Parquet' >> beam.io.WriteToParquet(output_location)

以上两种方式均可将批量处理后的List[T]类型PCollection转换为WriteToParquet支持的单个元素类型PCollection。

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