[Python] Calculate group mean for an int_range column in Polars dataframe

Update: This has been resolved in Polars. The code now runs without error.


I have the following code in polars:

import datetime
import polars as pl

df = pl.DataFrame(
{
"id": [1, 2, 1, 2, 1, 2, 3],
"date": [
datetime.date(2022, 1, 1),
datetime.date(2022, 1, 1),
datetime.date(2022, 1, 11),
datetime.date(2022, 1, 11),
datetime.date(2022, 2, 1),
datetime.date(2022, 2, 1),
datetime.date(2022, 2, 1),
],
"value": [1, 2, 3, None, 5, 6, None],
}
)

(df.group_by_dynamic("date", group_by="id", every="1mo", period="1mo", closed="both")
.agg(
pl.int_range(1, pl.len() + 1)
- pl.int_range(1, pl.len() + 1).filter(pl.col("value").is_not_null()).mean(),
)
)


But, when I run it, I got the following error which I don't quite understand.

pyo3_runtime.PanicException: index out of bounds: the len is 1 but the index is 1


The behavior I want to achieve is: for each group, create a natural sequence from 1 to the number of rows in that group, and subtract from it the average over non-null in the "value" column in that group. (return null if all "value" in that group are null).

To be more specific, the result I want is

shape: (5, 3)
┌─────┬────────────┬──────────────────┐
│ id ┆ date ┆ arange │
│ --- ┆ --- ┆ --- │
│ i64 ┆ date ┆ list[f64] │
╞═════╪════════════╪══════════════════╡
│ 1 ┆ 2022-01-01 ┆ [-1.0, 0.0, 1.0] │
│ 1 ┆ 2022-02-01 ┆ [0.0] │
│ 2 ┆ 2022-01-01 ┆ [-1.0, 2.0, 1.0] │
│ 2 ┆ 2022-02-01 ┆ [0.0] │
│ 3 ┆ 2022-02-01 ┆ [null] │
└─────┴────────────┴──────────────────┘


How can I achieve this?

Continue reading...