如何在Apache Beam DataFrame中无分片写入CSV至GCS？

方案一：单分区写入+GCS文件重命名 #

先通过repartition(1)强制将DataFrame转为单一分区，写入临时路径后，再通过GCS API将分片文件重命名为目标文件名：

from google.cloud import storage
import apache_beam as beam
from apache_beam.dataframe.io import read_csv

def rename_gcs_temp_file(temp_gcs_path, target_gcs_path):
    # 解析GCS路径
    temp_bucket_name, temp_blob_prefix = temp_gcs_path.replace('gs://', '').split('/', 1)
    target_bucket_name, target_blob_path = target_gcs_path.replace('gs://', '').split('/', 1)
    
    client = storage.Client()
    bucket = client.bucket(temp_bucket_name)
    
    # 查找临时路径下的CSV分片文件
    for blob in bucket.list_blobs(prefix=temp_blob_prefix):
        if blob.name.endswith('.csv'):
            # 重命名到目标路径
            bucket.rename_blob(blob, target_blob_path)
            # 清理临时文件
            bucket.delete_blobs(bucket.list_blobs(prefix=temp_blob_prefix))
            break

with beam.Pipeline(options=pipeline_options) as p:
    df = p | read_csv(known_args.input)
    df[column] = df.groupby(primary_key)[column].apply(lambda x: x.ffill().bfill())
    # 写入临时路径，强制单分区
    df.repartition(1).to_csv(f"{known_args.output}_temp", index=False, encoding='utf-8')

# 重命名临时文件到目标路径
rename_gcs_temp_file(f"{known_args.output}_temp", known_args.output)

注意：此方案适合数据量不大的场景，单分区写入可能会带来内存压力；若数据量极大，需评估worker的内存是否能承载全量数据。

方案二：转为PCollection后用WriteToText自定义分片规则 #

将DeferredDataFrame转换为CSV行的PCollection，通过WriteToText的num_shards=1和shard_name_template=''参数直接生成无后缀文件：

import apache_beam as beam
from apache_beam.dataframe.io import read_csv

def dataframe_to_csv_lines(df):
    # 将DataFrame转为CSV字符串并按行拆分
    csv_content = df.to_csv(index=False, encoding='utf-8')
    return csv_content.splitlines()

with beam.Pipeline(options=pipeline_options) as p:
    df = p | read_csv(known_args.input)
    df[column] = df.groupby(primary_key)[column].apply(lambda x: x.ffill().bfill())
    
    # 将全量DataFrame合并为单一分区并转为CSV行
    csv_lines = (
        df
        | beam.CombineGlobally(dataframe_to_csv_lines).with_output_types(str)
        | beam.FlatMap(lambda lines: lines)
    )
    
    # 写入CSV，禁用分片后缀
    csv_lines | beam.io.WriteToText(
        known_args.output,
        num_shards=1,
        shard_name_template='',
        # 若dataframe_to_csv_lines已包含表头，无需设置header参数
        # header="col1,col2,col3"
    )

注意：CombineGlobally会将全量数据集中到单个worker处理，数据量过大时可能引发OOM，需根据数据规模选择。

额外说明 #

如果数据量极大，强制单分区写入并非最优解，可与项目方确认：是否允许通过前缀匹配（如gs://path/to/file*.csv）访问分片文件，或是否可接受后续通过GCS对象组合（但CSV文件的表头重复问题需额外处理）。

内容的提问来源于stack exchange，提问作者pa-nguyen