本文介绍在 Polars 中高效生成基于多列组合（如 subject_id 和 timestamp）的唯一整数 ID 的方法，替代 Pandas 的 ngroup()，适用于数据分拆后通过 ID 关联回溯的场景。
本文介绍在 polars 中高效生成基于多列组合（如 subject_id 和 timestamp）的唯一整数 id 的方法，替代 pandas 的 `ngroup()`，适用于数据分拆后通过 id 关联回溯的场景。
在 Polars 中，由于没有内置的 ngroup() 类似接口，但可通过组合 with_row_index()、group_by().agg() 和 explode() 实现等效功能：为每个唯一的 (subject_id, timestamp) 组合分配一个从 0 开始递增的整数 ID，并保持原始行序。
以下是一个通用、可复用的实现方案（适配你的字段名）：
import polars as pl
# 假设原始 DataFrame 名为 df，含列：subject_id, timestamp, event, col1, col2
df_with_id = (
df
.with_row_index("original_idx") # 添加原始行索引（确保顺序可追溯）
.group_by(["subject_id", "timestamp"], maintain_order=True)
.agg(pl.col("original_idx")) # 聚合每组对应的原始索引列表
.with_row_index("group_id") # 为每个唯一组分配 0,1,2,... ID
.explode("original_idx") # 展开索引列表，恢复原始行数
.select("original_idx", "group_id", "subject_id", "timestamp")
.sort("original_idx") # 按原始顺序排列
.drop("original_idx") # 可选：移除临时索引，仅保留 group_id
)
# 将唯一 ID 加入原 DataFrame
df_final = df.with_columns(
group_id=pl.lit(None) # 占位
).join(
df_with_id,
on=["subject_id", "timestamp"],
how="left"
)✅ 关键说明：
- maintain_order=True 确保 group_by 结果按首次出现顺序编号，与原始数据逻辑一致；
- with_row_index() 是稳定且推荐的起点（优于 row_number()），避免因排序或过滤导致错位；
- 最终 group_id 是稠密整数（0,1,2,…），无间隙，适合后续 join 或 group_by("group_id") 操作；
- 若需全局唯一字符串 ID（如 UUID），可用 pl.col("subject_id").cast(pl.Utf8) + "_" + pl.col("timestamp").cast(pl.Utf8)，但整数 ID 更省内存、更高效。
? 最佳实践建议：
- 在分拆前一次性添加 group_id，再分别选取列构建两个子 DataFrame：
df_a = df_final.select("subject_id", "timestamp", "event", "group_id") df_b = df_final.select("subject_id", "timestamp", "col1", "col2", "group_id") - 后续对 df_a 或 df_b 单独聚合后，仍可用 group_id 精准 join 回关联结果，无需依赖原始顺序或重复键匹配。
该方法时间复杂度接近 O(n)，充分利用 Polars 的惰性计算与向量化能力，是生产环境中处理大规模分组标识任务的推荐范式。
