[Python] Find corresponding date of max value in a rolling window of each partition

Sample code:

import polars as pl

from datetime import date
from random import randint

df = pl.DataFrame({
"category": [cat for cat in ["A", "B"] for _ in range(1, 32)],
"date": [date(2025, 1, i) for _ in ["A", "B"] for i in range(1, 32)],
"value": [randint(1, 50) for _ in ["A", "B"] for _ in range(1, 32)]
})


I know how to find the maximum value in a rolling window in each partition by following:

df.with_columns(
value_max=pl.col("value").rolling_max_by("date", window_size="5d", closed="both").over(pl.col("category"))
)


shape: (62, 4)
┌──────────┬────────────┬───────┬───────────┐
│ category ┆ date ┆ value ┆ value_max │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ date ┆ i64 ┆ i64 │
╞══════════╪════════════╪═══════╪═══════════╡
│ A ┆ 2025-01-01 ┆ 42 ┆ 42 │
│ A ┆ 2025-01-02 ┆ 18 ┆ 42 │
│ A ┆ 2025-01-03 ┆ 33 ┆ 42 │
│ A ┆ 2025-01-04 ┆ 35 ┆ 42 │
│ A ┆ 2025-01-05 ┆ 46 ┆ 46 │
│ … ┆ … ┆ … ┆ … │
│ B ┆ 2025-01-27 ┆ 49 ┆ 49 │
│ B ┆ 2025-01-28 ┆ 22 ┆ 49 │
│ B ┆ 2025-01-29 ┆ 49 ┆ 49 │
│ B ┆ 2025-01-30 ┆ 32 ┆ 49 │
│ B ┆ 2025-01-31 ┆ 25 ┆ 49 │
└──────────┴────────────┴───────┴───────────┘


I also want to create another column and find the corresponding date of the value_max. For example, for category A, the value of this new column in date 2025-01-01 / 2025-01-02 / 2025-01-03 / 2025-01-04 is 2025-01-01, and 2025-01-05 when date 2025-01-05.

I currently use a self-join with condition to get the date. Is there another way to do that without using joining? Also, is the new logic achievable by using SQL?

Continue reading...