如何在Polars中复现Pandas的ngroup连续分组索引功能？

我需要在Polars中复现Pandas ngroup方法的行为——为按两列分组的DataFrame生成连续的组索引。熟悉R语言的话，这对应dplyr里的group_indices或更新的cur_group_id功能。

我试了几种方法都没成功，这些方法只会返回每组的行数，没法生成连续的组索引，示例代码及结果如下：

import polars as pl
import pandas as pd

df = pd.DataFrame(
    {
        "id": ["a", "a", "a", "a", "b", "b", "b", "b"],
        "cat": [1, 1, 2, 2, 1, 1, 2, 2],
    }
)

df_pl = pl.from_pandas(df)

# 期望的结果（Pandas的ngroup输出）
print(df.groupby(["id", "cat"]).ngroup())
# 0    0
# 1    0
# 2    1
# 3    1
# 4    2
# 5    2
# 6    3
# 7    3

# 尝试的方法1：仅统计每组行数
print(df_pl.select(pl.len().over("id", "cat")))
# ┌─────┐
# │ len │
# │ --- │
# │ u32 │
# ╞═════╡
# │ 2   │
# │ 2   │
# │ 2   │
# │ 2   │
# │ 2   │
# │ 2   │
# │ 2   │
# │ 2   │
# └─────┘

# 尝试的方法2：分组聚合后仍只得到行数
print(df_pl.group_by("id", "cat").agg(pl.len().alias("test")))
# shape: (4, 3)
# ┌─────┬─────┬──────┐
# │ id  ┆ cat ┆ test │
# │ --- ┆ --- ┆ ---  │
# │ str ┆ i64 ┆ u32  │
# ╞═════╪═════╪══════╡
# │ a   ┆ 1   ┆ 2    │
# │ a   ┆ 2   ┆ 2    │
# │ b   ┆ 1   ┆ 2    │
# │ b   ┆ 2   ┆ 2    │
# └─────┴─────┴──────┘

解决方法 #

方法1：使用factorize配合元组分组键

通过factorize对id和cat的组合进行编码，直接生成连续组索引并广播到每一行：

result = df_pl.with_columns(
    pl.factorize((pl.col("id"), pl.col("cat")))[0].alias("ngroup")
)
print(result)

输出：

shape: (8, 3)
┌─────┬─────┬────────┐
│ id  ┆ cat ┆ ngroup │
│ --- ┆ --- ┆ ---    │
│ str ┆ i64 ┆ u32    │
╞═════╪═════╪════════╡
│ a   ┆ 1   ┆ 0      │
│ a   ┆ 1   ┆ 0      │
│ a   ┆ 2   ┆ 1      │
│ a   ┆ 2   ┆ 1      │
│ b   ┆ 1   ┆ 2      │
│ b   ┆ 1   ┆ 2      │
│ b   ┆ 2   ┆ 3      │
│ b   ┆ 2   ┆ 3      │
└─────┴─────┴────────┘

方法2：分组后用dense_rank生成索引（Polars 0.18+）

利用窗口函数的dense_rank对分组键进行排序编码，再减1得到从0开始的索引：

result = df_pl.with_columns(
    pl.dense_rank().over(["id", "cat"]) - 1
)

方法3：先给唯一分组分配ID再合并

先提取唯一分组并添加行索引作为组ID，再通过连接映射回原数据：

# 给每个唯一分组分配ID
group_ids = df_pl.unique(["id", "cat"]).with_row_index("ngroup")
# 合并回原DataFrame
result = df_pl.join(group_ids, on=["id", "cat"])

以上三种方法都能生成和Pandas ngroup完全一致的连续组索引结果。

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