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。这样的RecordBatch除了数据(模式、潜在的填充/对齐(之外还有一些开销，因此对于只存储一行来说效率不高。

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

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

---

查看填充到每个数组64字节的效果(https://arrow.apache.org/docs/format/Columnar.html#buffer-对齐和填充(，让我们检查分配的总字节数与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 带来了巨大的开销