基于Apache Beam有状态处理补全时间序列PCollection缺失时间戳

需要为PCollection补全缺失的时间戳，思路如下：

• 使用状态创建所有可用时间步的列表
• 从列表推断出缺失的时间步
• 生成包含这些缺失时间步的PCollection
• 合并两个PCollection

但在第二步遇到问题，无法从状态中获取完整的时间步列表，尝试了如下Python代码但读取状态失败，示例中缺失的时间步为2022-03-01 00:02:00，请问如何获取完整的可用时间步列表？

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


class StatefulBufferingFn(beam.DoFn):
  BUFFER_STATE = BagStateSpec('buffer', StrUtf8Coder())

  def process(self, element,
              buffer_state=beam.DoFn.StateParam(BUFFER_STATE)):

    buffer_state.add(element[1]['timestamp'])
    yield element

data = [
    ('2022-03-01 00:00:00', {'timestamp': '2022-03-01 00:00:00', 'value': 1}),
    ('2022-03-01 00:01:00', {'timestamp': '2022-03-01 00:01:00', 'value': 2}),
    ('2022-03-01 00:03:00', {'timestamp': '2022-03-01 00:03:00', 'value': 3}),
    ('2022-03-01 00:04:00', {'timestamp': '2022-03-01 00:04:00', 'value': 4})
]

def read_final_state(state_pcoll):
  def print_state(timestamps):
    print('Final state:', list(timestamps))
  return state_pcoll | beam.CombineGlobally(lambda x: x).without_defaults() | beam.Map(print_state)


with beam.Pipeline(runner='DirectRunner') as pipeline:
    result, final_state = (
            pipeline
            | 'Create data' >> beam.Create(data)
            | 'apply_beam' >> beam.ParDo(StatefulBufferingFn()).with_outputs('final_state', main='output'))

    read_final_state(final_state)

原代码存在两个核心问题：

1. 未输出侧输出数据：代码定义了多输出，但StatefulBufferingFn的process方法仅生成主输出，从未向final_state侧输出写入内容，导致后续读取不到状态数据。
2. 状态分散存储：Apache Beam的用户状态默认绑定元素的key，原代码未做分组，每个元素的key是自身时间戳，状态被分散在多个key对应的容器中，无法获取全局完整的时间步列表。

要获取完整时间步列表，需先将所有数据聚合到同一key下，再在状态中累积时间戳，最后在合适时机输出完整列表，再推断缺失时间并补全。修正后的代码如下：

import apache_beam as beam
from apache_beam.coders import StrUtf8Coder
from apache_beam.transforms.userstate import BagStateSpec
from datetime import datetime, timedelta

class TimeStepAccumulatorFn(beam.DoFn):
    # 定义存储所有时间戳的Bag状态
    TIMESTAMP_BAG = BagStateSpec('timestamps', StrUtf8Coder())
    # 固定全局key，确保所有数据进入同一状态容器
    GLOBAL_KEY = 'global_key'

    def process(self, element, timestamp_bag=beam.DoFn.StateParam(TIMESTAMP_BAG)):
        # 将当前时间戳加入状态
        current_ts = element[1]['timestamp']
        timestamp_bag.add(current_ts)
        # 输出原始元素到主输出
        yield element

    def finish_bundle(self, timestamp_bag=beam.DoFn.StateParam(TIMESTAMP_BAG)):
        # Bundle处理完成时，输出完整时间戳列表到侧输出
        all_timestamps = list(timestamp_bag.read())
        if all_timestamps:
            yield beam.pvalue.TaggedOutput('all_timestamps', sorted(all_timestamps))

def infer_missing_timestamps(all_timestamps):
    # 解析时间戳为datetime对象
    dt_format = '%Y-%m-%d %H:%M:%S'
    dt_list = [datetime.strptime(ts, dt_format) for ts in all_timestamps]
    # 计算时间步长（默认取前两个时间的间隔）
    step = dt_list[1] - dt_list[0]
    # 生成完整时间序列
    full_dt_list = []
    current = dt_list[0]
    while current <= dt_list[-1]:
        full_dt_list.append(current)
        current += step
    # 找出缺失的时间戳并生成对应元素
    existing_dts = set(dt_list)
    missing_ts = [dt.strftime(dt_format) for dt in full_dt_list if dt not in existing_dts]
    return [('missing', {'timestamp': ts, 'value': 0}) for ts in missing_ts]

with beam.Pipeline(runner='DirectRunner') as pipeline:
    # 1. 生成原始数据并统一分组到全局key下
    raw_data = (
        pipeline
        | 'Create data' >> beam.Create(data)
        | 'Add global key' >> beam.Map(lambda x: (TimeStepAccumulatorFn.GLOBAL_KEY, x))
    )

    # 2. 累积时间戳，输出原始数据和完整时间列表
    processed_data, all_timestamps = (
        raw_data
        | 'Accumulate timestamps' >> beam.ParDo(TimeStepAccumulatorFn()).with_outputs('all_timestamps', main='processed')
    )

    # 3. 推断并生成缺失时间的元素
    missing_data = all_timestamps | 'Infer missing' >> beam.FlatMap(infer_missing_timestamps)

    # 4. 合并原始数据和缺失数据，整理格式并排序
    final_data = (
        (processed_data | 'Remove global key' >> beam.Map(lambda x: x[1]), missing_data)
        | 'Merge collections' >> beam.Flatten()
        | 'Sort by timestamp' >> beam.Sort(key=lambda x: x[1]['timestamp'])
    )

    # 输出结果验证
    final_data | 'Print result' >> beam.Map(print)

关键说明 #

• 全局分组：通过给所有元素添加固定global_key，确保所有时间戳累积到同一状态容器，获取完整时间列表。
• finish_bundle触发：在finish_bundle方法中读取状态并输出完整时间列表，适合批量处理场景；若为流式场景，可结合全局窗口+触发策略实现状态输出。
• 缺失时间推断：通过计算时间步长生成完整时间序列，对比现有时间戳找出缺失项，生成对应补全元素。

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