嘿，我来帮你搞定这个需求！你想要的就是给col1和col2组成的每个唯一分组分配一个连续的ID，和R data.table里的.GRP效果完全一致对吧？在Polars里有两种简单直接的方法可以实现：

方法一：使用ngroup()（最直接的方式） #

Polars的ngroup()方法专门用来给每个分组分配唯一的整数ID，默认从0开始计数，我们只需要加1就能得到和你示例里一样从1开始的ID：

import polars as pl

df = pl.DataFrame({
  'col1': ['A', 'A', 'A', 'B'],
  'col2': ['Z', 'Y', 'Z', 'Z']}
)

# 生成group_id列
result = df.with_columns(
    group_id=pl.col('col1').ngroup(by=['col1', 'col2']) + 1
)

print(result)

运行后你会得到完全符合预期的结果：

shape: (4, 3)
┌──────┬──────┬─────────┐
│ col1 ┆ col2 ┆ group_id│
│ ---  ┆ ---  ┆ ---     │
│ str  ┆ str  ┆ i32     │
╞══════╪══════╪═════════╡
│ A    ┆ Z    ┆ 1       │
│ A    ┆ Y    ┆ 2       │
│ A    ┆ Z    ┆ 1       │
│ B    ┆ Z    ┆ 3       │
└──────┴──────┴─────────┘

方法二：使用dense_rank()（通过组合键排名） #

另一种思路是把col1和col2打包成一个结构体，然后对这个结构体做密集排名，这样每个唯一的组合就会得到一个唯一的排名值，也就是我们需要的组ID：

result = df.with_columns(
    group_id=pl.struct(['col1', 'col2']).dense_rank()
)

这个方法同样能得到和示例一致的结果，而且不需要手动加1，因为dense_rank()默认从1开始计数。

补充说明 #

这两种方法都能完美满足你的需求：

• ngroup()是Polars专门为分组场景设计的API，语义更清晰，一看就知道是在生成组ID
• dense_rank()则更灵活，如果你需要基于其他组合键生成ID，这种方式也适用

按照你说的上下文，生成这个group_id后，后续做窗口操作时直接用它作为分组键就行，不用每次都重复写col1和col2，非常方便！

备注：内容来源于stack exchange，提问作者Samuel Allain