如何在Dataflow中实现GCS CSV文件并行读取以提升流水线性能

• 步骤融合导致任务未分发：Beam默认会将无Shuffle操作的连续步骤（MatchFiles→ReadMatches→FileToRowsFn）融合到同一执行单元，1000个文件记录被绑定在少数Worker进程中串行执行，无法触发节点扩容
• 自动扩缩容策略不匹配：Dataflow默认基于CPU使用率触发扩容，该场景为IO密集型，CPU使用率仅2%~3%，自动扩容逻辑不会新增Worker节点
• 你传入的普通字典参数不属于Beam SideInput范畴，不会影响并行化，无需担心

1. 打破步骤融合强制并行 #

在ReadMatches后插入beam.Reshuffle()操作，触发Shuffle打散文件任务，让每个文件处理任务可以分发到独立Worker节点并行执行，修改后代码如下：

additional_side_inputs = {'key1': 'value1', 'key2': 'value2'}  # etc.

p | 'Collect CSV files' >> MatchFiles(input_dir + "*.csv")
  | 'Read files' >> ReadMatches()
  | 'Reshuffle files for parallelism' >> beam.Reshuffle()
  | 'Parse contents' >> beam.ParDo(FileToRowsFn(), additional_side_inputs)
  | 'Compute average' >> beam.CombinePerKey(AverageCalculatorFn())

2. 调整Dataflow运行参数 #

提交流水线时添加以下参数适配IO密集场景：

• --autoscaling_algorithm=THROUGHPUT_BASED：改为吞吐量优先的扩缩容策略，基于待处理任务数而非CPU使用率扩容
• --num_workers=20：手动设置初始Worker数量，可根据预期并行度调整
• --max_num_workers=100：设置Worker数量上限，避免过度扩容

3. 优化文件读取性能 #

当前50MB文件读取耗时2分钟属于异常，可优化FileToRowsFn的读取缓冲配置：

class FileToRowsFn(beam.DoFn):
  def process(self, file_element, additional_side_inputs):
    with file_element.open() as csv_file:
      # 增加1MB读取缓冲，大幅提升大文件读取速度
      wrapper = TextIOWrapper(csv_file, encoding='utf-8', buffering=1024*1024)
      for row_id, *values in csv.reader(wrapper):
        yield row_id, values

4. （可选）使用内置IO简化实现 #

直接用Beam内置的ReadFromText读取通配符路径的CSV文件，内置实现已经做了并行拆分优化，无需手动处理文件IO：

from apache_beam.io import ReadFromText
import csv

class ParseRowFn(beam.DoFn):
  def process(self, line, additional_side_inputs):
    row_id, *values = next(csv.reader([line]))
    yield row_id, values

additional_side_inputs = {'key1': 'value1', 'key2': 'value2'}
p | 'Read all CSV' >> ReadFromText(input_dir + "*.csv")
  | 'Parse rows' >> beam.ParDo(ParseRowFn(), additional_side_inputs)
  | 'Compute average' >> beam.CombinePerKey(AverageCalculatorFn())

以上优化完成后，1000个文件的读取时间可从原有的1000*2分钟缩短至数分钟，整体流水线耗时可控制在1小时以内。

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