groupby и transform: групповые признаки на уровне строк и фильтрация групп

В прошлом уроке groupby давал сводную таблицу: одна строка на группу, агрегат в колонке. Это удобно для отчётов, но неудобно, когда нужно добавить групповой показатель обратно к каждой строке исходной таблицы. Именно здесь agg заканчивается и начинается transform.

transform: групповой агрегат на уровне строк

Когда вызываешь groupby(col).agg('mean'), pandas собирает группы и возвращает сжатую таблицу — одна строка на уникальное значение ключа. Индекс результата состоит из значений группировочного столбца, и длина результата равна числу уникальных групп, а не числу строк исходного DataFrame.

transform работает иначе. Он тоже разбивает данные на группы и вычисляет агрегат внутри каждой из них, но затем «разворачивает» результат обратно по исходному индексу. Каждая строка получает значение агрегата своей группы. Длина возвращаемой Series всегда равна len(df), индекс совпадает с исходным.

Представь таблицу с тремя сотрудниками из отдела HR и двумя из IT. После groupby('dept')['salary'].agg('mean') получаешь две строки — по одной на отдел. После groupby('dept')['salary'].transform('mean') — пять строк: каждый сотрудник несёт среднюю зарплату своего отдела рядом со своей строкой.

Это принципиальная разница: результат transform выровнен по исходному индексу и немедленно готов к присваиванию в новый столбец без каких-либо слияний или ручного сопоставления. Именно поэтому transform — стандартный инструмент, когда нужен групповой признак на уровне строки, а не сводная статистика.

Добавление группового показателя как нового столбца

Процедура укладывается в два шага: вызвать transform и присвоить результат новому столбцу.

import pandas as pd

df = pd.DataFrame({
    'dept': ['HR', 'IT', 'HR', 'IT', 'HR'],
    'salary': [60000, 120000, 70000, 110000, 80000]
})

df['dept_mean'] = df.groupby('dept')['salary'].transform('mean')
df['dept_sum']  = df.groupby('dept')['salary'].transform('sum')

print(df)

Вывод:

  dept  salary   dept_mean  dept_sum
0   HR   60000   70000.0    210000
1   IT  120000  115000.0    230000
2   HR   70000   70000.0    210000
3   IT  110000  115000.0    230000
4   HR   80000   70000.0    210000

Что здесь происходит. Pandas разбивает строки по значению dept, считает среднее и сумму внутри каждой группы, а затем возвращает значение обратно в каждую строку по её исходной позиции. HR получает среднее (60000 + 70000 + 80000) / 3 = 70000, IT — (120000 + 110000) / 2 = 115000. Длина результата — ровно 5, как у исходного df.

В строку присваивания df['dept_mean'] = ... результат transform ложится без каких-либо дополнительных операций — индексы совпадают, pandas выравнивает значения автоматически.

Аргумент func в transform — строковый псевдоним агрегата: 'mean', 'sum', 'max', 'min', 'std' и другие стандартные агрегатные имена, которые принимает pandas. Передавать туда NumPy-функцию или лямбду тоже можно, но строковые имена работают быстрее, потому что pandas использует оптимизированные пути вычисления.

Фильтрация строк по условию на групповом агрегате

До этого момента transform использовался для добавления нового столбца с групповым показателем. Но та же идея — результат той же длины, что исходный DataFrame — открывает второй сценарий: фильтрацию строк по условию на уровне группы, а не на уровне строки.

Представь задачу: есть таблица продаж по магазинам, и нужно оставить только строки из магазинов, чья суммарная выручка превышает пороговое значение. Наивное решение — посчитать groupby('store')['revenue'].agg('sum'), получить сжатую таблицу, отфильтровать там нужные магазины, а потом вернуться к исходному DataFrame и выкинуть строки через merge или isin. Это работает, но требует нескольких шагов.

С transform паттерн короче:

import pandas as pd

df = pd.DataFrame({
    'store': ['A', 'B', 'A', 'C', 'B', 'C', 'A'],
    'revenue': [100, 50, 200, 30, 80, 20, 150]
})

group_sum = df.groupby('store')['revenue'].transform('sum')
mask = group_sum > 200
df_filtered = df.loc[mask]
print(df_filtered)

Вывод:

  store  revenue
0     A      100
2     A      200
6     A      150

Что происходит по шагам. transform('sum') для каждой строки возвращает суммарную выручку её магазина: строки магазина A получают 450, строки B — 130, строки C — 50. Результат — Series длиной 7 с теми же индексами, что у df. Сравнение > 200 превращает её в булевую маску той же длины. df.loc[mask] отбирает только строки, для которых маска истинна, — то есть все строки магазина A.

Ключевое здесь: маска живёт в пространстве строк, а решение принято на уровне группы. Магазины B и C не проходят порог целиком — ни одна их строка не попадает в результат, даже если отдельные транзакции большие. Это именно групповая фильтрация, а не строчная.

Паттерн одинаков для любого агрегата. Хочешь оставить группы, где среднее выше порога — меняешь 'sum' на 'mean' и подбираешь нужный threshold. Хочешь отбросить группы с единственной строкой — transform('count') > 1. Маска в любом случае строится за одно выражение и напрямую передаётся в df.loc без промежуточных таблиц.

Типичная ошибка: agg вместо transform при создании группового признака

Когда нужно добавить в DataFrame столбец с групповым средним, иногда пишут через agg — по аналогии с тем, как только что считали сводные показатели. Код выглядит похоже, но результат оказывается полностью сломан.

import pandas as pd

df = pd.DataFrame({
    'dept': ['HR', 'IT', 'HR', 'IT', 'HR'],
    'salary': [60000, 120000, 70000, 110000, 80000]
})

# Неверный вариант: agg вместо transform
df['dept_mean_wrong'] = df.groupby('dept')['salary'].agg('mean')
print(df)

Вывод:

  dept  salary  dept_mean_wrong
0   HR   60000              NaN
1   IT  120000              NaN
2   HR   70000              NaN
3   IT  110000              NaN
4   HR   80000              NaN

Почему все значения — NaN. agg('mean') возвращает Series с двумя элементами и индексом ['HR', 'IT']: это метки групп, а не позиции строк. Исходный df имеет числовой индекс [0, 1, 2, 3, 4]. Когда pandas присваивает Series в столбец, он выравнивает значения по меткам индекса. Метка 'HR' не совпадает ни с одной из меток 0..4, метка 'IT' — тоже. Совпадений нет вообще, поэтому все пять строк получают NaN. Pandas не сообщает об ошибке — он молча выполняет присваивание, и баг легко пропустить без проверки результата.

Исправление прямое: заменить agg на transform.

# Верный вариант
df['dept_mean'] = df.groupby('dept')['salary'].transform('mean')
print(df[['dept', 'salary', 'dept_mean']])

Вывод:

  dept  salary  dept_mean
0   HR   60000    70000.0
1   IT  120000   115000.0
2   HR   70000    70000.0
3   IT  110000   115000.0
4   HR   80000    70000.0

transform возвращает Series длиной 5 с числовым индексом [0, 1, 2, 3, 4] — точно как у df. Выравнивание при присваивании проходит корректно, каждая строка получает среднее своей группы.

Правило простое: если нужен групповой показатель как столбец в исходной таблице — transform. Если нужна сводная статистика отдельной таблицей — agg.