Почему groupby по Series в pandas зависит от индексов

Разбираем, как pandas groupby использует Series как ключ: перед группировкой значения выравниваются по индексам. Пояснения, примеры и практики работы.

Группировка pandas Series по другой Series кажется простой — пока в дело не вступает выравнивание индексов. Если передать Series как группировщик и незаметно изменить её индекс, результат может показаться лишённым смысла. Но это ожидаемое поведение: перед использованием значений Series в качестве ключей групп pandas сначала выравнивает их по меткам индекса. Осознав это правило, вы поймёте всё, что видите на экране.

Воспроизводим проблему

Ниже мы группируем Series самой по себе, а затем по её же версии с переустановленным индексом. Во втором случае кажется, что «всё ломается».

import numpy as np
import pandas as pd
# базовая серия
data_s = pd.Series([10, 10, 20, 30, 30, 30], index=np.arange(6) + 2)
print(data_s)
# 2    10
# 3    10
# 4    20
# 5    30
# 6    30
# 7    30
# dtype: int64
# группировка по той же серии (ожидаемо)
bins_1 = data_s.groupby(data_s)
for grp_key, grp_vals in bins_1:
    print(f"Group: {grp_key}")
    print(grp_vals)
# Группа: 10
# 2    10
# 3    10
# dtype: int64
# Группа: 20
# 4    20
# dtype: int64
# Группа: 30
# 5    30
# 6    30
# 7    30
# dtype: int64
# те же значения как группировщик, но с переустановленным индексом (неожиданный результат)
key_s = data_s.reset_index(drop=True)
bins_2 = data_s.groupby(key_s)
for grp_key, grp_vals in bins_2:
    print(f"Group: {grp_key}")
    print(grp_vals)
# Группа: 20.0
# 2    10
# dtype: int64
# Группа: 30.0
# 3    10
# 4    20
# 5    30
# dtype: int64

Что происходит на самом деле

Поведение определяется тем, как pandas трактует аргумент “by” в groupby. В документации сказано:

by : mapping, function, label, pd.Grouper or list of such

Используется для определения групп в groupby. Если by — это функция, она вызывается для каждого значения индекса объекта. Если передан dict или Series, для определения групп будут использованы ИХ ЗНАЧЕНИЯ (значения Series сначала выравниваются; см. метод .align()).

Ключевое слово — выравниваются. Когда вы передаёте Series как группировщик, pandas сначала выравнивает её с объектом для группировки по меткам индекса, а затем использует выровненные значения как ключи групп. То есть группировка основана не на позициях, а на метках. Если метки не совпадают, на не пересекающихся метках вы получите NaN, и эти строки тихо исключатся из группировки, потому что NaN не является допустимым ключом группы.

В примере у исходной Series метки индекса 2, 3, 4, 5, 6, 7. У группировщика с переустановленным индексом — 0, 1, 2, 3, 4, 5. При выравнивании по меткам ключи получают только метки 2–5, а для 6–7 выходит NaN. Поэтому группы оказываются лишь 20.0 и 30.0, и ничего нет для меток 6 и 7.

Посмотрим выравнивание наглядно

Если явно показать выравнивание, картина становится очевидной. Объедините исходную Series с Series после сброса индекса по индексу, а затем сгруппируйте по выровненному столбцу.

# визуализируем выравнивание через объединение по индексу
aligned = (
    data_s.rename("col_src").to_frame()
    .join(key_s.rename("col_key"))
)
print(aligned)
#    col_src  col_key
# 2       10     20.0
# 3       10     30.0
# 4       20     30.0
# 5       30     30.0
# 6       30      NaN
# 7       30      NaN
# группировка по выровненному столбцу ключей
for grp_key, grp_df in aligned.groupby("col_key"):
    print(f"Group: {grp_key}")
    print(grp_df)
# Группа: 20.0
#    col_src  col_key
# 2       10     20.0
# Группа: 30.0
#    col_src  col_key
# 3       10     30.0
# 4       20     30.0
# 5       30     30.0

Как об этом думать

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

Почему это важно

Выравнивание по индексу — одна из самых мощных и тонких возможностей pandas. Легко принять результат за баг, когда библиотека делает ровно то, что обещает. Группировка с несовпадающими индексами может тихо отбросить данные или разнести строки по неожиданным корзинам. Цена — время на отладку «случайного» поведения, которое на самом деле детерминировано.

Выводы

При передаче Series в groupby значения используются только после выравнивания по индексу. Если группы выглядят странно, проверьте индексы с обеих сторон или воссоздайте шаг выравнивания через join, чтобы увидеть, по каким ключам pandas фактически группирует. Относитесь к меткам индекса как к объектам первого класса — и неприятных сюрпризов станет меньше.