Count consecutive True (or 1) values in a Boolean (or numeric) column with Polars?
17:19 09 Feb 2023

I am hoping to count consecutive values in a column, preferably using Polars expressions.

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

With the example data frame above, I would like to count the number of consecutive True values.

Below is an example using R's Data.Table package.

library(data.table)
dt <- data.table(value = c(T,T,T,F,F,T,F,F,T,T))
dt[, value2 := fifelse((1:.N) == .N & value == 1, .N, NA_integer_), by = rleid(value)]
dt

Desired result displayed as a Polars DataFrame

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    │
└───────┴──────┘

Any ideas how this would be done efficiently using Polars?

python dataframe window-functions python-polars run-length-encoding