如何在各分区的滚动窗口中获取最大值对应的日期

嘿，我来帮你搞定这个问题！针对你给出的这份Polars DataFrame，要按category分区，在滚动窗口里拿到value最大值对应的日期，咱们一步步来：

首先先确认你的示例DataFrame生成代码是这样的：

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)]
})

接下来分两种常见的滚动窗口场景给你解决方案：

1. 基于时间的滚动窗口（比如7天窗口） #

如果你的滚动窗口是按时间范围来划分的（比如最近7天），咱们可以用Polars的group_by结合rolling方法，配合arg_max定位最大值对应的日期：

# 先按日期排序，确保每个分区内的时间序列有序
result = df.sort("date").group_by("category").rolling(
    index_column="date",  # 指定时间索引列
    period="7d",          # 7天的滚动窗口范围
    closed="left"         # 窗口左闭右开，可根据需求调整为"both"等
).agg(
    # 找到窗口内value最大的那一行对应的日期
    max_value_date=pl.col("date").arg_max(pl.col("value")).first(),
    max_value=pl.col("value").max()  # 顺便输出最大值，方便验证结果
)

print(result)

这段代码里的关键点：

• sort("date")不能少，滚动窗口依赖有序的时间序列，不然结果会混乱
• arg_max(pl.col("value"))会定位到窗口内value最大的行对应的date，用.first()直接提取出目标日期
• 你可以随意调整period参数，比如改成"3d"就是3天窗口，完全适配你的业务需求

2. 基于行数的滚动窗口（比如前7行） #

如果你的滚动窗口是按行数定义的（比如每个窗口包含当前行及前6行，共7行），只需要把rolling里的period换成by参数就行：

result_by_rows = df.sort("date").group_by("category").rolling(
    index_column="date",
    by=7,                # 按7行划分滚动窗口
    closed="left"
).agg(
    max_value_date=pl.col("date").arg_max(pl.col("value")).first(),
    max_value=pl.col("value").max()
)

print(result_by_rows)

这样处理后，你就能轻松拿到每个分区在每个滚动窗口内，value最大值对应的日期啦，要是有其他窗口参数调整的需求，直接修改rolling里的参数就好~

备注：内容来源于stack exchange，提问作者Jonathan