PyArrow:递增地使用ParquetWriter,而无需将整个数据集保存在内存中(大于内存拼接文件)

2022-05-11 00:00:00 python parquet pyarrow apache-arrow

问题描述

我正在尝试将一个较大的拼图文件写入磁盘(大于内存)。我天真地认为我可以聪明地使用ParquetWriter和WRITE_TABLE递增地写入文件,比如(POC):

import pyarrow as pa
import pyarrow.parquet as pq
import pickle
import time

arrow_schema = pickle.load(open('schema.pickle', 'rb'))
rows_dataframe = pickle.load(open('rows.pickle', 'rb'))

output_file = 'test.parq'

with pq.ParquetWriter(
                output_file,
                arrow_schema,
                compression='snappy',
                allow_truncated_timestamps=True,
                version='2.0',  # Highest available schema
                data_page_version='2.0',  # Highest available schema
        ) as writer:
            for rows_dataframe in function_that_yields_data()
                writer.write_table(
                    pa.Table.from_pydict(
                            rows_dataframe,
                            arrow_schema
                    )
                )

但是,即使I正在生成块(在我的例子中类似于10,000行)并使用write_table,它仍然将整个数据集保留在内存中。

原来ParquetWriter将整个数据集保留在内存中,同时以增量方式写入磁盘。

有没有什么办法可以强制ParquetWriter不将整个数据集保存在内存中,或者有充分的理由根本不可能这样做?


解决方案

根据Arrow bug report中的分析,这可能是由元数据集合导致的,而元数据集合只能在文件关闭时刷新。

相关文章