如何在Polars中基于滚动窗口期计算分组Z-score？

希望基于滚动回溯期的数据计算分箱的Z-score。例如：当天[9:30-9:35)时段的访客量，需基于过去3天同一[9:30-9:35)时段访客数据的均值、标准差进行Z-score归一化。

当前用Polars的两种尝试均触发InvalidOperationError，求可行的计算方法。

尝试代码如下：

import polars as pl
import pandas as pd

def z_score(col: str, over: str, alias: str):
    # 计算按`over`分组的`col`的Z-score
    return (
        (pl.col(col)-pl.col(col).over(over).mean()) / pl.col(col).over(over).std()
    ).alias(alias)

df = pl.from_dict(
    {
        "timestamp": pd.date_range("2019-12-02 9:30", "2019-12-02 12:30", freq="30s").union(
            pd.date_range("2019-12-03 9:30", "2019-12-03 12:30", freq="30s")
        ),
        "visitors": [(e % 2) + 1 for e in range(722)]
    }
# 将时间截断为5分钟分箱，比如[9:30-9:35)转为930
).with_columns(
    pl.col("timestamp").dt.truncate(every="5m").dt.to_string("%H%M").cast(pl.Int32).alias("five_minute_bin")
).with_columns(
    pl.col("timestamp").dt.truncate(every="3d").alias("daytrunc")
)

# 尝试1：非滚动但报错（InvalidOperationError: window expression not allowed in aggregation）
# df.with_columns(
#     z_score("visitors", "five_minute_bin", "normalized").over("daytrunc")
# )

# 尝试2：滚动方式也报错（InvalidOperationError: window expression not allowed in aggregation）
# df.rolling(index_column="daytrunc", period="3i").agg(z_score("visitors", "five_minute_bin", "normalized"))

输入输出示例 #

含4天数据，每天有4个数据点分属两个时间分箱({0,0} - {0,1}), ({1,0} - {1,1})

输入：

Day 0: x_d0_{0,0}, x_d0_{0,1}, x_d0_{1,0}, x_d0_{1,1}
Day 1: x_d1_{0,0}, x_d1_{0,1}, x_d1_{1,0}, x_d1_{1,1}
Day 2: x_d2_{0,0}, x_d2_{0,1}, x_d2_{1,0}, x_d2_{1,1}
Day 3: x_d3_{0,0}, x_d3_{0,1}, x_d3_{1,0}, x_d3_{1,1}

输出：

Day 0: norm_x_d0_{0,0} = nan, norm_x_d0_{0,1} = nan, norm_x_d0_{1,0} = nan, norm_x_d0_{1,1} = nan
Day 1: norm_x_d1_{0,0} = nan, norm_x_d1_{0,1} = nan, norm_x_d1_{1,0} = nan, norm_x_d1_{1,1} = nan
Day 2: norm_x_d2_{0,0} = nan, norm_x_d2_{0,1} = nan, norm_x_d2_{1,0} = nan, norm_x_d2_{1,1} = nan
Day 3: norm_x_d3_{0,0} = (x_d3_{0,0} - np.mean([x_d0_{0,0}, x_d0_{0,1}, x_d1_{0,0}, ..., x_d3_{0,1}])) / np.std([x_d0_{0,0}, x_d0_{0,1}, x_d1_{0,0}, ..., x_d3_{0,1}]) , ... , 

之前的报错核心原因是：Polars不允许在聚合或窗口函数中嵌套另一个窗口表达式，比如你在over("daytrunc")里又嵌套了over("five_minute_bin")，这种嵌套逻辑会触发InvalidOperationError。

要实现需求，需要拆解步骤：先按分箱+日期聚合数据，再对每个分箱计算滚动统计量，最后合并回原数据计算Z-score。

完整实现代码 #

import polars as pl
import pandas as pd

# 构造包含4天的原始数据（补全示例中的4天数据）
df = pl.from_dict(
    {
        "timestamp": pd.date_range("2019-12-02 9:30", "2019-12-02 12:30", freq="30s").union(
            pd.date_range("2019-12-03 9:30", "2019-12-03 12:30", freq="30s")
        ).union(
            pd.date_range("2019-12-04 9:30", "2019-12-04 12:30", freq="30s")
        ).union(
            pd.date_range("2019-12-05 9:30", "2019-12-05 12:30", freq="30s")
        ),
        "visitors": [(e % 2) + 1 for e in range(722*4)]
    }
).with_columns(
    # 生成5分钟分箱标识
    pl.col("timestamp").dt.truncate(every="5m").dt.to_string("%H%M").cast(pl.Int32).alias("five_minute_bin"),
    # 提取日期维度，方便按天分组
    pl.col("timestamp").dt.date().alias("date")
)

# 步骤1：按日期+分箱聚合，得到每个分箱每天的访客均值（如果每个分箱每天只有1条数据，可跳过此步）
daily_bin_data = df.group_by("date", "five_minute_bin").agg(
    pl.col("visitors").mean().alias("daily_visitors")
)

# 步骤2：对每个分箱计算滚动3天的均值和标准差（closed="left"表示不包含当前天，符合示例需求）
rolling_stats = daily_bin_data.group_by("five_minute_bin", maintain_order=True).agg(
    pl.col("date"),
    pl.col("daily_visitors").rolling(window_size=3, closed="left").mean().alias("rolling_mean"),
    pl.col("daily_visitors").rolling(window_size=3, closed="left").std(ddof=0).alias("rolling_std"),
).explode(["date", "daily_visitors", "rolling_mean", "rolling_std"])

# 步骤3：合并滚动统计数据到原表，计算Z-score
result = df.join(
    rolling_stats.select("date", "five_minute_bin", "rolling_mean", "rolling_std"),
    on=["date", "five_minute_bin"],
    how="left"
).with_columns(
    ((pl.col("visitors") - pl.col("rolling_mean")) / pl.col("rolling_std")).alias("normalized_visitors")
)

# 查看结果示例（2019-12-05的930分箱数据）
print(result.filter((pl.col("date") == pl.date(2019,12,5)) & (pl.col("five_minute_bin") == 930)).head())

关键说明 #

1. 滚动窗口参数：closed="left"确保当前天的数据不纳入滚动计算，和你示例中Day3用Day0-Day2数据的逻辑一致；如果需要包含当前天，改为closed="both"。
2. 标准差计算：ddof=0计算的是总体标准差，若需要样本标准差（除以n-1），去掉该参数即可。
3. 数据聚合：如果你的原始数据中每个five_minute_bin每天只有1条记录，可以直接跳过group_by("date", "five_minute_bin").agg()步骤，直接对原数据的日期和分箱分组做滚动计算。

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