Python Polars如何计算自上次非零值以来的行数？

针对给定的Polars DataFrame，我们可以通过分组计数的方式生成目标dist列，同时满足支持.over()分组与兼容滚动窗口函数的要求：

实现代码 #

import polars as pl

df = pl.DataFrame({"a": [0, 29, 28, 4, 0, 0, 13, 0]})

result = df.with_columns(
    dist=pl.when(pl.col("a") == 0)
           .then(pl.int_range(1, pl.count() + 1).over(pl.col("a").ne(0).cumsum()))
           .otherwise(0)
)

print(result)

逻辑说明 #

1. 生成分组标识：通过pl.col("a").ne(0).cumsum()创建分组键——非0值会触发累加，连续的0会被归入同一分组，非0值各自独立成组。
2. 组内计数赋值：使用when/then/otherwise分支逻辑：
   • 当a为0时，在对应分组内用pl.int_range(1, pl.count() + 1)生成从1开始的连续计数；
   • 当a非0时，直接赋值为0。
3. 扩展性支持：该方案天然兼容.over()多列分组（只需在.over()中添加额外分组列），且不会破坏原数据的序列性，可与rolling_mean()等滚动窗口函数配合使用。

输出结果 #

运行代码后将得到目标DataFrame：

shape: (8, 2)
┌─────┬──────┐
│ a   ┆ dist │
│ --- ┆ ---  │
│ i64 ┆ i64  │
╞═════╪══════╡
│ 0   ┆ 1    │
│ 29  ┆ 0    │
│ 28  ┆ 0    │
│ 4   ┆ 0    │
│ 0   ┆ 1    │
│ 0   ┆ 2    │
│ 13  ┆ 0    │
│ 0   ┆ 1    │
└─────┴──────┘

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