Polars生成连续整数ID方法

本文揭秘了在 Polars 中如何仅用一行纯表达式链（无需临时变量、walrus 操作符或低效的 join）为满足条件的行（如 event == 1）高效生成连续整数 ID，核心是巧妙组合 pl.when().then() 与 pl.int_range().over() 实现原子化、向量化、可读性强的条件编号——不仅避免了常见 ColumnNotFoundError 和性能陷阱，更以 O(n) 时间复杂度和零冗余计算，展现了 Polars “向量化优先”的最佳实践，特别适合事件序列标记、状态追踪等高频数据工程场景。

本文介绍在 Polars 中通过纯表达式链（无需中间变量或 walrus 操作符）为满足条件的行（如 `event == 1`）分配连续整数 ID 的高效方法，利用 `pl.when().then()` 与 `pl.int_range().over()` 组合实现原子化、可链式调用的列构造。

在 Polars 中，用户常希望通过链式表达式（method chaining）完成复杂变换，但需注意：.join() 等操作内部的子表达式无法访问外部链中新增的列（如 .with_row_index() 添加的 "index"），因为每个子表达式（如 df.select(...).filter(...)）仍以原始 df 为起点独立执行——这正是报错 ColumnNotFoundError: index 的根本原因。

要真正实现“纯链式、无临时变量、无重复计算”的解决方案，推荐使用 Polars 原生的窗口聚合 + 条件表达式组合：

import polars as pl

df = pl.DataFrame(
    {
        "event": [0, 1, 1, 0],
        "foo": [1, 2, 3, 4],
        "boo": [2, 3, 4, 5],
    }
)

result = (
    df
    .with_columns(
        pl.when(pl.col("event") == 1)
        .then(pl.int_range(0, pl.len()).over("event"))
        .alias("event_id")
    )
)
print(result)

输出：

shape: (4, 4)
┌───────┬─────┬─────┬──────────┐
│ event ┆ foo ┆ boo ┆ event_id │
│ ---   ┆ --- ┆ --- ┆ ---      │
│ i64   ┆ i64 ┆ i64 ┆ i64      │
╞═══════╪═════╪═════╪══════════╡
│ 0     ┆ 1   ┆ 2   ┆ null     │
│ 1     ┆ 2   ┆ 3   ┆ 0        │
│ 1     ┆ 3   ┆ 4   ┆ 1        │
│ 0     ┆ 4   ┆ 5   ┆ null     │
└───────┴─────┴─────┴──────────┘

✅ 关键原理说明：

• pl.int_range(0, pl.len()) 生成从 0 到当前分组行数（不含）的整数序列；
• .over("event") 将其按 "event" 分组计算——由于我们只对 event == 1 的组赋值，而 event == 0 的组被 pl.when(...).then(...) 排除，因此仅 event=1 的连续行获得 0, 1, 2, ... 编号；
• pl.when(...).then(...) 实现条件掩码：不满足条件的行自动设为 null，语义清晰且无需 .join() 或显式索引对齐。

⚠️ 注意事项：

• pl.int_range().over() 依赖分组内行序；若原始数据顺序敏感，请确保未发生意外重排（如避免无序聚合后未 .sort()）；
• 若需保留原始 "index" 列用于后续关联，可先 .with_row_index() 再 .with_columns(...)，二者兼容；
• 此方案时间复杂度为 O(n)，远优于 .join() 的 O(n log n)（含哈希/排序开销），且内存更友好。

该模式体现了 Polars “向量化优先、避免显式循环与连接”的设计哲学，是处理条件编号、事件序列化、状态标记等场景的推荐范式。

以上就是《Polars生成连续整数ID方法》的详细内容，更多关于的资料请关注golang学习网公众号！