本文介绍两种在 Polars 中为多列组合生成稳定、连续整数型 group_id 的方法：一种基于结构体排序的简洁方案，另一种保留原始行序的精确方案，并附带代码示例与关键注意事项。
本文介绍两种在 polars 中为多列组合生成稳定、连续整数型 `group_id` 的方法：一种基于结构体排序的简洁方案，另一种保留原始行序的精确方案，并附带代码示例与关键注意事项。
在 Polars 中生成类似 R data.table 中 .GRP 的 group_id（即为每个唯一分组分配一个从 1 开始的连续整数 ID），是构建事件标识符、支持后续窗口计算或去重聚合的关键步骤。核心挑战在于：既要保证相同 (col1, col2) 组合获得相同 ID，又要控制 ID 分配顺序是否依赖原始数据顺序。以下是两种主流实现方式：
✅ 方案一：基于结构体排序（简洁通用，但顺序由字典序决定）
当分组 ID 的数值大小仅需唯一、连续，且不严格要求“首次出现的组分配最小 ID”时，可直接对组合字段构造结构体并调用 rank("dense")：
import polars as pl
df = pl.DataFrame({
'col1': ['A', 'A', 'A', 'B'],
'col2': ['Z', 'Y', 'Z', 'Z']
})
df_with_id = df.with_columns(
group_id = pl.struct("col1", "col2").rank("dense").cast(pl.UInt32)
)
输出：
shape: (4, 3) ┌──────┬──────┬──────────┐ │ col1 ┆ col2 ┆ group_id │ │ --- ┆ --- ┆ --- │ │ str ┆ str ┆ u32 │ ╞══════╪══════╪══════════╡ │ A ┆ Z ┆ 2 │ │ A ┆ Y ┆ 1 │ │ A ┆ Z ┆ 2 │ │ B ┆ Z ┆ 3 │ └──────┴──────┴──────────┘
⚠️ 注意：此方法按
struct字段的默认字典序排序（('A','Y') ），因此 <code>'A','Y'得到1，而非按其在原 DataFrame 中首次出现的位置。
✅ 方案二：基于行索引的稳定分组 ID（推荐用于事件建模）
若需严格保证“首次出现的 (col1,col2) 组合获得 1，第二次新组获得 2”，则必须锚定原始行序。推荐使用 with_row_index() + min().over() + rank("dense") 链式操作：
df_with_id = (
df.with_row_index("row_idx")
.with_columns(
group_id = pl.col("row_idx").min().over("col1", "col2").rank("dense").cast(pl.UInt32)
)
.drop("row_idx") # 清理临时索引列
)
输出（完全匹配问题中期望结果）：
shape: (4, 3) ┌──────┬──────┬──────────┐ │ col1 ┆ col2 ┆ group_id │ │ --- ┆ --- ┆ --- │ │ str ┆ str ┆ u32 │ ╞══════╪══════╪══════════╡ │ A ┆ Z ┆ 1 │ │ A ┆ Y ┆ 2 │ │ A ┆ Z ┆ 1 │ │ B ┆ Z ┆ 3 │ └──────┴──────┴──────────┘
该方案本质是：
- 为每行添加唯一、递增的
row_idx； - 对每个
(col1, col2)分组，取其所有行中最小的row_idx（即该组首次出现位置）； - 对这些“首次出现索引”进行密集排名（
rank("dense")），确保无跳号、无重复。
? 关键总结与建议
-
类型安全：显式
.cast(pl.UInt32)可避免rank()默认返回Float64类型，提升下游兼容性； -
性能考量：
over()是 Polars 高效的窗口操作，适用于百万级数据； -
扩展性：
over()支持任意数量分组列（如.over("col1", "col2", "col3")），无需修改逻辑； -
替代方案：
pl.col("col1").hash(seed=0).over("col2")等哈希法虽快但不保证 ID 连续，不适用于需排序/范围查询的场景； -
事件建模提示：生成的
group_id可直接用于pl.col("value").over("group_id")等窗口计算，替代冗长的多列over(...)表达式，大幅提升可读性与维护性。
