[Python] Polars-python: Time series shift creation using expressions

I am using Polars to perform time series analysis. My time series may not be complete, and the frequency of the dates may change by id. My objective is to obtain the column "result" below, using a with_columns expression.

import datetime as dt
import polars as pl

df = pl.DataFrame(
{
"id": ["1", "1", "1", "2", "2", "2"],
"reported_period": [
dt.date(2021, 3, 31),
dt.date(2021, 9, 30),
dt.date(2021, 12, 31),
dt.date(2021, 3, 31),
dt.date(2021, 6, 30),
dt.date(2021, 9, 30),
],
"variable": [10, 30, 40, 1, 3, 4],
"result": [None, None, 30, None, 1, 3],
}
)
shape: (6, 4)
┌─────┬─────────────────┬──────────┬────────┐
│ id ┆ reported_period ┆ variable ┆ result │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ date ┆ i64 ┆ i64 │
╞═════╪═════════════════╪══════════╪════════╡
│ 1 ┆ 2021-03-31 ┆ 10 ┆ null │
│ 1 ┆ 2021-09-30 ┆ 30 ┆ null │
│ 1 ┆ 2021-12-31 ┆ 40 ┆ 30 │
│ 2 ┆ 2021-03-31 ┆ 1 ┆ null │
│ 2 ┆ 2021-06-30 ┆ 3 ┆ 1 │
│ 2 ┆ 2021-09-30 ┆ 4 ┆ 3 │
└─────┴─────────────────┴──────────┴────────┘



So far, I thought about using a helper column (let's say I want just 1Q), and then loop over the columns I have, but it is not working as I cannot aggregate over the get() expression. I paste the code below. Could you please help me, or propose any alternative approach -but keeping a single expression-):


df = df.with_columns(
pl.col("reported_period")
.dt.offset_by("-1q")
.dt.month_end()
.alias("reported_period_1Q")
)

df.with_columns(
pl.coalesce(
(
pl.col("variable")
.get(pl.col("reported_period_1Q").index_of(date)
.over("id")).over("id")
.alias(f"rp_{date}_{id}")
for id, date in zip(
list(df.get_column("id")), list(df.get_column("reported_period_1Q"))
)
)
)
)

Continue reading...