Ускорить groupby.apply

Question

Пайплайн считает метрики по пользователям через groupby("user").apply(customfn) и занимает 40 минут. Как будете ускорять?

Техпрогон · Accepted Answer

Сначала измеряю: сколько групп, сколько строк, что делает функция. Почти всегда тормозит не pandas, а Python-вызов на каждую из миллионов групп. Профилирую: где время, сколько групп, что делает функция. Разбиваю функцию на встроенные агрегаты и transform. Условия через маски и where до группировки, а не внутри функции. Сортировка плюс shift, cumsum, diff в группе вместо цикла по строкам. Сократить данные: нужные колонки, типы category, фильтр до группировки. Если не хватает: Polars или DuckDB, вынос в SQL. Сначала измеряю: сколько групп, сколько строк, что делает функция. Почти всегда тормозит не pandas, а Python-вызов на каждую из миллионов групп. Дальше разбираю customfn на части: суммы, средние и счётчики переписываю на встроенные агрегаты через именованную агрегацию, условные метрики считаю так: сначала колонка-маска на всём DataFrame (df["paid"] = df["amount"].where(df["status"] == "ok", 0)), потом обычный sum по группам. Последовательную логику внутри группы заменяю сортировкой и групповыми shift, diff, cumsum. Уменьшаю данные: только нужные колонки, строки category, фильтр до группировки. Новую версию сверяю со старой на выборке через assertframeequal. Если и этого мало, переношу расчёт в DuckDB или Polars, они параллелят группировку по ядрам.

Ускорить groupby.apply

Вопрос

Попробуйте сами

Коротко

Что должно прозвучать в ответе

Что могут спросить следом

Разбор

Термины по теме

Ещё по теме «groupby и агрегаты»