Pyarrow：将流读入 pandas 数据帧高内存消耗

Question

我想先将流写入箭头文件，然后再将其读回 pandas 数据帧，尽可能减少内存开销。

批量写入数据非常好：

import pyarrow as pa
import pandas as pd
import random

data = [pa.array([random.randint(0, 1000)]), pa.array(['B']), pa.array(['C'])]
columns = ['A','B','C']
batch = pa.RecordBatch.from_arrays(data, columns)

with pa.OSFile('test.arrow', 'wb') as f:
    with pa.RecordBatchStreamWriter(f, batch.schema) as writer:
        for i in range(1000 * 1000):
            data = [pa.array([random.randint(0, 1000)]), pa.array(['B']), pa.array(['C'])]
            batch = pa.RecordBatch.from_arrays(data, columns)
            writer.write_batch(batch)

如上写入100万行速度很快，整个写入过程使用大约40MB内存。这很好。

然而，在生成大约 118MB 的最终数据帧之前，内存消耗高达 2GB，因此回读并不好。

我试过这个：

with pa.input_stream('test.arrow') as f:
    reader = pa.BufferReader(f.read())
    table = pa.ipc.open_stream(reader).read_all()
    df1 = table.to_pandas(split_blocks=True, self_destruct=True)

和这个，具有相同的内存开销：

with open('test.arrow', 'rb') as f:
   df1 = pa.ipc.open_stream(f).read_pandas()

数据帧大小：

print(df1.info(memory_usage='deep'))

Data columns (total 3 columns):
 #   Column  Non-Null Count    Dtype
---  ------  --------------    -----
 0   A       1000000 non-null  int64
 1   B       1000000 non-null  object
 2   C       1000000 non-null  object
dtypes: int64(1), object(2)
memory usage: 118.3 MB
None

我需要的是用 pyarrow 解决内存使用问题，或者建议我可以使用哪种其他格式来增量写入数据，然后将所有数据读入 pandas 数据帧并且不要太多内存开销。

Answer 1

您的示例使用了许多 RecordBatches，每一行都是 RecordBatches。这样的 RecordBatch 除了数据（模式，潜在的 padding/alignment）之外还有一些开销，因此对于仅存储单行来说效率不高。

当使用 read_all() 或 read_pandas() 读取文件时，它首先创建所有这些 RecordBatch，然后再将它们转换为单个 Table。然后开销加起来，这就是您所看到的。

RecordBatch 的推荐大小当然取决于确切的用例，但典型大小为 64k 到 1M 行。

要查看每个数组填充到 64 字节的效果 (https://arrow.apache.org/docs/format/Columnar.html#buffer-alignment-and-padding)，让我们检查分配的总字节数与 RecordBatch 表示的实际字节数：

import pyarrow as pa
 
batch = pa.RecordBatch.from_arrays(
    [pa.array([1]), pa.array(['B']), pa.array(['C'])],
    ['A','B','C']
)

# The size of the data stored in the RecordBatch
# 8 for the integer (int64), 9 for each string array (8 for the len-2 offset array (int32), 1 for the single string byte)
>>> batch.nbytes
26

# The size of the data actually being allocated by Arrow
# (5*64 for 5 buffers padded to 64 bytes)
>>> pa.total_allocated_bytes()
320

所以你可以看到，仅仅这个填充就已经给一个很小的 RecordBatch 带来了很大的开销

Pyarrow：将流读入 pandas 数据帧高内存消耗

Pyarrow: read stream into pandas dataframe high memory consumption

python

pandas

optimization

pyarrow

feather