如何在Apache Beam Python中将PCollection按N个元素分组

嘿，我刚好在Python Beam里折腾过类似的批量打包需求，给你分享两种实用的方案，适配不同的场景：

如果你的数据源是像文本文件这样的有界数据集，用CombineGlobally配合自定义的CombineFn是最简洁的方式，代码如下：

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

class BatchCombineFn(CombineFn):
    def __init__(self, batch_size):
        self.batch_size = batch_size

    def create_accumulator(self):
        # 初始化累加器，用来暂存元素
        return []

    def add_input(self, accumulator, element):
        # 把新元素加入累加器，凑够batch_size就返回当前批次并重置
        accumulator.append(element)
        if len(accumulator) == self.batch_size:
            batch = accumulator.copy()
            accumulator.clear()
            return batch
        return accumulator

    def merge_accumulators(self, accumulators):
        # 合并各个分片的累加器，处理跨分片的批量
        merged = []
        for acc in accumulators:
            if isinstance(acc, list):
                merged.extend(acc)
                # 从合并后的列表里拆分出完整批次
                while len(merged) >= self.batch_size:
                    batch = merged[:self.batch_size]
                    yield batch
                    merged = merged[self.batch_size:]
        # 最后返回剩余的不足一个批次的元素（如果需要保留的话）
        if merged:
            yield merged

    def extract_output(self, accumulator):
        # 返回最终剩余的元素（不需要的话可以返回None，后续过滤掉）
        return accumulator if accumulator else None

# 集成到你的现有管道里
p = beam.Pipeline(options=pipeline_options)
lines = p | 'File reading' >> beam.io.ReadFromText(known_args.input)

# 替换N为你需要的批量大小
batched_lines = lines | 'Batch into N elements' >> beam.CombineGlobally(BatchCombineFn(N)).with_output_types(list)

如果后续你的数据源变成了Kafka这样的无界数据流，上面的CombineGlobally就不太好用了（因为无界数据需要配合窗口），这时候可以用带状态的DoFn来实现：

class BatchDoFn(beam.DoFn):
    def __init__(self, batch_size):
        self.batch_size = batch_size
        # 定义状态存储，用来暂存当前未凑够批次的元素
        self.batch_state = beam.DoFn.StateSpec('batch_state', beam.coders.ListCoder(beam.coders.StrUtf8Coder()))

    def process(self, element, state=beam.DoFn.StateParam(batch_state)):
        current_batch = state.read() or []
        current_batch.append(element)
        
        if len(current_batch) == self.batch_size:
            # 凑够批次，清空状态并输出
            state.write([])
            yield current_batch
        else:
            # 没凑够，更新状态继续累计
            state.write(current_batch)

# 集成方式
batched_lines = lines | 'Batch with Stateful DoFn' >> beam.ParDo(BatchDoFn(N))

• 如果不需要最后剩余的不足N的元素，可以在BatchCombineFn的extract_output里返回None，然后加一步beam.Filter(lambda x: x is not None)过滤掉。
• 无界数据流场景下，建议配合FixedWindows或者SlidingWindows使用，避免状态无限累计。

内容的提问来源于stack exchange，提问作者Luis Felipe Muñoz