Python Apache Beam WriteToText如何设置分片文件最大大小

你使用的默认WriteToText调用方式没有显式配置分片规则，Beam会根据运行时的worker数量、数据分发情况自动生成分片数，不会主动感知单文件大小，因此容易出现单文件超出预期的问题，可以通过以下几种方式实现分片大小控制：

• 手动指定总分片数（适合数据量稳定的批任务）
  WriteToText原生支持num_shards参数，你可以提前估算任务输出的总数据量（注意算上gzip压缩后的体积折损，普通JSON文本gzip压缩比通常在3:1~8:1区间），用总数据量除以你期望的单文件大小上限，算出需要的总分片数传入即可，示例代码如下：

  # 例：期望单压缩文件不超过100MB，估算总压缩后数据量为3.4GB时，设34个分片即可
  (extract_data
       | 'SaveToGCS' >> beam.io.WriteToText(
           file_path_prefix=gcs_path,
           file_name_suffix=".json.gz",
           num_shards=34
       )
  )
  

  这个方案配置简单，但如果单批次数据量波动较大，提前估算的分片数不准，还是会出现单文件大小偏离预期的情况。

• 配置分片字节阈值（Beam 2.29+版本推荐，大小控制最精准）
  高版本Beam的文件写入底层基于FileIO实现，支持直接指定单分片的最大字节数，写入过程中达到阈值会自动切分新文件，不需要提前估算总数据量，示例代码如下：

  from apache_beam.io import fileio
  
  (extract_data
       # 将字符串编码为字节格式，若上游输出已经是字节可跳过该步骤
       | 'EncodeToBytes' >> beam.Map(lambda content: content.encode("utf-8"))
       | 'WriteToGCS' >> fileio.WriteToFiles(
           path=gcs_path,
           # 保持和原有WriteToText一致的文件命名规则
           file_naming=fileio.default_file_naming(suffix=".json.gz"),
           # 指定gzip压缩的文本写入格式
           sink=lambda: fileio.TextSink(compression_type="gzip"),
           # 单分片最大100MB，可根据自己的需求调整数值
           max_bytes_per_shard=100 * 1024 * 1024
       )
  )
  

  该配置下统计的是压缩后写入GCS的实际文件大小，控制精度远高于手动指定分片数的方案，适合数据量波动大的场景。

• 注意事项
  如果你继续使用默认自动分片+gzip压缩的逻辑，不要通过强行调大worker数的方式凑分片数，这种方式生成的分片大小离散度很高，很容易出现部分分片远大于平均值的情况。

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