如何在Polars中基于多列创建排名？

问题分析 #

你尝试用pl.col("era","pred").rank("ordinal")生成排名时，Polars会对每一列单独计算排名，最终生成两个名为rn的列，这就导致了duplicate name错误。而你需要的是基于era和pred组合的全局排名（先按era排序，再按pred排序后的连续排名）。

解决方案 #

方法一：用struct组合列后计算排名 #

将era和pred打包成结构体，Polars会按结构体的字段顺序（先era后pred）进行比较，生成全局唯一的排名列：

import polars as pl

df = pl.DataFrame(
    {   
        "era": ["01", "01", "02", "02", "03", "03"],
        "pred": [3,5,6,8,9,1]
    }
).with_row_index()

result = df.with_columns(
    rn = pl.struct("era", "pred").rank("ordinal")
)
print(result)

输出结果与你的期望完全一致：

shape: (6, 4)
┌───────┬─────┬──────┬─────┐
│ index ┆ era ┆ pred ┆ rn  │
│ ---   ┆ --- ┆ ---  ┆ --- │
│ u32   ┆ str ┆ i64  ┆ u32 │
╞═══════╪═════╪══════╪═════╡
│ 0     ┆ 01  ┆ 3    ┆ 1   │
│ 1     ┆ 01  ┆ 5    ┆ 2   │
│ 2     ┆ 02  ┆ 6    ┆ 3   │
│ 3     ┆ 02  ┆ 8    ┆ 4   │
│ 4     ┆ 03  ┆ 9    ┆ 6   │
│ 5     ┆ 03  ┆ 1    ┆ 5   │
└───────┴─────┴──────┴─────┘

方法二：排序后添加行号（保留原顺序） #

如果需要先排序再生成排名，同时保留原DataFrame的行顺序，可以先对排序后的数据集生成行号，再通过关联合并回去：

ranked_df = df.sort(["era", "pred"]).with_row_index(name="rn", offset=1)
result = df.join(ranked_df, on=["index", "era", "pred"])
print(result)

错误原因说明 #

pl.col("era","pred").rank("ordinal")的逻辑是对era和pred各自独立计算排名，相当于同时执行pl.col("era").rank().alias("rn")和pl.col("pred").rank().alias("rn")，两个列重名导致Polars抛出重复名称错误。

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