Pandas apply и векторизация: как писать быстрее и понятнее

Коротко

apply(axis=1) — хороший инструмент для сложного правила, которое трудно выразить иначе. Для арифметики, строковых методов, условий и дат сначала ищи операцию над всей Series. map и replace удобны для словаря значений, np.select — для нескольких условий.

• Начни с векторной операции и переходи к apply только при необходимости.
• Не используй axis=1, если условие зависит от одной колонки.
• Замеряй время на репрезентативном объёме данных.
• После переписывания сравни результат старой и новой версии на выборке.

Простая арифметика не требует apply

Построчный код выглядит привычно, но pandas уже умеет выполнять операцию над всей колонкой. Векторная версия лучше показывает формулу и не создаёт Python-вызов на каждую строку.

# Медленнее и сложнее читать на больших таблицах
orders['net_revenue'] = orders.apply(
    lambda row: row['revenue'] - row['discount'], axis=1
)

# Векторная версия
orders['net_revenue'] = orders['revenue'] - orders['discount']

Условия: np.select вместо вложенных if

Когда сегмент зависит от нескольких условий, np.select делает порядок правил явным. Условия проверяются сверху вниз, поэтому сначала ставь более специфичное правило. Значение default обязательно: иначе неизвестный случай легко потеряется.

import numpy as np

conditions = [
    orders['revenue'].ge(5000),
    orders['revenue'].ge(1500),
    orders['revenue'].ge(0),
]
choices = ['high', 'medium', 'low']
orders['order_segment'] = np.select(conditions, choices, default='invalid')

map и str работают над колонкой

Словарь каналов, нормализация регистра и поиск части строки не требуют построчной функции. Такие операции проще покрыть тестами и легче заметить, если неизвестное значение превратилось в пропуск.

channel_names = {'paid_social': 'paid', 'paid_search': 'paid'}
orders['channel_group'] = orders['channel'].map(channel_names)
orders['channel_group'] = orders['channel_group'].fillna('other')

orders['campaign'] = (
    orders['campaign'].astype('string').str.strip().str.lower()
)

Когда apply оправдан

Если правило использует несколько значений строки, вызывает сложную доменную функцию или строит небольшой текстовый объект, apply может быть самым понятным решением. Не нужно переписывать всё ради скорости. Сначала измерь узкое место и зафиксируй ожидаемый результат.

Как сравнить варианты на реальных данных

Разница между apply и векторизацией особенно заметна на сотнях тысяч строк, но измерять нужно на объёме, похожем на рабочий. Маленький DataFrame может быть слишком быстрым для обоих вариантов, а синтетические данные без пропусков не покажут цену обработки плохих строк. Зафиксируй вход, прогоняй каждую версию несколько раз и сравни не только время, но и результат.

Сначала напиши эталонную реализацию, которую легко прочитать. Затем сделай векторную версию и сравни колонки через pd.testing.assert_series_equal или assert_frame_equal. Если значения различаются из-за округления, опиши допустимую погрешность явно. Нельзя считать переписывание успешным, если оно просто поменяло NaN на ноль или по-другому обработало отрицательную сумму.

Профилируй весь pipeline, а не одну строку в вакууме. Иногда apply занимает 10% времени, а 90% уходит на чтение CSV или сортировку. В таком случае переписывание функции не изменит пользовательский опыт, зато усложнит код.

import pandas as pd
from time import perf_counter

started = perf_counter()
old = orders.apply(lambda row: row['revenue'] - row['discount'], axis=1)
old_seconds = perf_counter() - started

started = perf_counter()
new = orders['revenue'] - orders['discount']
new_seconds = perf_counter() - started
pd.testing.assert_series_equal(old, new, check_names=False)
print({'apply': old_seconds, 'vectorized': new_seconds})

axis=1 скрывает контракт строки

apply(axis=1) передаёт в функцию Series, где значения могут иметь смешанные типы. Внутри лямбды легко получить строку вместо числа, а ошибка проявится только на конкретной строке. Кроме того, функция начинает зависеть от названий колонок и их наличия, но этот контракт часто нигде не записан. Если apply остаётся, вынеси функцию с понятной сигнатурой и проверь обязательные поля до запуска.

Ещё одна ловушка — использование состояния снаружи: словаря, текущей даты или переменной, которая была изменена в предыдущей ячейке. Один и тот же DataFrame тогда даёт разные результаты в зависимости от порядка запуска ноутбука. Передавай справочники явно или используй map, merge и чистую функцию, чтобы результат зависел только от входных данных.

Векторизация не отменяет необходимость думать о типах. Выражение с числом и pd.NA может привести к nullable-типу, а сравнение строк с числом — вернуть неожиданный результат или ошибку. Проверка dtypes, доли пропусков и диапазона значений должна идти до сложных условий.

Условия должны быть взаимоисключающими

В np.select условия проверяются в заданном порядке. Если правило revenue >= 0 стоит перед revenue >= 5000, дорогой заказ попадёт в low и ошибка будет выглядеть правдоподобно. Пиши условия от специфичного к общему и добавляй тесты на границы: ровно 1500, ровно 5000, отрицательное значение и пропуск.

Если условия пересекаются намеренно, объясни приоритет в комментарии или в тексте статьи. Иногда полезнее сначала создать отдельные boolean-колонки: is_large, is_returned, is_new_user, а затем собрать итоговый сегмент. Это занимает больше строк, зато позволяет проверить каждую часть правила и построить диагностический срез.

Для бизнес-сегмента обязательно оставь значение unknown или invalid. Default low удобен для демонстрации, но в отчёте он смешивает реальные низкие значения с отрицательными и нераспознанными. Отдельная категория быстрее показывает, что источник прислал новый код.

conditions = [
    orders['revenue'].ge(5000),
    orders['revenue'].ge(1500),
    orders['revenue'].ge(0),
]
choices = ['high', 'medium', 'low']
orders['segment'] = np.select(conditions, choices, default='invalid')

assert orders.loc[orders['revenue'].eq(5000), 'segment'].eq('high').all()
assert orders.loc[orders['revenue'].eq(-1), 'segment'].eq('invalid').all()

Строки и даты: accessor лучше универсальной функции

У pandas есть специальные accessor-ы str и dt, потому что они знают о типе Series и работают над колонкой. Через str.strip, str.lower и str.contains нормализация становится последовательной. Но перед этим нужно привести пропуски и убедиться, что колонка действительно строковая. Иначе числа и None начнут обрабатываться не так, как ожидает лямбда.

С датами та же идея: pd.to_datetime выполняется один раз, после чего используются .dt.date, .dt.month, .dt.dayofweek или .dt.to_period. Не вызывай to_datetime внутри apply для каждой строки. Это дороже и скрывает настройку часового пояса и ошибок парсинга.

Если строковое правило сложное, сначала приведи данные к каноническому виду, затем примени простые маски. Две-три именованные маски легче проверить в таблице, чем одна длинная лямбда с пятью тернарными операторами.

Практическое решение: понятный код против микросекунд

В аналитике оптимизация имеет смысл, когда она меняет время получения ответа или стоимость запуска. Для отчёта на 5 000 строк читаемый apply может быть разумным выбором, если правило действительно построчное. Для лога на 20 миллионов событий даже простая арифметика должна быть векторной или вынесенной в SQL. Ориентиром служит не культ скорости, а масштаб задачи и цена ожидания.

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

• Сначала определить зерно и типы входных колонок.
• Выбрать векторную операцию, если правило относится к одной колонке.
• Проверить границы и unknown на маленьком контрольном наборе.
• Замерить реальный объём и сравнить результат со старой версией.
• Оставить apply, если он честнее выражает сложное правило.