Прокатка по группе в полярах ⇐ Python

[Anonymous]

Предположим, у меня есть следующий df:

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

import numpy as np
import polars as pl

df = pl.DataFrame({
"date": pl.date_range(pl.date(2022, 1, 1), pl.date(2022, 1, 30), interval="1d", eager=True),
"group": ["A", "B", "C"]*10,
"total": np.random.randint(100, size=30)
})

Теперь я хочу подсчитать общую сумму по группе за временной интервал в 5 дней. Что-то вроде комбинирования группировки с прокруткой.
Интересно, есть ли более элегантный и эффективный способ сделать это, чем следующий код:

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

df_list = []

for table_group in df["group"].unique():
temp_df = (
df
.filter(pl.col("group")==table_group)
.sort("date")
.rolling("date", period="5d")
.agg(pl.col("total").sum().name.prefix("roll_"))
)

temp_df2 = (
df
.filter(pl.col("group")==table_group)
.sort("date")
.select(
pl.all(),
temp_df.get_column("roll_total")
)
)
df_list.append(temp_df2)

pl.concat(df_list)

Моя основная проблема заключается в том, что у меня миллионы строк и более 10 000 групп (и более 1 столбца, который я хочу суммировать)...
Поэтому этот код не так элегантен и потребляет больше памяти, чем мне хотелось бы.
>

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