Увеличение значения в зависимости от условия ⇐ Python

[Anonymous]

Я хочу увеличить значение столбца на основе определенного условия в кадре данных Polars, учитывая при этом, сколько раз это условие было выполнено.
Пример данных. >
import polars as pl

df = pl.DataFrame({
"before": [0, 0, 0, 0, 0, 0, 0, 0, 0],
"cdl_type": ["REC", "REC", "GEC", None, None, "GEC", None, "REC", "GEC"],
})

Текущий подход.
df = df.with_columns(
a=(
pl.when(pl.col("cdl_type").is_in(["GEC", "REC"])).then(
pl.int_ranges(
pl.col("cdl_type")
.is_in(["REC", "GEC"])
.rle()
.struct.field("len")
).flatten()
)
.when(pl.col('cdl_type').is_null().and_(pl.col('cdl_type').shift(1).is_not_null()))
.then(pl.lit(1))
.otherwise(0)
)
)

Ожидаемый результат.
┌────────┬──────────┬───────┐
│ before ┆ cdl_type ┆ after │
│ --- ┆ --- ┆ --- │
│ i64 ┆ str ┆ i64 │
╞════════╪══════════╪═══════╡
│ 0 ┆ REC ┆ 0 │
│ 0 ┆ REC ┆ 1 │
│ 0 ┆ GEC ┆ 2 │
│ 0 ┆ null ┆ 3 │
│ 0 ┆ null ┆ 0 │
│ 0 ┆ GEC ┆ 0 │
│ 0 ┆ null ┆ 1 │
│ 0 ┆ REC ┆ 0 │
│ 0 ┆ GEC ┆ 1 │
└────────┴──────────┴───────┘


Подробнее здесь: https://stackoverflow.com/questions/791 ... -condition