从单个 PCollection 写入多个文件（Beam-Python）

Question

我正在处理多个文件输入，并希望将它们输入数据流管道。 但是，我想保留输出以匹配输入的数量。 假设我们有三个不同的文件：

 gcs_files = ['gs://bucket/<file1_dir>', 'gs://bucket/<file2_dir>', 'gs://bucket/<file3_dir>']

我们希望将这些文件输入一个简单的 Read->Write 管道，如下所示：

 import apache_beam as beam
 from apache_beam.io import ReadAllFromText
 from apache_beam.io import WriteToText
 from apache_beam.options.pipeline_options import PipelineOptions
 
 #Please assume I am using default pipeline options
 print_files = (p | beam.Create(gcs_files) | ReadAllFromText() | WriteToText(<output_path>, shard_name_template = '', file_name_suffix = '.json')

这个管道的目的很简单，但我正在努力转换输入文件的格式。 因此，我们期望每个输入文件都有一个对应的转换文件。 当前设置的问题是输出都是一个文件。 有什么办法可以保持原始文件的完整性？

Answer 1

您可以ReadAllFromText(with_filename=True) ，然后它为您提供元素的键值对，其中键是原始文件名。 但是，在转换为值之后，您必须使用WriteToFiles根据原始文件名动态写入目标。

或者，如果您的 gcs_files 的大小很小，您可以在构建时完成所有操作：

pcolls = {}
for gcs_file in gcs_files:
    pcolls[gcs_file] = (p 
        | f'Create {gcs_file}' >> beam.Create([gcs_file]) 
        | f'Read {gcs_file}' >> ReadAllFromText()
        | f'Your Transform {gcs_file}' >> YourTransform()
        | f'Write {gcs_file}' >> WriteToText())

Answer 2

根据您的要求，您可以创建一个 DoFn 对象，该对象可以传递给ParDo ，用于转换输入 PCollection。 DoFn 包含必须应用于输入集合的处理逻辑。ReadAllFromText 用于从 PCollection 中读取，它可以将输入作为键值对，WriteToText 可用于对本文档中给出的输出文件进行分片。您可以参考代码此链接中给出。

但是，如果您正在分组执行转换，例如，如果您想获得一些 Avro 格式的转换和一些 CSV 格式的转换，那么您可以使用WriteToFiles 。