Polars中按分组内计数对类别排名并透视的更优实现方案问询

我现在有一个包含月份、类别和ID的Polars DataFrame，想要按每个月内的类别出现频率进行排名，然后将结果透视，展示每个排名位置在不同月份对应的类别。

我的DataFrame数据： #

import polars as pl
df = pl.DataFrame(
    {
        "ID": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17],
        "MONTH": [1, 2, 1, 2, 1, 2, 2, 2, 1, 1, 1, 2, 2, 2, 1, 1, 2],
        "CATEGORY": [
            "C", "B", "A", "C", "B", "A", "C", "C", "A", "B", "A", "A", "C", "C", "A", "B", "B",
        ],
    }
)

期望输出： #

shape: (3, 3)
┌───────┬─────┬─────┐
│ PLACE ┆ 1   ┆ 2   │
│ ---   ┆ --- ┆ --- │
│ i64   ┆ str ┆ str │
╞═══════╪═════╪═════╡
│ 1     ┆ A   ┆ C   │
│ 2     ┆ B   ┆ A   │
│ 3     ┆ C   ┆ B   │
└───────┴─────┴─────┘

目前可行的解决方案： #

(
    df.group_by(pl.col("MONTH"), pl.col("CATEGORY"))
    .agg(pl.col("ID").len().alias("COUNT"))
    .sort(by=["MONTH", "COUNT"], descending=True)
    .with_columns(pl.lit(1).alias("aux"))
    .with_columns(pl.col("aux").cum_sum().over(["MONTH"]).alias("PLACE"))
    .pivot(index="PLACE", values="CATEGORY", on="MONTH", sort_columns=True)
)

这个方案能得到正确结果，但我感觉在Polars里应该有更简洁直接的方式来实现这个排名+透视的操作，有没有更优雅的方案呢？

更优实现方案 #

当然有！我们可以利用Polars原生的rank函数直接在分组内计算排名，省去手动创建辅助列求和的步骤，代码会更简洁直观：

(
    df.group_by(["MONTH", "CATEGORY"])
    .agg(count=pl.len())
    .with_columns(
        PLACE=pl.col("count").rank(descending=True, method="ordinal").over("MONTH")
    )
    .pivot(index="PLACE", values="CATEGORY", on="MONTH", sort_columns=True)
    .sort("PLACE")
)

代码细节解释：

1. 分组计数简化：用pl.len()替代原方案的pl.col("ID").len()，不需要依赖具体的ID列，代码通用性更强。
2. 原生排名函数：直接用rank函数在MONTH分组内计算排名，descending=True表示按计数从高到低排序，method="ordinal"保证相同计数时生成连续的排名（本案例中每个月的类别计数无重复，也可以省略该参数，默认就是ordinal）。
3. 透视排序：和原方案逻辑一致，最后按PLACE排序确保结果顺序符合预期。

如果想要更极致的简洁，还可以省略计数列的显式命名，直接使用分组后默认的len列名：

(
    df.group_by(["MONTH", "CATEGORY"])
    .agg(pl.len())
    .with_columns(
        PLACE=pl.col("len").rank(descending=True).over("MONTH")
    )
    .pivot(index="PLACE", values="CATEGORY", on="MONTH", sort_columns=True)
    .sort("PLACE")
)

对比原方案的优势：

• 去掉了手动创建aux辅助列再累加的冗余步骤，逻辑更贴合业务需求，可读性大幅提升
• 代码更紧凑，减少了不必要的中间列操作
• 通用型更强，不依赖DataFrame中的特定列（比如原方案的ID列）

运行这段代码后，得到的结果和你期望的输出完全一致，而且整体逻辑更直接易懂。

内容来源于stack exchange