Calculate group mean for an int_range column in Polars dataframe
10:04 17 Feb 2023

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?

python dataframe time-series python-polars