Как ускорить apply(). Часть 1

Привет! На связи Кристина Желтова, ментор курса «ML-инженер» 👋🏻

Представьте себе ситуацию: вы обрабатываете датасет с 1 млн строк. Пишете простой код:

df['new_col'] = df.apply(lambda x: x['col1'] + x['col2'], axis=1)

Запускаете. Проходит минута, вторая, пятая... Наконец после 15 минут вычисления завершились.

Потом приходит коллега, делает:

df['new_col'] = df['col1'].values + df['col2'].values

Запускает и получает результат за 0.15 секунды.

Что произошло? 🤯

Именно так выглядит разница между apply() и векторизацией, и это одно из самых узких мест в обработке данных.

Почему apply() такой медленный?

apply() работает примерно так:

Для каждой строки датасета (1 млн штук!):
➖ Взять одну строку;
➖ Передать в функцию lambda;
➖ Вызвать функцию (потратив время и «силы» на вызов — overhead);
➖ Получить результат;
➖ Вернуть в датасет.

Результат: 1 млн * overhead = ОЧЕНЬ медленно.

Векторизация же работает так:

➖ Взять весь столбец col1;
➖ Взять весь столбец col2;
➖ Сложить их поэлементно (в оптимизированном коде на C);
➖ Результат: одна операция, очень быстро.

Итого, в apply() 1 млн функций вызовов, а в векторизации 1 операция!

Проведём небольшой реальный тест:

import pandas as pd
import numpy as np
import time

# Создаём датасет
df = pd.DataFrame({
    'col1': np.random.rand(1_000_000),
    'col2': np.random.rand(1_000_000)
})

# Медленный apply()
start = time.time()
df['result_slow'] = df.apply(lambda x: x['col1'] + x['col2'], axis=1)
slow_time = time.time() - start
print(f"apply(): {slow_time:.2f} секунд")

# Быстрая векторизация
start = time.time()
df['result_fast'] = df['col1'].values + df['col2'].values
fast_time = time.time() - start
print(f"Векторизация: {fast_time:.4f} секунд")

# Разница
print(f"Ускорение: {slow_time / fast_time:.0f}x")

Как результат — ускорение в сотни раз и даже больше!

📊 Simulative