Ранжирование категорий по месяцам в Polars: rank over и pivot

Как в Polars посчитать частоты категорий по месяцам, присвоить ранги с rank over и развернуть результат через pivot. Идиоматичный пример кода и нюансы ничьих.

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

Пример набора данных и базовый подход

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

import polars as pl

records = pl.DataFrame(
    {
        "rec_id": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17],
        "period": [1, 2, 1, 2, 1, 2, 2, 2, 1, 1, 1, 2, 2, 2, 1, 1, 2],
        "label": [
            "C",
            "B",
            "A",
            "C",
            "B",
            "A",
            "C",
            "C",
            "A",
            "B",
            "A",
            "A",
            "C",
            "C",
            "A",
            "B",
            "B",
        ],
    }
)

(
    records.group_by(pl.col("period"), pl.col("label"))
    .agg(pl.col("rec_id").len().alias("freq"))
    .sort(by=["period", "freq"], descending=True)
    .with_columns(pl.lit(1).alias("flag"))
    .with_columns(pl.col("flag").cum_sum().over(["period"]).alias("rank_idx"))
    .pivot(index="rank_idx", values="label", on="period", sort_columns=True)
)

Это решает задачу: сначала мы считаем вхождения, сортируем по частоте внутри каждого месяца, имитируем счётчик по месяцу с помощью накопительной суммы, а затем делаем сводную развёртку. Работает, но не слишком прямо отражает идею «ранга внутри группы».

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

Суть задачи — ранжировать количества категорий по месяцам. Хитрость с накопительной суммой даёт индекс, похожий на ранг, но в Polars уже есть встроенная операция ранжирования. Недостающее звено — применить ранг внутри каждого месяца, и именно это делает оконное выражение с over. Есть и тонкость с обработкой равенств: метод "ordinal" разрешает ничьи произвольно. На практике при равных счетчиках в столбце конкретного месяца можно увидеть, например, порядок CBA вместо CAB.

Более простой вариант с rank и over

Ниже — более лаконичный и идиоматичный конвейер, который считает, ранжирует по месяцам и разворачивает результат в одной цепочке.

(
    records.group_by("period", "label")
    .len()
    .with_columns(
        pl.col("len").rank("ordinal", descending=True).over("period").alias("rank_idx")
    )
    .sort("rank_idx")
    .pivot("period", index="rank_idx", values="label", sort_columns=True)
)

Практическая деталь: результирующий rank_idx имеет тип u32, а не i64; если для последующей логики важен тип целого, можно выполнить приведение в конце.

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

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

Выводы

Для ранжирования внутри групп в Polars предпочтительнее rank с over, а не построение рангов косвенно через накопительные суммы. Посчитайте по (месяц, категория), ранжируйте счётчики по убыванию внутри месяца, отсортируйте по этому рангу и сделайте сводную развёртку, чтобы выровнять категории по рангам между месяцами. Помните, что ordinal-ранжирование разбирается с ничьими произвольно, и при необходимости приводите тип ранга к нужному в итоговой схеме.