如何用Polars高效统计布尔列中连续True值的长度？

需求说明 #

需要统计Polars DataFrame中布尔列的连续True值长度，仅在每组连续True的最后一行标注该组的长度，其余行（包括False行）均为null，最终结果格式如下：

shape: (10, 2)
┌───────┬──────┐
│ value ┆ len  │
│ ---   ┆ ---  │
│ bool  ┆ u32  │
╞═══════╪══════╡
│ true  ┆ null │
│ true  ┆ null │
│ true  ┆ 3    │
│ false ┆ null │
│ false ┆ null │
│ true  ┆ 1    │
│ false ┆ null │
│ false ┆ null │
│ true  ┆ null │
│ true  ┆ 2    │
└───────┴──────┘

示例DataFrame #

首先创建测试用的DataFrame：

import polars as pl
df = pl.DataFrame(
   {"value": [True,True,True,False,False,True,False,False,True,True]}
)

Polars实现方案 #

借助Polars的窗口函数和rle_id功能，可以高效实现需求，代码如下：

result = df.with_columns(
    # 为连续相同的value生成分组ID
    group_id=pl.rle_id("value"),
    # 判断当前行是否为所在分组的最后一行
    is_last=pl.int_range(0, pl.count()).over("group_id") == (pl.count() - 1).over("group_id")
).with_columns(
    # 仅在value为True且是组内最后一行时填充连续长度，否则为null
    len=pl.when(pl.col("value") & pl.col("is_last"))
          .then(pl.count().over("group_id"))
          .cast(pl.UInt32)
          .alias("len")
).drop("group_id", "is_last")

print(result)

代码解释 #

1. pl.rle_id("value")：生成连续相同值的分组ID，和R data.table中的rleid(value)功能一致，用于将连续的相同值归为一组。
2. 标记组内最后一行：通过窗口函数pl.int_range生成组内行号，和组内总行数减1对比，判断是否为组内最后一行。
3. 填充连续长度：使用when/then条件逻辑，仅在value为True且是组内最后一行时，填入该组的总行数（即连续True的长度），其余情况保留null，同时转换类型为UInt32匹配期望结果。

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