迭代 Pyarrow Table 的最快方法

Fastest way to iterate Pyarrow Table

我正在使用 Pyarrow 库来优化 Pandas DataFrame 的存储。我需要尽可能快地逐行处理 pyarrow Table 而不将其转换为 pandas DataFrame(它不适合内存)。 Pandas 有 iterrows()/iterrtuples() 方法。除了 for 循环和索引寻址之外,是否有任何快速迭代 Pyarrow Table 的方法?

软件目前完全没有针对这个用例进行优化。我建议使用 Cython 或 C++ 或逐行与数据交互。如果您还有其他问题,请联系开发者邮件列表 dev@arrow.apache.org

这段代码对我有用:

for batch in table.to_batches():
    d = batch.to_pydict()
    for c1, c2, c3 in zip(d['c1'], d['c2'], d['c3']):
        # Do something with the row of c1, c2, c3

如果您有一个大型 parquet 数据集拆分为多个文件,这似乎相当快且内存效率高。

import argparse
import pyarrow.parquet as pq
from glob import glob


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('parquet_dir')
    return parser.parse_args()


def iter_parquet(dirpath):
    for fpath in glob(f'{dirpath}/*.parquet'):
        tbl = pq.ParquetFile(fpath)

        for group_i in range(tbl.num_row_groups):
            row_group = tbl.read_row_group(group_i)

            for batch in row_group.to_batches():
                for row in zip(*batch.columns):
                    yield row


if __name__ == '__main__':
    args = parse_args()

    total_count = 0
    for row in iter_parquet(args.parquet_dir):
        total_count += 1
    print(total_count)