如何在Python Polars中对DataFrame指定列按分组计算百分位排名

问题背景 #

给定如下Polars DataFrame：

df = pl.DataFrame(
    {   
        "era": ["01", "01", "02", "02", "03", "03"],
        "pred1": [1, 2, 3, 4, 5,6],
        "pred2": [2,4,5,6,7,8],
        "pred3": [3,5,6,8,9,1],
        "something_else": [5,4,3,67,5,4],
    }
)
pred_cols = ["pred1", "pred2", "pred3"]
ERA_COL = "era"

需要实现和以下Pandas代码等价的功能：按era分组，对pred_cols列计算百分位排名（rank(pct=True)）：

df[list(pred_cols)] = df.groupby(ERA_COL, group_keys=False).apply(
    lambda d: d[list(pred_cols)].rank(pct=True)
)

Polars的rank函数没有pct参数，此前参考的方案返回数组而非单个浮点数，不符合需求。

解决方案 #

核心逻辑是手动实现百分位计算：(排名 - 1) / (分组内元素数量 - 1)，这和Pandas默认method='average'的rank(pct=True)逻辑完全一致。

利用Polars的分组窗口函数实现，代码如下：

import polars as pl

df = pl.DataFrame(
    {   
        "era": ["01", "01", "02", "02", "03", "03"],
        "pred1": [1, 2, 3, 4, 5,6],
        "pred2": [2,4,5,6,7,8],
        "pred3": [3,5,6,8,9,1],
        "something_else": [5,4,3,67,5,4],
    }
)
pred_cols = ["pred1", "pred2", "pred3"]
ERA_COL = "era"

result_df = df.with_columns(
    [
        (pl.col(col).rank().over(ERA_COL) - 1) / (pl.count().over(ERA_COL) - 1)
        for col in pred_cols
    ]
)

print(result_df)

结果说明 #

运行后输出的pred1、pred2、pred3列数值与Pandas代码结果完全一致：

shape: (6, 5)
┌─────┬───────┬───────┬───────┬────────────────┐
│ era ┆ pred1 ┆ pred2 ┆ pred3 ┆ something_else │
│ --- ┆ ---   ┆ ---   ┆ ---   ┆ ---            │
│ str ┆ f64   ┆ f64   ┆ f64   ┆ i64            │
╞═════╪═══════╪═══════╪═══════╪════════════════╡
│ 01  ┆ 0.0   ┆ 0.0   ┆ 0.0   ┆ 5              │
│ 01  ┆ 1.0   ┆ 1.0   ┆ 1.0   ┆ 4              │
│ 02  ┆ 0.0   ┆ 0.0   ┆ 0.0   ┆ 3              │
│ 02  ┆ 1.0   ┆ 1.0   ┆ 1.0   ┆ 67             │
│ 03  ┆ 0.0   ┆ 0.0   ┆ 1.0   ┆ 5              │
│ 03  ┆ 1.0   ┆ 1.0   ┆ 0.0   ┆ 4              │
└─────┴───────┴───────┴───────┴────────────────┘

代码逻辑拆解 #

• pl.col(col).rank().over(ERA_COL)：按era分组，计算目标列每个值的排名（默认平均排名，与Pandas默认规则对齐）
• pl.count().over(ERA_COL)：获取每个分组的总行数
• 公式(rank - 1)/(n - 1)：将排名转换为0到1区间的百分位值，完全匹配Pandasrank(pct=True)的计算逻辑

内容的提问来源于stack exchange，提问作者jbssm