Apache Beam Python SDK如何通过fileio模块写入压缩文件至GCS

你对Apache Beam 2.16.0版本fileio模块的判断完全正确。

• 该版本的fileio.WriteToFiles原生没有内置压缩写入能力，_MoveTempFilesIntoFinalDestinationFn类的注释也明确标注：框架仅负责临时文件到最终存储路径的移动操作，压缩、编码这类扩展处理逻辑需要开发者自行实现。

可直接落地的压缩写入实现方案 #

不需要从零重构写入流程，利用WriteToFiles开放的sink、file_naming扩展点，就能快速接入压缩能力，完全兼容你现有的窗口、分片配置。

第一步：实现支持压缩的自定义Sink #

继承fileio.FileSink重写IO方法，把普通文件流包装成压缩写入流即可，以最常用的gzip格式为例：

import gzip
from apache_beam.io import fileio

class GzipLogSink(fileio.FileSink):
    def __init__(self, compresslevel=6):
        self.compresslevel = compresslevel
        self._compress_wrapper = None

    def open(self, raw_file_handle):
        self._compress_wrapper = gzip.GzipFile(
            fileobj=raw_file_handle,
            mode="wb",
            compresslevel=self.compresslevel
        )

    def write(self, log_record):
        # 适配上游输出的JSON字符串，转字节后写入
        if isinstance(log_record, str):
            log_record = log_record.encode("utf-8")
        self._compress_wrapper.write(log_record + b"\n")

    def flush(self):
        self._compress_wrapper.flush()

    def close(self):
        self._compress_wrapper.close()

第二步：配置压缩文件命名规则 #

给输出文件加上.gz后缀，方便后续读取时自动识别压缩格式：

def gzip_log_naming(window, pane, shard_idx, total_shards, compression, dest):
    # 可根据自身业务调整文件名前缀、时间格式
    time_part = window.start.to_utc_datetime().strftime("%Y%m%d_%H00")
    return f"{time_part}-{shard_idx:05d}-of-{total_shards:05d}.json.gz"

第三步：修改原有写入逻辑 #

替换默认的Sink和命名规则即可，原有窗口、分片参数不需要改动：

_ = (logs 
    | 'Window' >> beam.WindowInto(beam.window.FixedWindows(60*60))
    | 'Convert to JSON' >> beam.ParDo(ConvertToJson())
    | 'Write logs to GCS file' >> fileio.WriteToFiles(
        path = gsc_output_path, 
        shards=1, 
        max_writers_per_bundle=0,
        sink=lambda _: GzipLogSink(),
        file_naming=gzip_log_naming
    )
)

补充说明 #

• 如果后续升级到Beam 2.28.0及以上版本，WriteToFiles已经原生内置压缩支持，只需要额外传入compression="gzip"参数即可，不需要再自定义Sink。
• 上述自定义Sink的逻辑可以通用适配bzip2、zstd等其他压缩格式，只需要把gzip流替换成对应压缩库的流写入对象即可。
• 写入GCS的gzip压缩日志，存储成本通常比未压缩文件低70%左右，用Beam、GCS原生客户端读取时会自动解压，不需要额外做格式适配。

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