Возможно ли в Polars «сбросить» cumsum() при определенном условии? ⇐ Python

[Anonymous]

Мне нужно суммировать столбец b, пока a не станет истинным. После этого совокупная сумма будет начинаться с этой строки и так далее.

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

import polars as pl

df = pl.from_repr("""
┌───────┬─────┐
│ a     ┆ b   │
│ ---   ┆ --- │
│ bool  ┆ i64 │
╞═══════╪═════╡
│ false ┆ 1   │
│ false ┆ 2   │
│ true  ┆ 3   │
│ false ┆ 4   │
└───────┴─────┘
""")

Ожидаемый результат:

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

┌───────┬─────┬─────┐
│ a     ┆ b   ┆ sum │
│ ---   ┆ --- ┆ --- │
│ bool  ┆ i64 ┆ i64 │
╞═══════╪═════╪═════╡
│ false ┆ 1   ┆ 1   │
│ false ┆ 2   ┆ 3   │
│ true  ┆ 3   ┆ 3   │
│ false ┆ 4   ┆ 7   │
└───────┴─────┴─────┘

Могу ли я сделать это на Polars, не повторяя каждую строку?


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