本文介绍在 Polars 中通过纯表达式链（无需中间变量或 walrus 操作符）为满足条件的行（如 event == 1）分配连续整数 ID 的高效方法，利用 pl.when().then() 与 pl.int_range().over() 组合实现原子化、可链式调用的列构造。
本文介绍在 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 “向量化优先、避免显式循环与连接”的设计哲学，是处理条件编号、事件序列化、状态标记等场景的推荐范式。
