如何将SQL的row_number()窗口函数改写为Python Polars实现？

你的代码问题在于错误地将整个DataFrame的行计数和排序操作嵌套在窗口函数中，Polars的窗口函数是基于表达式级别的操作，不需要先对整个DataFrame做修改，直接通过pl.row_number()结合.over()方法就能实现SQL中的row_number() over (...)逻辑。

正确实现代码 #

import polars as pl

df = pl.DataFrame(
    {
        "a": ["a", "b", "a", "b", "c"],
        "b": ["Z", "Z", "T", "Z", "T"],
        "c": [1, 2, 1, 3, 3],
        "d": [5, 4, 3, 2, 1],
    }
)

result = df.with_columns(
    pl.row_number().over(
        partition_by=["a", "b"],
        order_by=[pl.col("c") * pl.col("d").desc(nulls_last=True)]
    ).alias("rn")
)

print(result)

代码对应SQL逻辑说明 #

• pl.row_number() 对应SQL中的 row_number()
• .over(partition_by=["a", "b"]) 对应SQL中的 partition by a,b
• order_by=[pl.col("c") * pl.col("d").desc(nulls_last=True)] 对应SQL中的 order by c*d desc nulls last，通过.desc(nulls_last=True)直接指定降序且空值后置的规则

运行结果 #

shape: (5, 5)
┌─────┬─────┬─────┬─────┬─────┐
│ a   ┆ b   ┆ c   ┆ d   ┆ rn  │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ str ┆ i64 ┆ i64 ┆ u32 │
╞═════╪═════╪═════╪═════╪═════╡
│ a   ┆ Z   ┆ 1   ┆ 5   ┆ 1   │
│ b   ┆ Z   ┆ 2   ┆ 4   ┆ 1   │
│ b   ┆ Z   ┆ 3   ┆ 2   ┆ 2   │
│ a   ┆ T   ┆ 1   ┆ 3   ┆ 1   │
│ c   ┆ T   ┆ 3   ┆ 1   ┆ 1   │
└─────┴─────┴─────┴─────┴─────┘

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