pandas 变量赋值后出现内存不足错误

Question

我有一个非常大的 pandas 数据框，想从中采样行进行建模，我遇到了这样的内存不足错误：

MemoryError: Unable to allocate 6.59 GiB for an array with shape (40, 22117797) and data type float64

这个错误很奇怪，因为我不需要分配这么大的内存，因为我的采样数据帧只有原始数据的 1%。下面是我的代码。

具体来说，原始数据有2000万行，大部分是np.float64数据。使用 pyarrow 从 parquet 文件加载数据后，jupyter 内核占用大约 3 GB 内存。使用“d0['r_%s'%(t)] = d0.col0”变量赋值后，内核占用 6 GB。但是，一旦我运行采样命令“d0s = d0.iloc[id1,:]”，内存就会上升到 13 GB，程序会由于上面的内存不足错误而停止。

下面的代码是使用 pandas 1.2.3.

在 16GB 内存机器上重现错误的最小工作示例

import pandas as pd
import numpy as np

d0 = pd.DataFrame(np.random.rand(22117797, 12))

for t in range(30):

     d0['r_%s'%(t)] = d0[0]

id1 = np.random.randint(low = 0, high = d0.shape[0], size = round(d0.shape[0]*0.01))

d0s = d0.iloc[id1,:]

请注意，如果我直接生成一个大数据帧，以下代码不会产生错误：

import pandas as pd
import numpy as np

d0 = pd.DataFrame(np.random.rand(22117797, 42))

id1 = np.random.randint(low = 0, high = d0.shape[0], size = round(d0.shape[0]*0.01))

d0s = d0.iloc[id1,:]

Answer 1

我发现错误是由于执行了合并操作 pandas。具体来说，在使用“d0['r_%s'%(t)] = d0[0]”进行变量赋值后，d0存储在13个块中，即：13个连续内存space，可以检查使用命令

d0._data.nblocks

一旦我运行命令“d0s = d0.iloc[id1,:]”, pandas 试图将 13 个块合并为 1 个块，所以这个新版本 d0 与1 块将占用很多内存 space，我的内存用完了。所以弹出内存不足的错误。可以使用较小的数组检查单块合并，如下所示：

import pandas as pd
import numpy as np

d0 = pd.DataFrame(np.random.rand(22117797, 12))

for t in range(10):

     d0['r_%s'%(t)] = d0[0]

d0._data.nblocks

id1 = np.random.randint(low = 0, high = d0.shape[0], size = round(d0.shape[0]*0.01))

d0s = d0.iloc[id1,:]

d0._data.nblocks

为了解决这个问题，我使用了另一种方式来获取值而不触发合并操作，如下所示：

d0s = pd.concat([d0.iloc[id1,col] for col in d0.columns], axis = 1)

希望这对遇到类似问题的其他人有所帮助。

pandas 变量赋值后出现内存不足错误

pandas out of memory error after variable assignment

python

pandas

dataframe

out-of-memory