Dataflow中CSV转JSON格式问题及Pipeline优化咨询

1. 生成数组包裹的JSON格式 #

当前输出是每行一个独立JSON对象，要生成数组格式，需先将所有JSON对象收集为一个列表，再序列化为标准JSON数组。

修改Pipeline，添加合并与格式化步骤：

import json
import apache_beam as beam
from apache_beam.io import ReadFromText, WriteToText

class ConvertCsvToJson(beam.DoFn):
    def process(self, element):
        col1, col2, col3 = element.split(',')
        json_obj = {
            'col1': col1.strip(),
            'col2': col2.strip(),
            'col3': col3.strip()
        }
        yield json_obj

class FormatJsonArray(beam.DoFn):
    def process(self, elements):
        # 将列表序列化为带缩进的JSON数组
        yield json.dumps(list(elements), indent=2)

def run_pipeline(input_file, output_file):
    with beam.Pipeline() as pipeline:
        (pipeline
         | 'Read CSV' >> ReadFromText(input_file, skip_header_lines=1)
         | 'Convert to JSON' >> beam.ParDo(ConvertCsvToJson())
         | 'Combine into List' >> beam.CombineGlobally(lambda elems: list(elems))
         | 'Format JSON Array' >> beam.ParDo(FormatJsonArray())
         | 'Write JSON' >> WriteToText(output_file, shard_name_template='')  # 避免生成分片文件
        )

关键说明：

• CombineGlobally 将所有分散的JSON对象合并为一个列表
• FormatJsonArray 负责将列表转成标准JSON数组字符串
• shard_name_template='' 确保输出为单个完整文件，而非多个分片

2. 基于CSV表头自动生成键名 #

无需硬编码col1/col2，可先读取CSV表头，再将表头与每行数据配对生成JSON键值对。同时建议用csv模块解析，避免逗号在引号内的解析错误：

import csv
import json
from io import StringIO
import apache_beam as beam
from apache_beam.io import ReadFromText, WriteToText

class ConvertCsvToJsonWithHeader(beam.DoFn):
    def process(self, element):
        _, (header_list, lines) = element
        for line in lines:
            # 用csv模块解析行，兼容带引号的逗号分隔值
            reader = csv.reader(StringIO(line))
            cols = next(reader)
            # 表头与列值一一对应生成JSON对象
            json_obj = dict(zip(header_list, [col.strip() for col in cols]))
            yield json_obj

class FormatJsonArray(beam.DoFn):
    def process(self, elements):
        yield json.dumps(list(elements), indent=2)

def run_pipeline(input_file, output_file):
    with beam.Pipeline() as pipeline:
        # 单独读取CSV表头行
        header = (
            pipeline
            | 'Read Header' >> ReadFromText(input_file, limit=1)
            | 'Parse Header' >> beam.Map(lambda line: line.strip().split(','))
        )
        # 读取数据行并与表头关联
        (pipeline
         | 'Read CSV Lines' >> ReadFromText(input_file, skip_header_lines=1)
         | 'Add Key for Join' >> beam.WithKeys(lambda _: 'data')  # 统一key用于关联表头
         | 'Join with Header' >> beam.CoGroupByKey()
         | 'Convert to JSON' >> beam.ParDo(ConvertCsvToJsonWithHeader())
         | 'Combine into List' >> beam.CombineGlobally(lambda elems: list(elems))
         | 'Format JSON Array' >> beam.ParDo(FormatJsonArray())
         | 'Write JSON' >> WriteToText(output_file, shard_name_template='')
        )

关键说明：

• 通过CoGroupByKey将表头与所有数据行合并，确保每个转换步骤都能获取表头信息
• 使用csv.reader解析行数据，处理包含逗号的字段（如"Smith, John"）

3. 优化Pipeline运行速度 #

针对5分钟的耗时，可从以下几点优化：

• 调整并行度：在PipelineOptions中设置worker数量，提升并行处理能力：

  from apache_beam.options.pipeline_options import PipelineOptions, WorkerOptions
  
  options = PipelineOptions()
  worker_options = options.view_as(WorkerOptions)
  worker_options.worker_count = 8  # 根据机器CPU核心数调整
  

• 替换低效解析方式：用csv模块替代split(',')，解析更高效且容错性强
• 合并小输入文件：如果输入是大量小CSV文件，先合并为大文件，减少IO开销
• 资源配置升级：使用Dataflow运行时，选择更高配置的机器类型（如n2-standard-4），提升单worker处理能力
• 减少冗余步骤：尽量将多个操作合并到一个DoFn中，减少数据在Pipeline中的传递次数

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