head в Polars group_by.agg(): согласованность столбцов

Как Polars сохраняет выравнивание при head(n) в group_by().agg(): первые n строк выбираются синхронно для нескольких столбцов. Альтернатива: group_by().head(n).

Когда вы агрегируете в Polars и применяете head(n) к нескольким столбцам внутри одного group_by().agg(), действительно ли значения берутся из тех же исходных строк? Это важный практический момент: если вы выбираете первые K элементов по группе сразу из нескольких полей, ожидается, что они будут согласованы.

Пример

В примере ниже данные группируются по ключу, а затем head(2) применяется сразу к трём столбцам. Цель — сохранить выравнивание: чтобы topic, vec и flag происходили из одних и тех же двух строк внутри каждой группы.

import polars as pl

# игрушечные данные
tbl = pl.DataFrame({
    "grp": ["A", "A", "A", "B", "B"],
    "topic": ["i1", "i2", "i3", "i4", "i5"],
    "vec": ["e1", "e2", "e3", "e4", "e5"],
    "flag": ["a1", "a2", "a3", "a4", "a5"],
})

res = (
    tbl.group_by("grp")
    .agg([
        pl.col("topic").head(2).alias("topics"),
        pl.col("vec").head(2).alias("vecs"),
        pl.col("flag").head(2).alias("flags"),
    ])
)

print(res)

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

Внутри каждой группы Polars сохраняет порядок строк. Поэтому при применении head к столбцу внутри agg() берутся первые n строк именно этой группы. Поскольку внутренний порядок в группе стабилен, выполнение той же операции для нескольких столбцов выбирает значения из одних и тех же исходных индексов строк.

Внутри каждой группы порядок строк всегда сохраняется, независимо от этого аргумента.

Следовательно, при таком использовании head риска «рассинхрона» нет. Для контраста: sample() может нарушить выравнивание, если выбирать несколько столбцов внутри agg().

Решение и более надёжная альтернатива

Подход с head внутри agg работает корректно и обеспечивает ожидаемое выравнивание. Если вам удобнее оперировать строками, а не списками, можно использовать и group_by(...).head(n), который возвращает первые n строк каждой группы. Тогда вопрос согласованности отпадает сам собой — вы напрямую получаете нужные строки по группам.

import polars as pl

# те же данные
tbl = pl.DataFrame({
    "grp": ["A", "A", "A", "B", "B"],
    "topic": ["i1", "i2", "i3", "i4", "i5"],
    "vec": ["e1", "e2", "e3", "e4", "e5"],
    "flag": ["a1", "a2", "a3", "a4", "a5"],
})

# возвращает первые 2 строки внутри каждой группы (а не из всего датафрейма)
rows = tbl.group_by("grp").head(2)
print(rows)

Легко предположить, что head(n) без agg работает по всему DataFrame, но в таком виде он ограничивается каждой группой и берёт первые n строк именно в ней.

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

В задачах feature engineering, ранжирования и извлечения top-K часто требуется брать первые n элементов из нескольких столбцов синхронно. Понимание того, что head сохраняет выравнивание внутри групп, позволяет уверенно агрегировать без последующих исправлений. А если нужен результат в виде строк, а не списков, group_by(...).head(n) делает вывод простым и предсказуемым.

Итоги

Применение head(n) к нескольким столбцам внутри group_by().agg() сохраняет согласованность, поскольку порядок в группе сохраняется, а head выбирает первые строки каждой группы. Для ещё более наглядного результата group_by(...).head(n) возвращает первые n строк в каждой группе именно строками, а не списками. Помните, что не все операции обладают этим свойством: например, выборка нескольких столбцов с помощью sample() внутри agg() может их рассинхронизировать.