Beam-Python单个PCollection如何输出多文件匹配输入文件数量

你当前的流水线写法中，ReadAllFromText默认会把所有输入文件的文本行打散为无来源标记的扁平化PCollection，完全丢失了每行记录所属的原文件关联关系，后续WriteToText接收到的是混在一起的全量数据，自然会合并输出，无法和输入文件形成一一对应关系。

核心要做两个关键调整：

• 读取阶段保留每条记录对应的原文件标识，不要丢失来源关联
• 写入阶段按原文件维度做分桶隔离，不同文件的数据分别写入独立的输出文件，避免混写

以下是可直接运行的实现代码，全程流式处理，不需要提前把整个文件加载到内存，适配大文件场景：

import apache_beam as beam
from apache_beam.io import ReadAllFromText, fileio
from apache_beam.options.pipeline_options import PipelineOptions

def map_input_to_output(original_gcs_path: str, output_base_dir: str) -> str:
    # 从原GCS路径提取文件名，拼接生成对应的输出文件路径
    original_filename = original_gcs_path.rsplit("/", 1)[-1]
    return f"{output_base_dir.rstrip('/')}/{original_filename}.json"

p = beam.Pipeline(options=PipelineOptions())
gcs_files = ['gs://bucket/<file1_dir>', 'gs://bucket/<file2_dir>', 'gs://bucket/<file3_dir>']
OUTPUT_STAGING_DIR = 'gs://bucket/<output_staging_dir>'
FINAL_OUTPUT_DIR = 'gs://bucket/<final_output_dir>'

(
    p
    | "初始化输入文件列表" >> beam.Create(gcs_files)
    # 开启with_filename参数，读取返回(原文件路径, 文本行)的元组，保留来源信息
    | "逐行读取文件并打来源标记" >> ReadAllFromText(with_filename=True)
    # 给每行数据附加对应的输出路径作为分桶标记
    | "标记每行的输出目的地" >> beam.Map(
        lambda file_line_pair: (
            map_input_to_output(file_line_pair[0], FINAL_OUTPUT_DIR),
            file_line_pair[1]
        )
    )
    # 按输出路径分桶写入，每个分桶对应一个独立输出文件，和输入文件一一对应
    | "按目标文件分桶写入" >> fileio.WriteToFiles(
        path=OUTPUT_STAGING_DIR,
        destination=lambda tagged_record: tagged_record[0],
        sink=lambda dest: fileio.TextSink(),
        file_naming=lambda window, pane, shard_idx, total_shards, compression, dest_path: dest_path.split("/")[-1],
        shards=1  # 强制每个分桶只生成1个文件分片，不自动拆分
    )
)

p.run().wait_until_finish()

• shards=1和你之前配置的shard_name_template=''作用一致，保证每个输入文件对应生成1个输出文件，不会被Beam框架自动拆分为多个分片。
• 如果你需要对读取到的文本行做格式转换（比如转JSON结构），直接在「标记每行的输出目的地」步骤之后加Map转换逻辑即可，不会破坏文件和记录的对应关系。
• 小文件场景下也可以用fileio.MatchAll()+fileio.ReadMatches()的简化写法，代码更简洁，但大文件场景下优先用上面的逐行流式读取方案，避免Worker内存溢出。

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