如何使用 Arrow 对 CSV 进行分块？

Question

我要做什么

我正在使用 PyArrow 读取一些 CSV 并将它们转换为 Parquet。我阅读的一些文件有很多列并且内存占用很高（足以使运行该作业的机器崩溃）。我试图在读取 CSV 的同时以类似于 Pandas read_csv 和 chunksize 的工作方式对文件进行分块。

例如，这就是 pandas 中分块代码的工作方式：

chunks = pandas.read_csv(data, chunksize=100, iterator=True)

# Iterate through chunks
for chunk in chunks:
    do_stuff(chunk)

我想将类似的功能移植到 Arrow

我尝试过的事情

我注意到 Arrow 有 ReadOptions，其中包含一个 block_size 参数，我想也许我可以像这样使用它：

# Reading in-memory csv file
arrow_table = arrow_csv.read_csv(
    input_file=input_buffer,
    read_options=arrow_csv.ReadOptions(
        use_threads=True,
        block_size=4096
    )
)

# Iterate through batches
for batch in arrow_table.to_batches():
    do_stuff(batch)

由于这个 (block_size) 似乎没有返回迭代器，我的印象是这仍然会使 Arrow 读取内存中的整个表，从而重新创建我的问题。

最后，我知道我可以先使用 Pandas 读取 csv 并对其进行分块，然后转换为箭头表。但我试图避免使用 Pandas 并且只使用 Arrow。

如果需要，我很乐意提供更多信息

Answer 1

您要查找的函数是pyarrow.csv.open_csv，它返回一个pyarrow.csv.CSVStreamingReader。批次的大小将由您注意到的block_size 选项控制。完整示例：

import pyarrow as pa
import pyarrow.parquet as pq
import pyarrow.csv

in_path = '/home/pace/dev/benchmarks-proj/benchmarks/data/nyctaxi_2010-01.csv.gz'
out_path = '/home/pace/dev/benchmarks-proj/benchmarks/data/temp/iterative.parquet'

convert_options = pyarrow.csv.ConvertOptions()
convert_options.column_types = {
    'rate_code': pa.utf8(),
    'store_and_fwd_flag': pa.utf8()
}

writer = None
with pyarrow.csv.open_csv(in_path, convert_options=convert_options) as reader:
    for next_chunk in reader:
        if next_chunk is None:
            break
        if writer is None:
            writer = pq.ParquetWriter(out_path, next_chunk.schema)
        next_table = pa.Table.from_batches([next_chunk])
        writer.write_table(next_table)
writer.close()

此示例还强调了流式 CSV 阅读器引入的挑战之一。它需要返回具有一致数据类型的批次。但是，在解析 CSV 时，您通常需要推断数据类型。在我的示例数据中，文件的前几 MB 具有 rate_code 列的整数值。在批处理中间的某处，该列有一个非整数值（在本例中为*）。要解决此问题，您可以像我在这里所做的那样预先指定列的类型。