pandas.unique()的好奇内存消耗

Question

在分析我的算法的内存消耗时,我很惊讶有时对于较小的输入需要更多的内存.

这一切都归结为以下用法pandas.unique():

import numpy as np
import pandas as pd
import sys

N=int(sys.argv[1])

a=np.arange(N, dtype=np.int64)
b=pd.unique(a)

有N=6*10^7需要3.7GB的峰值内存,但与N=8*10^7"唯一" 3GB.

扫描不同的输入大小会产生以下图表:

出于好奇和自我教育:如何反直觉的行为(即较小的输入大小更多的内存)左右N=5*10^7,N=1.3*10^7来解释？

---

以下是在Linux上生成内存消耗图的脚本:

pandas_unique_test.py:

import numpy as np
import pandas as pd
import sys

N=int(sys.argv[1])    
a=np.arange(N, dtype=np.int64)
b=pd.unique(a)

show_memory.py:

import sys
import matplotlib.pyplot as plt   
ns=[]
mems=[]
for line in sys.stdin.readlines():
    n,mem = map(int, line.strip().split(" "))
    ns.append(n)
    mems.append(mem)
plt.plot(ns, mems, label='peak-memory')
plt.xlabel('n')
plt.ylabel('peak memory in KB')
ymin, ymax = plt.ylim()
plt.ylim(0,ymax)
plt.legend()
plt.show()

run_perf_test.sh:

WRAPPER="/usr/bin/time -f%M" #peak memory in Kb
N=1000000
while [ $N -lt 100000000 ]
do
   printf "$N "
   $WRAPPER python pandas_unique_test.py $N
   N=`expr $N + 1000000`
done 

现在:

sh run_perf_tests.sh  2>&1 | python show_memory.py

Answer 1

让我们来看看...

pandas.unique 说它是"基于哈希表的独特".

它调用此函数来获取数据的正确哈希表实现,即htable.Int64HashTable.

哈希表初始化为size_hint=值向量的长度.这意味着kh_resize_DTYPE(table, size_hint)被召唤.

这些函数在这里khash.h被定义(模板化).

它似乎(size_hint >> 5) * 4 + (size_hint) * 8 * 2为桶分配内存字节(可能更多,也许更少,我可能会在这里).

然后,HashTable.unique()被叫.

它分配一个空的Int64Vector,从128开始,它们在填充时似乎会增加四倍.

然后迭代你的值,弄清楚它们是否在哈希表中; 如果没有,它们会被添加到哈希表和向量中.(这是向量可能增长的地方;由于大小提示,哈希表不需要增长.)

最后,使NumPy ndarray指向向量.

所以,呃,我认为你看到矢量大小在某个阈值上翻了两倍(如果我的深夜数学表明,这应该是,

>>> [2 ** (2 * i - 1) for i in range(4, 20)]
[
    128,
    512,
    2048,
    8192,
    32768,
    131072,
    524288,
    2097152,
    8388608,
    33554432,
    134217728,
    536870912,
    2147483648,
    8589934592,
    34359738368,
    137438953472,
    ...,
]

希望这能为事情带来一些启示:)

Answer 2

@AKX 答案解释了为什么内存消耗在跳跃中增加，但没有解释为什么它可能会随着更多元素而减少 - 这个答案填补了空白。

pandas使用khash-map 来查找唯一元素。创建哈希映射时，数组中的元素数量用作提示：

def unique(values):
    ...
    table = htable(len(values))
    ...

然而，提示的含义是“映射中将有n个值”：

cdef class {{name}}HashTable(HashTable):

    def __cinit__(self, int64_t size_hint=1):
        self.table = kh_init_{{dtype}}()
        if size_hint is not None:
            size_hint = min(size_hint, _SIZE_HINT_LIMIT)
            kh_resize_{{dtype}}(self.table, size_hint)

然而 khash-map 理解它，因为我们至少需要n桶（而不是我们需要放置元素的地方n）：

SCOPE void kh_resize_##name(kh_##name##_t *h, khint_t new_n_buckets)
...

关于 khash-map 中的存储桶数量，有两个重要的实现细节：

• 桶的数量是2 的幂，
• 最多可以占用 77% 的桶，否则大小将增加一倍（此处+此处）

后果是什么？我们来看一个有 1023 个元素的数组：

• 哈希映射在开始时将有 1024 个桶（大于元素数量的最小的 2 次幂），但仅够容纳 ca. 800 个元素（即 1024 个元素中的 77%）。
• 添加约后。800个元素，大小将调整为2048个元素（2的下一个幂），这意味着峰值消耗将是3072（同时需要新旧数组）。

对于具有 1025 个元素的数组会发生什么？

• 哈希映射在开始处将有 2048 个存储桶（大于元素数量的最小的 2 的幂），并且足够的位置仅用于 ca。1600 个元素（即 2048 个元素的 77%）。
• 不会重新散列，峰值消耗将保持在2048，因此需要更少的内存。

每次数组大小加倍时都会发生这种内存消耗模式。这就是我们观察到的情况。

较小影响的解释：

• 内存消耗的每秒跳跃并不伴随着减少：唯一的元素存储在其大小增加四倍的向量中（如此可见，每秒大小增加一倍）。必须复制元素，这样内存就被提交/使用了两倍，从而隐藏了哈希映射的较小使用。
• 消耗在其间线性增加：随着元素被添加到唯一的向量中，越来越多的内存页面被提交：使用的np.resize零附加内存不提交它（至少在Linux上）。

这是一个小实验，表明它np.zeros(...)不会提交内存，而只会保留它：

import numpy as np
import psutil
process = psutil.Process()
old = process.memory_info().rss
a=np.zeros(10**8)
print("commited: ", process.memory_info().rss-old)
# commited:  0, i.e. nothign
a[0:100000] = 1.0
print("commited: ", process.memory_info().rss-old)
# commited:  2347008, more but not all

a[:] = 1.0
print("commited: ", process.memory_info().rss-old)
# commited:  799866880, i.e. all

注意：a=np.full(10**8, 0.0)将直接提交内存。