基于Polars实现按午夜截断的滚动分组均值计算

原始数据构建 #

import polars as pl
import numpy as np
from datetime import datetime

df = pl.DataFrame({'ts': pl.datetime_range(datetime(2020, 1, 1), datetime(2020, 1, 10), '1h', eager=True)})
df = df.with_columns(value=pl.Series(np.arange(len(df))))

DataFrame预览：

shape: (217, 2)
┌─────────────────────┬───────┐
│ ts                  ┆ value │
│ ---                 ┆ ---   │
│ datetime[μs]        ┆ i64   │
╞═════════════════════╪═══════╡
│ 2020-01-01 00:00:00 ┆ 0     │
│ 2020-01-01 01:00:00 ┆ 1     │
│ 2020-01-01 02:00:00 ┆ 2     │
│ 2020-01-01 03:00:00 ┆ 3     │
│ …                   ┆ …     │
│ 2020-01-09 21:00:00 ┆ 213   │
│ 2020-01-09 22:00:00 ┆ 214   │
│ 2020-01-09 23:00:00 ┆ 215   │
│ 2020-01-10 00:00:00 ┆ 216   │
└─────────────────────┴───────┘

需求说明 #

对每一行数据，计算该行时间的3天前午夜至该行当天午夜区间内所有行的value列均值。示例：

• 针对行2020-01-09 23:00:00，需计算ts ≥ 2020-01-06 00:00:00 且 ts < 2020-01-09 00:00:00的行的value均值。

预期输出 #

shape: (217, 2)
┌─────────────────────┬───────┐
│ ts                  ┆ value │
│ ---                 ┆ ---   │
│ datetime[μs]        ┆ f64   │
╞═════════════════════╪═══════╡
│ 2020-01-01 00:00:00 ┆ null  │
│ 2020-01-01 01:00:00 ┆ null  │
│ 2020-01-01 02:00:00 ┆ null  │
│ 2020-01-01 03:00:00 ┆ null  │
│ …                   ┆ …     │
│ 2020-01-09 21:00:00 ┆ 155.5 │
│ 2020-01-09 22:00:00 ┆ 155.5 │
│ 2020-01-09 23:00:00 ┆ 155.5 │
│ 2020-01-10 00:00:00 ┆ 179.5 │
└─────────────────────┴───────┘

对应预期计算示例：

# 2020-01-09 23:00:00行的均值计算逻辑
df.filter(
    (pl.col("ts") >= datetime(2020, 1, 6)) & (pl.col("ts") < datetime(2020, 1, 9))
)["value"].mean()

# 2020-01-10 00:00:00行的均值计算逻辑
df.filter(
    (pl.col("ts") >= datetime(2020, 1, 7)) & (pl.col("ts") < datetime(2020, 1, 10))
)["value"].mean()

解决方案 #

方法1：区间匹配+分组聚合 #

该方法先计算每行的区间边界，再通过半连接聚合区间内的均值，最后匹配回原表，效率较高：

# 计算每行的区间边界：3天前午夜、当天午夜
df = df.with_columns(
    current_midnight=pl.col("ts").dt.date().cast(pl.Datetime),
    three_days_ago_midnight=pl.col("ts").dt.date().cast(pl.Datetime) - pl.duration(days=3)
)

# 聚合每个区间的均值
interval_means = df.select("three_days_ago_midnight", "current_midnight").unique().join(
    df,
    how="left",
    condition=pl.col("ts").is_between(pl.col("three_days_ago_midnight"), pl.col("current_midnight"), closed="left")
).group_by("three_days_ago_midnight", "current_midnight").agg(
    mean_value=pl.col("value").mean()
)

# 合并结果到原表
result = df.join(interval_means, on=["three_days_ago_midnight", "current_midnight"]).select(
    "ts", "mean_value"
).rename({"mean_value": "value"})

print(result)

方法2：滚动窗口函数 #

利用Polars的滚动窗口功能，结合时间范围过滤，代码更简洁：

result = df.with_columns(
    value=pl.col("value").rolling(
        offset=-pl.duration(days=3),
        window_size=pl.duration(days=3),
        by="ts",
        closed="left",
        filter=pl.col("ts") < pl.col("ts").dt.date().cast(pl.Datetime)
    ).mean()
)

print(result)

内容的提问来源于stack exchange，提问作者ignoring_gravity