Совокупная сумма с group_by ⇐ Python

[Anonymous]

Предположим, у меня есть следующий DataFrame с количеством рождений в каждом штате за каждый год:

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

df = pl.DataFrame(
{
"state": ["CA", "CA", "CA", "TX", "TX", "TX"],
"year": [2010, 2011, 2012, 2010, 2011, 2012],
"births": [10, 20, 30, 40, 50, 60],
}
)

Предположим, что население в 2009 году равно 0 во всех штатах.
Предположим, что число смертей также равно 0 во всех штатах за все годы.
Я хотел бы рассчитать население в каждом штате за каждый год, используя совокупную сумму рождений на уровне штата по годам.
Я придумал:

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

df.group_by("state").agg(pl.col("births").cum_sum().alias("pop")).explode(
"pop"
).with_columns(year=df["year"])

Но для этого требуется добавить обратно столбец года из предыдущего набора данных, что не очень чисто.
Есть ли лучший и более чистый способ?
В документации для cum_sum нет примера с group_by

Подробнее здесь: https://stackoverflow.com/questions/798 ... h-group-by