如何在Apache Beam中实现10条记录的滚动窗口？

阿华AIGC实验室

2026-5-15

嘿，完全懂你刚接触Apache Beam时的困惑！毕竟Pandas里一行my_data_frame.rolling(10).sum().dropna()就能搞定的滑动记录窗口，到Beam里好像突然找不到直接对应的API——这其实是因为两者的设计思路差异：Pandas是单机内存批处理，所有数据都在内存里，能轻松实现基于位置的滑动；而Beam是分布式流/批通用框架，默认窗口是为流场景设计的时间驱动窗口，但咱们还是有办法实现你要的基于连续ID的滑动记录窗口的！

下面给你两种方案，分别适合流处理（实时数据）和批处理（离线数据）场景：

方案1：流/批通用（推荐）——用状态（State）实现滑动窗口

这是最符合Beam设计模式的方法，既能处理实时流数据，也能处理离线批数据，而且适合大数据量场景。核心思路是用Beam的State特性维护最近的9条记录，每来一条新记录就更新状态，然后计算最近10条的总和。

代码示例：

import apache_beam as beam
from apache_beam.transforms.userstate import BagStateSpec

class SlidingWindowSum(beam.DoFn):
    # 定义状态：保存最近的数值，用FloatCoder编码（如果你的值是其他类型可替换）
    STATE_SPEC = BagStateSpec('recent_values', beam.coders.FloatCoder())
    WINDOW_SIZE = 10  # 滑动窗口大小

    def process(self, element, recent_values=beam.DoFn.StateParam(STATE_SPEC)):
        # 假设你的每条记录是(id, value)的元组（如果结构不同可调整）
        current_id, current_value = element
        
        # 把当前值加入状态
        recent_values.add(current_value)
        
        # 取出状态里的所有值，确保不超过窗口大小
        values_list = list(recent_values.read())
        if len(values_list) > self.WINDOW_SIZE:
            # 因为输入是按ID排序的，所以最早加入的就是最旧的记录，直接移除
            values_list.pop(0)
            # 更新状态（先清空再重新添加）
            recent_values.clear()
            for val in values_list:
                recent_values.add(val)
        
        # 当状态里的元素数量刚好等于窗口大小时，输出总和
        if len(values_list) == self.WINDOW_SIZE:
            total_sum = sum(values_list)
            # 输出窗口范围和总和，格式可以自己调整
            yield (f"window_{current_id-9}_to_{current_id}", total_sum)

def run_pipeline():
    with beam.Pipeline() as p:
        # 模拟按ID排序的输入数据（替换成你的数据源）
        sorted_records = p | "生成排序记录" >> beam.Create([(i, i+1) for i in range(20)])
        
        # 给所有记录分配同一个键，让它们进入同一个DoFn实例处理（确保状态能连续维护）
        keyed_records = sorted_records | "添加全局键" >> beam.Map(lambda x: ('global_group', x))
        
        # 应用滑动窗口求和逻辑
        sliding_sums = keyed_records | "滑动窗口求和" >> beam.ParDo(SlidingWindowSum())
        
        # 输出结果
        sliding_sums | "打印结果" >> beam.Map(print)

if __name__ == "__main__":
    run_pipeline()

方案2：批处理专用——直接生成滑动子列表求和

如果你的数据是离线批处理场景，而且数据量不大（能在单机内存里放下），可以用更简单的方法：把所有数据收集成一个列表，然后生成滑动子列表再求和。

代码示例：

import apache_beam as beam

def generate_sliding_windows(elements, window_size=10):
    # 提取所有记录的数值（假设输入是(id, value)格式，可根据实际调整）
    values = [elem[1] for elem in elements]
    # 生成滑动子列表并计算总和
    for i in range(len(values) - window_size + 1):
        window_values = values[i:i+window_size]
        yield (f"window_{i}_to_{i+window_size-1}", sum(window_values))

def run_batch_pipeline():
    with beam.Pipeline() as p:
        sorted_records = p | "生成排序记录" >> beam.Create([(i, i+1) for i in range(20)])
        
        # 把所有记录合并成一个列表（仅适合批处理，大数据量慎用）
        sliding_results = sorted_records | "合并成列表" >> beam.CombineGlobally(generate_sliding_windows)
        
        # 展开结果并打印
        sliding_results | "展开输出" >> beam.FlatMap(lambda x: x) | beam.Map(print)

if __name__ == "__main__":
    run_batch_pipeline()