如何用Polars实现Postgres中ROW_NUMBER() OVER(PARTITION BY)等价功能？

最优方案：使用Polars内置row_number()函数 #

Polars 0.17及以上版本提供了专门对应SQL ROW_NUMBER() OVER (PARTITION BY ...)的pl.row_number()窗口函数，完全替代你原来的繁琐写法：

import polars as pl

df = pl.DataFrame(
    {"foo": [2,1,3,7,6,32], "tech": ['python', 'python', 'java','rust','c','rust'], "service": ["a","a","a","b","c","b"]}
)

# 直接实现ROW_NUMBER按tech和service分组
result = df.with_columns(
    pl.row_number().over(['tech', 'service']).alias('rn')
)

print(result)

执行结果和你原来的代码完全一致：

shape: (6, 4)
┌─────┬────────┬─────────┬─────┐
│ foo ┆ tech   ┆ service ┆ rn  │
│ --- ┆ ---    ┆ ---     ┆ --- │
│ i64 ┆ str    ┆ str     ┆ u32 │
╞═════╪════════╪═════════╪═════╡
│ 2   ┆ python ┆ a       ┆ 1   │
│ 1   ┆ python ┆ a       ┆ 2   │
│ 3   ┆ java   ┆ a       ┆ 1   │
│ 7   ┆ rust   ┆ b       ┆ 1   │
│ 6   ┆ c      ┆ c       ┆ 1   │
│ 32  ┆ rust   ┆ b       ┆ 2   │
└─────┴────────┴─────────┴─────┘

旧版本Polars兼容方案（低于0.17） #

如果你的Polars版本较低，无法使用row_number()，可以用pl.int_range结合窗口计数实现：

result = df.with_columns(
    pl.int_range(1, pl.count().over(['tech', 'service']) + 1).flatten().alias('rn')
)

这个方法通过count()获取每个分组的行数，生成对应长度的连续整数序列，再通过flatten()映射到每一行，效果和原生函数一致。

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