Apache Beam（GCP Dataflow）中PCollection复用是否会重复计算？能否缓存？

我在GCP Dataflow上使用Apache Beam，希望多次复用一个PCollection，但担心会重复计算这个计算成本高昂的PCollection。我在Apache Beam文档中找不到“materialize”或“cache”转换。

代码示例：

import apache_beam as beam

# Set up a pipeline and read in a PCollection
p = beam.Pipeline()
input_data = p | beam.io.ReadFromText('input.txt')
reused_data = input_data | beam.Map(some_expensive_function)
  
# Write the outputs to different files
reused_data | beam.io.WriteToText('output1.txt')
reused_data | beam.io.WriteToText('output2.txt')

# Run the pipeline
p.run()

请问这段代码运行时，会重复计算数据还是缓存数据？如果机器内存不足会怎样？

回答 #

• 这段代码会重复计算reused_data对应的PCollection。Apache Beam基于DAG（有向无环图）执行，当一个PCollection被多个下游转换引用时，默认会重新跑一遍上游所有计算逻辑（包括some_expensive_function）来满足每个下游分支的需求。
• 若机器内存不足：如果some_expensive_function的计算过程或生成的中间数据占满了Worker节点的可用内存，会直接触发OOM（内存溢出）错误，导致Worker进程崩溃，进而可能引发Pipeline重试甚至直接失败。

避免重复计算的实用方案 #

要复用计算成本高的PCollection，得显式把它持久化到外部存储，常用方法有两种：

1. 写入临时持久化存储：把计算结果先写到GCS、BigQuery这类外部存储，再从存储读取复用

   import apache_beam as beam
   
   p = beam.Pipeline()
   input_data = p | beam.io.ReadFromText('input.txt')
   reused_data = input_data | beam.Map(some_expensive_function)
   
   # 先把计算结果写入GCS临时路径
   reused_data | beam.io.WriteToText('gs://your-temp-bucket/reused_data_temp')
   # 从临时存储读取，生成可复用的PCollection
   cached_data = p | beam.io.ReadFromText('gs://your-temp-bucket/reused_data_temp')
   
   # 下游直接用缓存后的PCollection
   cached_data | beam.io.WriteToText('output1.txt')
   cached_data | beam.io.WriteToText('output2.txt')
   
   p.run()
   

2. 使用Reshuffle转换：针对Dataflow Runner，Reshuffle会自动把数据写入Dataflow的临时存储，间接实现持久化避免重复计算（但会引入shuffle开销，适合快速实现的场景）

   reused_data = input_data | beam.Map(some_expensive_function) | beam.Reshuffle()
   

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