Любопытное потребление памяти pandas.unique()

Профилируя потребление памяти в моем алгоритме, я был удивлен, что иногда для небольших входов требуется больше памяти.

Все это сводится к следующему использованию 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.

Сканирование различных размеров ввода дает следующий график:

enter image description here

Из любопытства и самообразования: как можно объяснить противоречивое поведение (т.е. Больше памяти для меньшего размера ввода) вокруг 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

Ответы

Ответ 1

Посмотрим...

pandas.unique говорит, что это "уникальный хэш-стол".

Он вызывает эту функцию для получения правильной реализации хеш-таблицы для ваших данных, а именно htable.Int64HashTable.

Хэш-таблица инициализируется size_hint= длина вашего вектора значений. Это означает, что kh_resize_DTYPE(table, size_hint).

Эти функции определены (templated) здесь в 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,
    ...,
]

Надеюсь, это проливает свет на вещи :)