时间戳

首页 > 解决方案 > 如何使用 Arrow 对 CSV 进行分块？

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

问题描述

我想要做什么

我正在使用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 具有包含参数的ReadOptionsblock_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)

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

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

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

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

推荐阅读