Фильтрация сгруппированных данных в pandas по индексу без второго groupby

Разбираем, как фильтровать сгруппированные данные в pandas: индексация вместо DataFrameGroupBy.filter, без повторной группировки и с подсчетом выживших групп.

Фильтрация сгруппированных данных в pandas по подмножеству исходных значений индекса кажется простой задачей, пока не становится нужно одновременно сохранить исходную группировку и получить отфильтрованный вид, а заодно сравнить, сколько групп переживут фильтр. Загвоздка в том, что DataFrameGroupBy.filter удаляет целые группы по агрегирующему предикату, тогда как здесь цель — оставить только конкретные строки внутри каждой группы по индексу и уже затем рассуждать о получившихся группах.

Постановка задачи

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

import pandas as pd

frame = pd.DataFrame(
    data={
        "g0": ["foo", "foo", "bar", "bar"],
        "g1": ["baz", "baz", "baz", "qux"],
        "data": [0.1, 0.3, 0.4, 0.2],
    },
    index=["a", "b", "c", "d"],
)

bunches = frame.groupby(by=["g0", "g1"], sort=False)
keep_idx = ["a", "b", "d"]

Почему DataFrameGroupBy.filter здесь не подходит

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

Решение 1: Сначала отфильтровать, затем сгруппировать

Если нужен только отфильтрованный сгруппированный вид, самый прямой путь — сузить DataFrame по индексу и уже результат сгруппировать. Так мы получаем корректное количество групп после фильтрации, поскольку группы без оставшихся строк просто не попадут в объект группировки.

filtered_bunches = frame[frame.index.isin(keep_idx)].groupby(by=["g0", "g1"], sort=False)

Подход простой и даёт аккуратный объект группировки, отражающий только сохранённые индексы.

Решение 2: Сохранить исходные группы и получить отфильтрованный вид без повторной группировки

Когда важно сохранить исходную группировку и избежать второго groupby, можно итерироваться по уже существующим группам и для каждой вырезать нужные строки по индексу. Так исходная группировка остаётся нетронутой, а на выходе — DataFrame каждой группы, отфильтрованный по индексу.

filtered_parts = [part[part.index.isin(keep_idx)] for _, part in bunches]

Если нужно именно количество групп, оставшихся непустыми после фильтрации, просто игнорируйте пустые срезы. По числу групп это будет эквивалентно повторной группировке уже отфильтрованного фрейма.

nonempty_parts = [
    part[part.index.isin(keep_idx)]
    for _, part in bunches
    if not part[part.index.isin(keep_idx)].empty
]

original_group_count = len(bunches)
filtered_group_count = len(nonempty_parts)

Как считать количество групп после фильтрации

Возьмём пример выше. До фильтрации есть три группы. При keep_idx = ["a", "b", "d"] группа, содержащая только индекс "c", исчезает после фильтрации — остаётся две группы. При keep_idx = ["a", "c", "d"] общая для "a" и "b" группа сохраняется, потому что присутствует "a", и группа для "c" тоже остаётся, так что всего по-прежнему три.

Зачем это нужно

Выборка по индексу внутри сгруппированных данных полезна в конвейерах, где требуется частичное сохранение групп, а также при сравнении структуры данных до и после фильтра. Булева индексация по исходному индексу точно передаёт задуманное и избегает семантики групповой фильтрации, при которой иначе удаляются целые группы.

Выводы

Если нужен только отфильтрованный сгруппированный вид, сначала фильтруйте по индексу и группируйте один раз. Если важно сохранить исходные группы и одновременно получить отфильтрованный ракурс без второго groupby, переиспользуйте исходный GroupBy и нарежьте каждый фрагмент. Когда требуется сравнить, сколько групп пережило фильтр, считайте только непустые срезы — так вы получите тот же результат, что и при повторной группировке отфильтрованного DataFrame.