Count consecutive True (or 1) values in a Boolean (or numeric) column with Polars? - python-polars

Question

I am hoping to count consecutive values in a column, preferably using Polars expressions.
import polars
df = pl.DataFrame(
{"values": [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 example output 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
value
value2
TRUE
NA
TRUE
NA
TRUE
3
FALSE
NA
FALSE
NA
TRUE
1
FALSE
NA
FALSE
NA
TRUE
NA
TRUE
2
Any ideas who this would be done efficiently using Polars?
[EDIT with a new approach]
I got it working with the code below, but hoping there is a more efficient way. Anyone know the default struct/dictionary field names from value_counts?
(
df.lazy()
.with_row_count()
.with_column(
pl.when(pl.col("value") == False).then(
pl.col("row_nr")
).fill_null(
strategy = "forward"
).alias("id_consecutive_Trues")
)
.with_column(
pl.col("id_consecutive_Trues").value_counts(sort = True)
)
.with_column(
(
pl.col("id_consecutive_Trues").arr.eval(
pl.element().struct().rename_fields(["value", "count"]).struct.field("count")
).arr.max()
- pl.lit(1)
).alias("max_consecutive_true_values")
)
.collect()
)

Answer 1

It can be thought of as a groupby operation.
A common way to generate group IDs for consecutive values is to check != .shift() and use the .cumsum()
I've put each step into its own .with_columns() here but it could be simplified:
(
df
.with_columns(
(pl.col("values") != pl.col("values").shift(+1))
.alias("change"))
.with_columns(
pl.col("change").cumsum()
.alias("group"))
.with_columns(
pl.count().over("group"))
.with_columns(
(pl.col("values") != pl.col("values").shift(-1))
.alias("keep"))
.with_columns(
pl.when(pl.col("values") & pl.col("keep"))
.then(pl.col("count"))
.alias("values2"))
)
shape: (10, 6)
┌────────┬────────┬───────┬───────┬───────┬─────────┐
│ values | change | group | count | keep | values2 │
│ --- | --- | --- | --- | --- | --- │
│ bool | bool | u32 | u32 | bool | u32 │
╞════════╪════════╪═══════╪═══════╪═══════╪═════════╡
│ true | true | 1 | 3 | false | null │
│ true | false | 1 | 3 | false | null │
│ true | false | 1 | 3 | true | 3 │
│ false | true | 2 | 2 | false | null │
│ false | false | 2 | 2 | true | null │
│ true | true | 3 | 1 | true | 1 │
│ false | true | 4 | 2 | false | null │
│ false | false | 4 | 2 | true | null │
│ true | true | 5 | 2 | false | null │
│ true | false | 5 | 2 | true | 2 │
└────────┴────────┴───────┴───────┴───────┴─────────┘
One possible way to write it in a "less verbose" manner:
df.with_columns(
pl.when(
pl.col("values") &
pl.col("values").shift_and_fill(-1, False).is_not())
.then(
pl.count().over(
(pl.col("values") != pl.col("values").shift())
.cumsum()))
)
shape: (10, 2)
┌────────┬───────┐
│ values | count │
│ --- | --- │
│ bool | u32 │
╞════════╪═══════╡
│ true | null │
│ true | null │
│ true | 3 │
│ false | null │
│ false | null │
│ true | 1 │
│ false | null │
│ false | null │
│ true | null │
│ true | 2 │
└────────┴───────┘