Apache Beam入门性能问题排查：速度慢与protobuf超2GB报错

1. 解决exceeds maximum protobuf size of 2GB报错 #

这个错误直接来源于beam.combiners.ToList()操作——它会把整个PCollection的所有元素合并成一个单一对象，当处理几十万行数据时，这个对象的大小很容易突破Protobuf的2GB限制。

• 优化：移除ToList()操作，直接对每个元素执行保存逻辑，改用官方分布式写入API替代自定义的hacky_save：

  # 替换原来的final_docs和res部分
  final_docs = (
      {"features": pages_pcoll, "lines": lines_to_keep}
      | "group_features_and_lines_by_url" >> beam.CoGroupByKey()
      | beam.FlatMap(_remove_lines_from_text, min_num_sentences=min_num_sentences)
  )
  # 用官方分布式写入替代hacky_save
  final_docs | "Write to JSON" >> beam.io.WriteToJson(
      file_path_prefix="output/c4_docs",
      file_name_suffix=".json",
      shard_name_template="-SSSSS-of-NNNNN"
  )
  

2. 修复beam.Create(pages)的性能瓶颈 #

如果pages是包含几十万条数据的大列表，beam.Create(pages)会把所有数据一次性加载到Driver进程中，导致启动慢、Driver内存压力大，还无法利用分布式并行加载能力。

• 优化：将pages存储为JSONL格式文件，改用beam.io.ReadFromText分布式加载：

  # 假设pages已经导出为jsonl文件
  pages_pcoll = (
      pipeline
      | beam.io.ReadFromText("input/pages_data.jsonl")
      | beam.Map(json.loads)  # 解析每行的JSON数据
      | beam.Map(lambda x: (x[0], x[1]))  # 转成原有的(string, dict)键值对
  )
  

3. 修正CoGroupByKey的错误用法 #

代码中{"features": pages, "lines": lines_to_keep}直接将Python列表pages传入CoGroupByKey，会导致pages的所有数据被广播到每个Worker节点，造成严重的数据冗余和内存浪费，这是性能慢的核心原因之一。

• 优化：必须将pages转换为PCollection（参考上面的pages_pcoll），确保两个输入都是分布式的键值对PCollection：

  final_docs = (
      {"features": pages_pcoll, "lines": lines_to_keep}
      | "group_features_and_lines_by_url" >> beam.CoGroupByKey()
      | beam.FlatMap(_remove_lines_from_text, min_num_sentences=min_num_sentences)
  )
  

4. 调整运行时参数提升并行效率 #

当前使用的--autoscalingAlgorithm=BASIC扩缩容逻辑偏保守，适合CPU密集型任务，对于批处理数据场景，更推荐使用吞吐量驱动的扩缩容：

• 替换参数：--autoscalingAlgorithm=THROUGHPUT_BASED --maxWorkers=32
• 若使用Dataflow，可添加--diskSizeGb=50（根据数据量调整），避免Worker节点磁盘不足影响性能。

额外建议 #

• 检查_emit_url_to_lines和_remove_lines_from_text函数，确保没有在函数内部执行全局状态修改、大对象创建等串行操作，保证每个元素的处理是独立可并行的。
• 启用Beam的监控功能（如Dataflow控制台），查看Worker的CPU、内存使用情况，定位是否有节点资源瓶颈。

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