如何在Beam中实现本地partition内按Key分组（替代全局GroupByKey）

当然有办法搞定你要的「只在本地Partition内按Key分组」的需求！Beam虽然没有直接提供一个叫LocalGroupByKey的原生变换，但我们可以用两种简单的方式达成目标：

方法1：用Combine.perKey配合自定义CombineFn实现本地分组 #

如果你不需要跨分区的全局聚合，只是想在每个本地Partition内收集同一Key下的所有元素（输出per-key-per-window的Iterable），可以用Combine.perKey搭配一个只做元素收集的CombineFn，再加上禁用跨分区洗牌的配置。

举个Python的例子：

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

class LocalGroupCombineFn(CombineFn):
    def create_accumulator(self):
        return []
    
    def add_input(self, accumulator, element):
        accumulator.append(element)
        return accumulator
    
    def merge_accumulators(self, accumulators):
        merged = []
        for acc in accumulators:
            merged.extend(acc)
        return merged
    
    def extract_output(self, accumulator):
        return accumulator

# 在Pipeline中使用
with beam.Pipeline() as p:
    (p
     | beam.Create([("a", 1), ("b", 2), ("a", 3), ("c", 4), ("b", 5)])
     # 设置shuffle_mode为LOCAL，强制只在本地Partition内完成分组
     | beam.CombinePerKey(LocalGroupCombineFn(), shuffle_mode=beam.transforms.combiners.ShuffleMode.LOCAL)
     | beam.Map(print)
    )

核心就是ShuffleMode.LOCAL这个参数——它会告诉Beam不要把数据跨分区搬运洗牌，只在元素所在的本地Partition内完成Key的分组聚合，完全符合你的需求。

方法2：先标记Partition，再按「分区+Key」分组 #

另一种更直观的思路：先给每个元素打上它所在Partition的标记，然后按「Partition ID + Key」这个复合键做全局GroupByKey，最后再把Partition ID剥离掉，就能得到每个Key在本地Partition内的分组结果了。

Python示例：

import apache_beam as beam

with beam.Pipeline() as p:
    # 第一步：给每个元素带上所在Partition的标记
    tagged_elements = (p
     | beam.Create([("a", 1), ("b", 2), ("a", 3), ("c", 4), ("b", 5)])
     | beam.Map(lambda x: (beam.io.utils.get_current_partition(), x))
    )
    
    # 第二步：按(Partition ID, Key)做全局分组
    grouped = (tagged_elements
     | beam.Map(lambda x: ((x[0], x[1][0]), x[1][1]))
     | beam.GroupByKey()
    )
    
    # 第三步：去掉Partition ID，得到本地分组结果
    local_grouped = (grouped
     | beam.Map(lambda x: (x[0][1], list(x[1])))
     | beam.Map(print)
    )

这个方法逻辑简单，不用写自定义CombineFn，但要注意：获取当前Partition ID的API在不同Beam Runner下的兼容性可能有差异，使用前最好先测试下你的目标Runner是否支持。

额外提醒 #

• 这两种方式都不会触发跨Partition的数据洗牌，性能比全局GroupByKey好很多，但只能拿到单个Partition内的Key分组结果，完全符合你的要求。
• 如果你用的是Java SDK，思路是完全一样的：可以用Combine.perKey()配合ShuffleMode.LOCAL，或者通过标记分区的方式实现。

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