如何为Polars DataFrame的group列生成分组编号列？

给定如下Polars DataFrame：

df = pl.DataFrame({'group': [1, 1, 1, 3, 3, 3, 4, 4]})

数据结构：

shape: (8, 1)
┌───────┐
│ group │
│ ---   │
│ i64   │
╞═══════╡
│ 1     │
│ 1     │
│ 1     │
│ 3     │
│ 3     │
│ 3     │
│ 4     │
│ 4     │
└───────┘

需求是新增一列group_i，为group列的每个唯一分组分配从0开始的连续编号，期望输出：

shape: (8, 2)
┌───────┬─────────┐
│ group ┆ group_i │
│ ---   ┆ ---     │
│ i64   ┆ i64     │
╞═══════╪═════════╡
│ 1     ┆ 0       │
│ 1     ┆ 0       │
│ 1     ┆ 0       │
│ 3     ┆ 1       │
│ 3     ┆ 1       │
│ 3     ┆ 1       │
│ 4     ┆ 2       │
│ 4     ┆ 2       │
└───────┴─────────┘

你当前的实现逻辑较为繁琐：

df.with_columns(((pl.col('group')!=pl.col('group').shift()).cast(pl.Int64).cum_sum()-1).alias('group_i'))

这里提供两种更简洁直观的实现方式：

方法1：使用factorize() #

factorize()是Polars专门用于将离散值转换为连续整数编码的方法，返回的第一个元素就是从0开始的编号序列：

df.with_columns(pl.col('group').factorize()[0].alias('group_i'))

方法2：使用rank()的稠密排名 #

通过rank(method='dense')可以得到每个分组的连续排名（从1开始），减去1后即可得到从0开始的编号：

df.with_columns((pl.col('group').rank(method='dense').cast(pl.Int64) - 1).alias('group_i'))

这两种方法都利用了Polars内置的向量化操作，代码更易读，性能也更稳定。

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