Не пытайтесь измерить использование памяти в Pandas

Представьте: у вас есть файл с данными, которые вы хотите обработать в Pandas. Хочется быть уверенным, что память не закончится. Как оценить использование памяти с учетом размера файла?

Все эти оценки могут как занижать, так и завышать использование памяти. На самом деле оценивать использование памяти просто не стоит. А если конкретнее, в этой статье я:

• покажу широкий диапазон использования памяти ещё до обработки, только во время загрузки данных;
• расскажу о других подходах — измерении и передаче файла по частям.

"Обработка" — слишком широкое понятие

Есть ди хорошая эвристика, которая исходит из размера файла набора данных? Трудно сказать, ведь обработка данных может означать многое. Вот две крайности, когда вы:

• вычисляете максимальное значение столбца, и в этом случае память потрачена только на то, чтобы загрузить данные;
• выполняете перекрестное соединение между двумя столбцами длиной M и N, — объем памяти будет M×N; удачи вам, если это большие числа.

Итак, начнем с того, что оценка зависит от конкретной ситуации. Нужно понять код и детали реализации Pandas. Обобщенные умножения вряд ли будут полезны.

Но полезна была бы некоторая основная информация: объем памяти для данного DataFrame или Series. Для примеров выше это ключевой фактор при оценке использования памяти.

Какова же хорошая эвристика оценки использования памяти для загрузки набора данных по размеру файла? Ответ: она зависит от ситуации.

Использование памяти может оказаться намного меньше размера файла

Иногда использование памяти будет намного размера входного файла. Давайте сгенерируем CSV-файл в 1 000 000 строк с тремя числовыми столбцами; первый столбец — в диапазоне от 0 до 100, второй — от 0 до 10 000, третий — от 0 до 1 000 000.

import csv
from random import randint

with open("3columns.csv", "w") as csvfile:
    writer = csv.writer(csvfile)
    writer.writerow(["A", "B", "C"])
    for _ in range(1_000_000):
        writer.writerow(
            [
                randint(0, 100),
                randint(0, 10_000),
                randint(0, 1_000_000_000),
            ]
        )

Файл весит 18 МБ. Теперь загрузим данные в dtype соответствующего размера и измерим использование памяти:

import pandas as pd
import numpy as np

df = pd.read_csv(
    "3columns.csv",
    dtype={"A": np.uint8, "B": np.uint16, "C": np.uint32},
)
print(
    "{:.1f}".format(df.memory_usage(deep=True).sum() / (1024 * 1024))
)

Итоговое использование памяти при запуске кода выше — 6.7MB, а это 0,4 от размера исходного файла.

Может Parquest поможет?

При представлении чисел в виде удобочитаемого текста (CSV) данные могут занимать больше байтов, чем представление в памяти. Например, \~90% случайных значений от 0 до 1 000 000 будут состоять из 6 цифр, поэтому в CSV потребуется 6 байтов. Однако в памяти для каждого np.int32 требуется всего 4 байта. Кроме того, в CSV есть запятые и новые строки.

Так что, возможно, другой, лучший формат файла может дать нам способ оценить использование памяти с учетом размера файла? Файл Parquet, например, хранит данные так, чтобы они точнее соответствовали представлению в памяти.

Эквивалент Parquet для примера выше составляет 7 МБ. По сути, столько же, сколько использование памяти. Так является ли размер файла Parquet хорошим индикатором в смысле оценки использования памяти?

Использование памяти может быть намного больше размера файла

Размер файла Parquest может вводить в заблуждение

Загрузим еще один файл Parquet из предыдущей статьи с размером в 20 МБ. Использование памяти при загрузке составляет… 300 МБ.

Одна из причин в том, что Parquet может сжимать данные на диске. В нашем исходном примере использовались случайные данные, а это означает, что сжатие помогает не сильно. Данные этого конкретного файла гораздо более структурированы, а значит, сжимаются гораздо лучше. Делать обобщения о том, насколько хорошо будут сжиматься данные, трудно; но, конено, можно представить себе случаи, когда данные сжимаются сильнее.

Другая причина, по которой данные могут быть меньше размера в памяти — строковое представление.

Строки — это весело

Может ли несжатый файл Parquet помочь нам оценить использование памяти? Для числовых данных разумно ожидать соотношения 1 к 1. Для строк… не обязательно.

Давайте создадим несжатый Parquet со строковым столбцом; строки в нём вроде таких — "A B A D E":

from itertools import product
import pandas as pd

alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"

data = {
    "strings": [
        " ".join(values) for values in product(
            *([alphabet] * 5)
        )
    ]
}
df = pd.DataFrame(data)
print(
    "{:.1f}".format(df.memory_usage(deep=True).sum() / (1024 * 1024))
)
df.to_parquet("strings.parquet", compression=None)

Размер файла составляет 148 МБ, сжатия нет. Использование памяти — 748 МБ, в 5 раз больше. Разница в том, что Pandas и Parquet по-разному представляют строки; Parquest использует UTF-8, как и Arrow.

Использование строковой памяти Python также зависит от того, происходит ли [интернирование] (https://docs.python.org/3/library/sys.html#sys.intern) (деталь реализации CPython, которая может меняться в разных версиях Python), а также в том, записаны ли ваши строки в ASCII, содержат ли китайские символы или эмоджи.

>>> import sys
>>> sys.getsizeof("0123456789")
59
>>> sys.getsizeof("012345678惑")
94
>>> sys.getsizeof("012345678