Apache Beam：窗口关闭后延迟执行DoFn的编码异常排查

我有实时数据以同步方式发布并读取到Dataflow管道中，将数据按1秒固定窗口做累积式窗口处理后写入Firestore数据库，前端监听该数据库并自动拉取新数据快照，但目前前端数据到达不同步，延迟情况不一致。

我研究了有状态和定时函数，认为可解决该问题，因未批量处理数据，我理解应使用定时DoFn，但不知如何实现。数据已按Key分组并通过自定义函数聚合，传入DoFn的是包含多种数据结构的嵌套Python字典。

当前代码如下：

WriteToFirestore DoFn 代码 #

# class to write to firestore using DoFn
class WriteToFirestore(beam.DoFn):
    TIMER = userstate.TimerSpec('timer', userstate.TimeDomain.WATERMARK)
    BUFFER_STATE = userstate.BagStateSpec('buffer', WindowedValueCoder(MapCoder(str, str)))

    def setup(self):
        from google.cloud import firestore
        self.firestore_client = firestore.Client(project='project ID')

    def process(self, element,
            timer=beam.DoFn.TimerParam(TIMER),
            window=beam.DoFn.WindowParam,
            buffer_state=beam.DoFn.StateParam(BUFFER_STATE)):
        logging.info("PROCESSING FUNCTION")

        timer.set(window.end + Duration(0.5)) # expire 0.5 seconds after window closes

        buffer_state.add(element)

    @userstate.on_timer(TIMER)
    def expiry(self):
        logging.info("FIRED")
        #collection_ref = self.firestore_client.collection('collection id')
        #collection_ref.add(event)

Pipeline 步骤代码 #

# define pipeline steps
data = (
    pipeline
    | 'Read from PubSub' >> beam.io.ReadFromPubSub(topic=self.topic_path)
    | 'Proto to Dict' >> beam.Map(lambda pb_message: proto_to_dict(pb_message, msgs_pb2.Message))
    | 'Add Key' >> beam.Map(lambda message: add_group_key(message, SETTINGS))
    | 'Window' >> beam.WindowInto(
        window.FixedWindows(self.window_length),
        accumulation_mode=beam.trigger.AccumulationMode.ACCUMULATING,
    )
    | 'Group' >> beam.GroupByKey()
    | 'Fuse' >> beam.Map(lambda group_messages: fuse(group_messages, SETTINGS, NODES))
)
        
# branch 1: Write to BigQuery
data | 'Write to BigQuery' >> beam.io.WriteToBigQuery(
    'table ID',
     dataset='dataset ID',
     schema=TABLE_SCHEMA,
     write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND,
     create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED
)
        
# branch 2: Publish to PubSub
data | 'Write to Firestore' >> beam.ParDo(WriteToFirestore())

运行时抛出**"Number of components doesn't match number of coders"**错误，我认为仅需一个组件，请问我哪里出错了？

1. 编码错误的根源 #

你在定义BUFFER_STATE时犯了两个关键错误：

• 错误使用WindowedValueCoder包裹MapCoder：BagStateSpec只需要直接对应元素类型的编码器，Dataflow会自动处理窗口值的封装，无需手动指定WindowedValueCoder。
• MapCoder(str, str)无法适配嵌套字典：该编码器仅支持键值均为字符串的单层字典，无法处理你传入的嵌套数据结构。

2. 修正后的WriteToFirestore DoFn #

方案一：使用JSON编码器（适配嵌套字典） #

如果你的嵌套字典可序列化为JSON，这是最简单的解决方案：

from apache_beam.coders import coders

class WriteToFirestore(beam.DoFn):
    TIMER = userstate.TimerSpec('timer', userstate.TimeDomain.WATERMARK)
    # 替换为适配嵌套结构的JsonCoder
    BUFFER_STATE = userstate.BagStateSpec('buffer', coders.JsonCoder())

    def setup(self):
        from google.cloud import firestore
        self.firestore_client = firestore.Client(project='project ID')

    def process(self, element,
                timer=beam.DoFn.TimerParam(TIMER),
                window=beam.DoFn.WindowParam,
                buffer_state=beam.DoFn.StateParam(BUFFER_STATE)):
        logging.info("PROCESSING FUNCTION")
        timer.set(window.end + Duration(0.5))
        buffer_state.add(element)

    @userstate.on_timer(TIMER)
    def expiry(self, buffer_state=beam.DoFn.StateParam(BUFFER_STATE)):
        logging.info("FIRED")
        # 读取并批量写入Firestore
        elements = list(buffer_state.read())
        collection_ref = self.firestore_client.collection('collection id')
        for elem in elements:
            collection_ref.add(elem)
        # 清空状态避免重复写入
        buffer_state.clear()

方案二：自定义编码器（处理特殊类型） #

如果数据包含JSON无法序列化的类型（如datetime），可自定义编码器：

from apache_beam.coders import Coder
import pickle

class NestedDictCoder(Coder):
    def encode(self, value):
        return pickle.dumps(value)

    def decode(self, value):
        return pickle.loads(value)

class WriteToFirestore(beam.DoFn):
    TIMER = userstate.TimerSpec('timer', userstate.TimeDomain.WATERMARK)
    BUFFER_STATE = userstate.BagStateSpec('buffer', NestedDictCoder())

    # 其余方法同方案一

3. 前端延迟不一致的额外优化 #

修复编码错误后，可进一步优化前端数据同步问题：

• 切换定时器时间域：若存在数据迟到情况，将WATERMARK改为REAL_TIME，确保定时器准时触发：

  TIMER = userstate.TimerSpec('timer', userstate.TimeDomain.REAL_TIME)
  

• 调整窗口触发器：使用AfterWatermark()确保每个窗口只输出最终聚合结果，减少Firestore写入频率：

  | 'Window' >> beam.WindowInto(
      window.FixedWindows(self.window_length),
      accumulation_mode=beam.trigger.AccumulationMode.ACCUMULATING,
      trigger=beam.trigger.AfterWatermark(),
  )
  

• 批量写入Firestore：使用firestore.Client().batch()批量提交写入请求，降低网络开销，提升一致性。

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