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%的bucket，否則大小會翻倍（ 這里+ 這里）

后果是什么？ 讓我們看一個有 1023 個元素的數組：

• 哈希映射在開始時將有 1024 個桶（比元素數大的 2 的最小冪），但足夠的地方僅用於 ca。 800 個元素（即 77% 來自 1024）。
• 添加ca后。 800個元素，大小將調整為2048個元素（2的下一次冪），這意味着峰值消耗將是3072 （同時需要新舊陣列）。

一個有 1025 個元素的數組會發生什么？

• 哈希映射將在開頭有 2048 個桶（比元素數大的 2 的最小冪），並且足夠的位置僅用於 ca。 1600 個元素（即 77% 來自 2048）。
• 不會重新散列，峰值消耗將保持在2048 ，因此需要更少的內存。

這種內存消耗模式將隨着數組大小的每增加一倍而發生。 這是我們觀察到的。

對較小影響的解釋：

• 內存消耗的每一秒跳躍並不伴隨着減少：獨特的元素存儲在其大小四倍的向量中（這樣可見，大小每增加一倍）。 元素必須被復制，這樣兩倍的內存被提交/使用，從而隱藏了哈希映射的較小使用。
• 消耗在兩者之間線性增加：隨着元素被添加到唯一的向量中，越來越多的內存頁面被提交：使用np.resize將不提交的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)將直接提交內存。