Как сократить размеры датафрейма в разы: учим Pandas сидеть на диете

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

import pandas as pd
import numpy as np
np.random.seed(0)

N = 100
M = 20

df = pd.DataFrame({'dt':pd.date_range(start='2021-01-01', periods=M), 'cat1':[str(it) for it in range(M)],
                  'cat2':[str(it) for it in range(M)]})\
      .join(pd.DataFrame(np.random.normal(0, 1, size=(M, N)), columns=[f'float_{it}' for it in range(N)]))\
      .join(pd.DataFrame(np.random.randint(1, 10, size=(M, N)), columns=[f'int_{it}' for it in range(N)]))

Для наглядности будем выводить диагностическую информацию:

Размер потребляемой датафреймом памяти с методами:

info - среди прочей информации выводит занимаемую память
memory_usage - поколоночное потребление памяти в байтах

Отличия в датафреймах, как compare на статистической информации из describe:

df1.describe().round(5).compare(df2.describe().round(5))

df.info()
df.memory_usage().sample(2, random_state=55)

to_numeric

В библиотеке Pandas имеется родной инструмент - функция to_numeric, которая для численных колонок помогает найти более экономичный подтип:

sh_df = df.copy()

float_cols = df.dtypes.astype(str).loc[lambda x: x.str.contains(pat='float', na=False, regex=True)].index
int_cols = df.dtypes.astype(str).loc[lambda x: x.str.contains(pat='int', na=False, regex=True)].index

for cols, col_type in zip([float_cols, int_cols], ['float', 'integer']):
    sh_df[cols] = sh_df[cols].apply(pd.to_numeric, downcast=col_type)

sh_df.info()
df.describe().round(5).compare(sh_df.describe().round(5))

Для агрегаций датафреймов разница только для минимума колонки float_37.

df_shrink

В fastai.tabular.core имеется функция df_shrink, которая располагает дополнительными возможностями, а именно, задавать колонки к которым преобразование типов не применять, преобразовывать ли object в category, int в uint. Под капотом всю работу по нахождению оптимальных типов делает df_shrink_dtypes, а df_shrink с полученным словарем только вызывает df.astype(...).

Входные параметры df_shrink_dtypes/df_shrink:

obj2cat - надо ли колонки object преобразовывать в тип category
int2uint - надо ли тип int преобразовывать в незнаковый uint (если нет отрицательных значений)
skip - список колонок, для которых преобразования типов не надо вычислять

from fastai.tabular.core import df_shrink

sh_df = df_shrink(df, obj2cat=True, skip=['cat1'], int2uint=True)

sh_df.info()

df.describe().round(5).compare(sh_df.describe().round(5))

manual df_shrink_dtypes

Отмечу, что оба указанных выше метода допускают преобразование чисел с плавающей запятой только до типа float32. Если хотите предусмотреть конвертацию и до float16, можно модифицировать код df_shrink_dtypes, добавив в словарь typemap np.float16:

def df_shrink_dtypes(df, skip=[], obj2cat=True, int2uint=False):
    "Return any possible smaller data types for DataFrame columns. Allows `object`->`category`, `int`->`uint`, and exclusion."

    # 1: Build column filter and typemap
    excl_types, skip = {'category','datetime64[ns]','bool'}, set(skip)

    typemap = {'int'   : [(np.dtype(x), np.iinfo(x).min, np.iinfo(x).max) for x in (np.int8, np.int16, np.int32, np.int64)],
               'uint'  : [(np.dtype(x), np.iinfo(x).min, np.iinfo(x).max) for x in (np.uint8, np.uint16, np.uint32, np.uint64)],
               'float' : [(np.dtype(x), np.finfo(x).min, np.finfo(x).max) for x in (np.float16, np.float32, np.float64, np.longdouble)]
              }
    if obj2cat: typemap['object'] = 'category'  # User wants to categorify dtype('Object'), which may not always save space
    else:       excl_types.add('object')

    new_dtypes = {}
    exclude = lambda dt: dt[1].name not in excl_types and dt[0] not in skip

    for c, old_t in filter(exclude, df.dtypes.items()):
        t = next((v for k,v in typemap.items() if old_t.name.startswith(k)), None)

        if isinstance(t, list): # Find the smallest type that fits
            if int2uint and t==typemap['int'] and df[c].min() >= 0: t=typemap['uint']
            new_t = next((r[0] for r in t if r[1]<=df[c].min() and r[2]>=df[c].max()), None)
            if new_t and new_t == old_t: new_t = None
        else: new_t = t if isinstance(t, str) else None

        if new_t: new_dtypes[c] = new_t
    return new_dtypes

t_d = df_shrink_dtypes(df, obj2cat=True, skip=['cat1'], int2uint=True)

sh_df = df.astype(t_d)

В целом функция осуществляет следующие шаги:

определяет словарь отображения типов, в котором хранит диапазоны значений подтипов
для каждой колонки, если она не попадает в список исключений (тип не в списке таких, например, 'datetime64[ns]','bool', 'category' и имя не в skip), находит минимальный подтип (такой, что значения минимума и максимума удовлетворяют его пределам).

Если у вас в проекте не предусмотрена работа с нейросетями и не установлены fastai и torch, можно просто скопировать код df_shrink_dtypes.

#pytorch #torch #fastai #types #pandas #dtypes #shrink #df_shrink_dtypes