Собесов

Сценарий: apply против векторизации — реальный бенчмарк

PythonВекторизация и performanceСредняяMiddle

Условие

В df 5 млн строк. Нужно посчитать колонку bonus:

bonus = amount * 0.1 если segment == 'premium' иначе amount * 0.05

Сравни 4 варианта по скорости: apply, iterrows, np.where, Series.where.

Решение

Бенчмарк

import numpy as np
import pandas as pd
 
df = pd.DataFrame({
    'amount': np.random.rand(5_000_000) * 1000,
    'segment': np.random.choice(['premium', 'basic'], 5_000_000),
})
 
# 1. iterrows — катастрофа, ~ минуты
def f_iter():
    out = []
    for _, r in df.iterrows():
        out.append(r['amount'] * 0.1 if r['segment'] == 'premium' else r['amount'] * 0.05)
    return out
 
# 2. apply по строкам — медленно, ~ десятки секунд
def f_apply():
    return df.apply(lambda r: r['amount']*0.1 if r['segment']=='premium' else r['amount']*0.05, axis=1)
 
# 3. np.where — быстро, ~ 50 мс
def f_np():
    return np.where(df['segment'] == 'premium', df['amount']*0.1, df['amount']*0.05)
 
# 4. Series.where — почти как np.where
def f_swhere():
    return (df['amount']*0.05).where(df['segment'] != 'premium', df['amount']*0.1)

Порядок ускорения

iterrows  >> apply(axis=1)  >>> np.where ≈ Series.where ≈ маска

Векторные операции в 100–1000 раз быстрее apply(axis=1).

Альтернативы для сложных условий

# Несколько условий: np.select
conds = [df['segment'] == 'premium', df['segment'] == 'gold']
choices = [df['amount']*0.10, df['amount']*0.15]
df['bonus'] = np.select(conds, choices, default=df['amount']*0.05)

Подводные камни

  1. apply(axis=1) — Python-уровень цикла, не векторизация. Даже с swifter/pandarallel обычно проигрывает np.where.
  2. iterrows создаёт новые Series на каждую строку — медленно и теряет dtype.
  3. df['col'].apply(simple_func) (по столбцу) часто быстрее, чем по строкам, но всё равно проигрывает векторам.
  4. Для строк (str) — df['s'].str.upper() уже векторизована, не нужен apply.
  5. Когда нужна тяжёлая нелинейная логика — numba (@numba.njit) или numpy-векторы.

Эталонный ответ

np.where / np.select / маска быстрее apply(axis=1) в 100–1000 раз. iterrows не использовать. Векторизуйте сначала, оптимизируйте потом.

Хочешь увидеть разбор?

Зарегистрируйся бесплатно — откроется развёрнутое решение этой задачи и ещё 4 на выбор.

Зарегистрироваться и увидеть разбор
Уже есть аккаунт? Войти