如何在Polars中实现pandas的ngroup功能并分配组ID

需求说明 #

给DataFrame的指定分组列分配连续整数ID（从0开始），对应pandas中groupby().ngroup()的功能，且方案需支持百万级规模数据集的高效处理。

示例数据集 #

import polars as pl

df = pl.DataFrame(
     {"group_a": ["a", "aa", "a"], "group_b": ["b", "bb", "b"], "val": [1, 2, 3]}
)

最优实现方案 #

方案1：使用factorize()（推荐）

factorize()是Polars专门用于生成唯一分组ID的矢量化方法，内存效率极高，适合大规模数据：

df = df.with_columns(
    # factorize返回元组，第一个元素是每行对应的分组ID（从0开始）
    pl.col(["group_a", "group_b"]).factorize()[0].alias("new_group")
)

方案2：使用窗口函数dense_rank()

通过窗口函数实现连续排名，调整起始值后与ngroup效果一致：

df = df.with_columns(
    # dense_rank默认从1开始，减1转为0起始的ID
    (pl.dense_rank().over(["group_a", "group_b"]) - 1).alias("new_group")
)

执行结果 #

两种方案均会得到期望输出：

shape: (3, 4)
┌─────────┬─────────┬─────┬───────────┐
│ group_a ┆ group_b ┆ val ┆ new_group │
│ ---     ┆ ---     ┆ --- ┆ ---       │
│ str     ┆ str     ┆ i64 ┆ i64       │
╞═════════╪═════════╪═════╪═══════════╡
│ a       ┆ b       ┆ 1   ┆ 0         │
│ aa      ┆ bb      ┆ 2   ┆ 1         │
│ a       ┆ b       ┆ 3   ┆ 0         │
└─────────┴─────────┴─────┴───────────┘

临时方案的问题分析 #

你之前的临时方案通过unique提取分组、生成ID后join回去，存在核心缺陷：

• 需要生成中间唯一分组数据集，当百万级数据包含大量唯一分组时，会占用大量内存；
• join操作的时间和内存开销极高，容易引发OOM（内存不足）崩溃。

上述两种最优方案均为矢量化操作，无需额外中间数据，完全在Polars查询引擎中高效处理，能够轻松应对大规模数据集。

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