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

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

我想做什么

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

例如这是分块代码在 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 并对其进行分块，然后转换为 Arrow 表。但我尽量避免使用 Pandas，只使用 Arrow。

如果需要，我很乐意提供其他信息

解决方法

您要查找的函数是 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 列的整数值。在批处理中间的某个地方，该列有一个非整数值（在本例中为 *）。要解决此问题，您可以像我在这里所做的那样预先指定列的类型。