August 19, 2022

Как метод apply из Pandas тормозит вычисления и чем его заменять

При работе с данными свыше 1G вы обязательно ощутите проблему медленной работы метода apply. Поэтому изучайте специализированные возможности библиотеки Pandas и проводите операции на уровне целых колонок, apply же в силу универсальности будет выполнять аналогичные шаги намного медленнее. Это правило проявляется не только в случае использования метода для построчной обработки датафрейма, но и поколоночных операций. Рассмотрим кейс преобразования типов колонок с apply и без него. Сначала сгенерируем датафрейм:

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

df = pd.DataFrame(np.random.randint(500, size=(10000000, 5)))
df.head()

По умолчанию типы колонок установлены в int64:

df.dtypes

Пусть мы хотим сэкономить память и преобразовать типы к оптимальному формату:

num_cols = df.select_dtypes(include=np.number).columns
%%time
df1 = df.copy()
df1[num_cols] = df1[num_cols].apply(pd.to_numeric, downcast='integer')
df1.dtypes

А теперь сделаем обработку колонок в цикле без apply:

%%time
df2 = df.copy()
for col in num_cols:
    df2[col] = pd.to_numeric(df2[col], downcast='integer')

Получили примерно 20% выигрыш по времени. Но данный кейс еще "цветочки", так как apply применялся к 5 колонкам. В большинстве случаев его используют для построчной обработки датафрейма (а строк обычно кратно больше, чем столбцов). Рассмотрим такой подход для подсчета суммы значений по строкам:

%%time
def get_sum(x):
    return sum(x)

df1=df.copy()
df1['sum'] = df1.apply(lambda x: get_sum(x), axis=1)

Цикл справляется с той же задачей более чем в 80 раз быстрее:

%%time
df2 = df.copy()

df2['sum'] = 0
for i in range(5):
    df2['sum'] = df2['sum'] + df2[i] 

Однако еще лучше для проведения типичных операций использовать готовые функции Pandas:

%%time
df3 = df.copy()
df3['sum'] = df3.sum(axis=1)

В данном случае удалось получить еще более впечатляющую разницу - более чем в 100 раз.

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