Как сбросить значение cum_sum на основе условия в полярах ⇐ Python

[Anonymous]

У меня есть следующий DF, на котором я пытаюсь создать новый столбец, который будет cum_sum.

Код: Выделить всё

import polars as pl
pl.Config(tbl_rows=20)
df= pl.LazyFrame({
"a" : [0,1,2,3,4,5,6,7,8,9],
"b" : [33,35,33,32,33,37,33,29,34,36],
"c" : [33,35,35,35,35,37,37,37,37,37],
})
df = df.with_columns(pl.when(pl.col('b') == pl.col('c')).then(0).otherwise(1).alias('temp'))
df = df.with_columns(pl.when(pl.col('temp') == 0).then(0).otherwise(pl.col('temp').cum_sum()).alias('c_sum_n'))

В настоящее время я получаю следующий результат:

Код: Выделить всё

>>> df.collect()
shape: (10, 5)
┌─────┬─────┬─────┬──────┬─────────┐
│ a   ┆ b   ┆ c   ┆ temp ┆ c_sum_n │
│ --- ┆ --- ┆ --- ┆ ---  ┆ ---     │
│ i64 ┆ i64 ┆ i64 ┆ i32  ┆ i32     │
╞═════╪═════╪═════╪══════╪═════════╡
│ 0   ┆ 33  ┆ 33  ┆ 0    ┆ 0       │
│ 1   ┆ 35  ┆ 35  ┆ 0    ┆ 0       │
│ 2   ┆ 33  ┆ 35  ┆ 1    ┆ 1       │
│ 3   ┆ 32  ┆ 35  ┆ 1    ┆ 2       │
│ 4   ┆ 33  ┆ 35  ┆ 1    ┆ 3       │
│ 5   ┆ 37  ┆ 37  ┆ 0    ┆ 0       │
│ 6   ┆ 33  ┆ 37  ┆ 1    ┆ 4       │
│ 7   ┆ 29  ┆ 37  ┆ 1    ┆ 5       │
│ 8   ┆ 34  ┆ 37  ┆ 1    ┆ 6       │
│ 9   ┆ 36  ┆ 37  ┆ 1    ┆ 7       │
└─────┴─────┴─────┴──────┴─────────┘

но желаемый результат, который я ищу,

Код: Выделить всё

>>> df.collect()
shape: (10, 5)
┌─────┬─────┬─────┬──────┬─────────┐
│ a   ┆ b   ┆ c   ┆ temp ┆ c_sum_n │
│ --- ┆ --- ┆ --- ┆ ---  ┆ ---     │
│ i64 ┆ i64 ┆ i64 ┆ i32  ┆ i32     │
╞═════╪═════╪═════╪══════╪═════════╡
│ 0   ┆ 33  ┆ 33  ┆ 0    ┆ 0       │
│ 1   ┆ 35  ┆ 35  ┆ 0    ┆ 0       │
│ 2   ┆ 33  ┆ 35  ┆ 1    ┆ 1       │
│ 3   ┆ 32  ┆ 35  ┆ 1    ┆ 2       │
│ 4   ┆ 33  ┆ 35  ┆ 1    ┆ 3       │
│ 5   ┆ 37  ┆ 37  ┆ 0    ┆ 0       │
│ 6   ┆ 33  ┆ 37  ┆ 1    ┆ 1       │
│ 7   ┆ 29  ┆ 37  ┆ 1    ┆ 2       │
│ 8   ┆ 34  ┆ 37  ┆ 1    ┆ 3       │
│ 9   ┆ 36  ┆ 37  ┆ 1    ┆ 4       │
└─────┴─────┴─────┴──────┴─────────┘

i,e, когда col b == col c, я хотел бы сбросить значение cum_sum до нуля и начать все сначала. Любая помощь, будь то поляры Python или поляры ржавчины, будет полезна.

Подробнее здесь: https://stackoverflow.com/questions/788 ... -in-polars