python - Apache Beam 将每个标记的输出写入单独的文件

Question

我根据输入数据元素之一(日期)标记输入元素。

class TagElementsWithDate(beam.DoFn):
    def process(self, element):
        dt = element['date'].replace('-', '')[:6]
        yield pvalue.TaggedOutput(dt, element)

input_data = p | 'Read Input' >>  beam.io.Read(beam.io.BigQuerySource(query='select id, date from `project.dataset.tablename`', use_standard_sql=True))

tagged_data = input_data | 'tag data' >> beam.ParDo(TagElementsWithDate()).with_outputs()

tagged_data 是 DoOutputsTuple。我希望对此进行迭代并将每个标记数据写入一个单独的文件。

Answer 1

您需要编写自己的 DoFn。有点像

from apache_beam.io.textio import _TextSink


class WriteEachKeyToText(beam.DoFn):
    def __init__(self, file_path_prefix=str):
        super().__init__()
        self.file_path_prefix = file_path_prefix

    def process(self, kv):
        key = kv[0]
        elements = kv[1]
        sink = _TextSink(self.file_path_prefix, file_name_suffix=f"{key}.json")

        writer = sink.open_writer("prefix", self.file_path_prefix)
        for e in elements:  # values
            writer.write(e)

然后你可以像这样使用它:

output_path = "/some/path/"
tagged_data | beam.ParDo(WriteEachKeyToText(output_path))