使用 Apache Beam 将文件写入 Parquet 中的动态目标 Python SDK

Question

我正在尝试通过WriteToFiles class 使用动态目的地编写 Parquet 文件。

我什至发现了一些进一步开发的例子，比如这个，他们构建了一个自定义的 Avro 文件接收器。

我目前正在尝试使用pyarrow库编写一个 Parquet 接收器，该接收器可以以分布式方式管理写入操作，类似于 WriteToParquet PTransform 的方式。

class ParquetFileSink(fileio.FileSink):
    def __init__(self, schema, codec='deflate'):
        self._schema = schema
        self._codec = codec
        self.writer = None

    def open(self, fh):
        # This is called on every new bundle.
        self.writer = pq.ParquetWriter(
            fh,
            self._schema,
            compression=self._codec,
            use_deprecated_int96_timestamps=False
        )

    def write(self, record):
        # This is called on every element.
        row = pa.Table.from_pandas(
            pd.DataFrame(record), schema=self._schema, preserve_index=False
        )
        self.writer.write_table(row)

    def flush(self):
        pass

这里的主要问题是，据我所知，不可能将无界 PCollections 写为 Parquet 文件，所以如果我尝试使用以下 class 通过记录写入，我会收到在关闭文件处理程序上写入错误，或者某些文件根本没有创建。 我还尝试使用GroupByKey PTransform 编写批处理，但是由于无法关闭pyarrow.parquet.ParquetWriter object，文件最终只被部分写入并被损坏。 此外，这种策略并不安全，因为批处理可能非常大，并且将它们写为单个文件并不是一个好主意。

I can see that this problem is being faced in the class apache_beam.io.parquetio._ParquetSink , but I don't think this can be directly applied to the WriteToFiles class as I can't see how to fully manage file handlers with it.

Answer 1

我遇到了类似的问题，最后我编写了一个可以与WriteToFiles一起使用的ParquetSink 。 因此，它会根据您的配置批量处理 memory 中的记录。 我已经使用它在依赖于记录中的字段的批处理过程中创建动态文件，但我认为它也可以与流式管道一起使用，尽管我尚未对其进行测试。

您可以在此要点中找到代码

class ParquetSink(fileio.FileSink):
    def __init__(self,
                file_path_prefix,
                schema,
                row_group_buffer_size=64 * 1024 * 1024,
                record_batch_size=1000,
                codec='none',
                use_deprecated_int96_timestamps=False,
                file_name_suffix='',
                num_shards=0,
                shard_name_template=None,
                mime_type='application/x-parquet'):
        self._inner_sink = beam.io.parquetio._create_parquet_sink(
            file_path_prefix,
            schema,
            codec,
            row_group_buffer_size,
            record_batch_size,
            use_deprecated_int96_timestamps,
            file_name_suffix,
            num_shards,
            shard_name_template,
            mime_type
        )
        self._codec = codec
        self._schema = schema
        self._use_deprecated_int96_timestamps = use_deprecated_int96_timestamps

    def open(self, fh):
        self._pw = pyarrow.parquet.ParquetWriter(
            fh,
            self._schema,
            compression=self._codec,
            use_deprecated_int96_timestamps=self._use_deprecated_int96_timestamps)

    def write(self, record):
        self._inner_sink.write_record(self._pw, record)

    def flush(self):
        if len(self._inner_sink._buffer[0]) > 0:
            self._inner_sink._flush_buffer()
        if self._inner_sink._record_batches_byte_size > 0:
            self._inner_sink._write_batches(self._pw)

        self._pw.close()



   def parquet_compatible_filenaming(suffix=None):
    def _inner(window, pane, shard_index, total_shards, compression, destination):
        return fileio.destination_prefix_naming(suffix )(
            window, pane, shard_index, total_shards, compression, destination).replace(":", ".")

    return _inner


def get_parquet_pipeline(pipeline_options, input, output):
    with beam.Pipeline(options=pipeline_options) as p:
        lines = (p 
                    | 'Read' >> beam.io.ReadFromParquet(file_pattern=input)
                    | 'Transform' >> beam.Map(lambda x: { 'some_key': x['some_key'], 'raw': x})
                    | 'Write to Parquet' >> fileio.WriteToFiles(
                                                path=str(output),
                                                destination=lambda x: x["some_key"],
                                                sink=lambda x: ParquetSink(
                                                                    file_path_prefix=output,
                                                                    file_name_suffix=".parquet",
                                                                    codec="snappy",
                                                                    schema=pyarrow.schema([
                                                                        pyarrow.field("some_key", pyarrow.string()),
                                                                        pyarrow.field("raw", pyarrow.string())
                                                                    ])),
                                                    file_naming=parquet_compatible_filenaming(suffix=".parquet")
                                                )
        )

Answer 2

parquet 格式针对批量写入数据进行了优化。 因此，它不适合流式传输，您可以在流式传输中一一接收记录。 在您的示例中，您在镶木地板文件中逐行编写，这是非常低效的。

我建议以适合于逐行附加数据的格式保存您的数据，然后进行常规作业，将这些数据批量移动到镶木地板文件中。

或者你可以像apache_beam.io.parquetio._ParquetSink一样。 它将 memory 中的记录保存在缓冲区中，并时不时地批量写入。 但是这样一来，如果您的应用程序崩溃，您就有丢失缓冲区中记录的风险。