Polars中滚动窗口内按多维度分组计算均值的实现方法

数据集说明 #

每个Node每小时对应一行数据，BidType为计算列：

• Price > 5 时为 Buy
• Price < 2 时为 Sell
• 其余情况为 None

示例数据集：

需求 #

按Node、Hour、BidType分组，对Price计算过去100天内的滚动均值，要求覆盖所有Node、Hour、BidType的组合（即每个日期每个Node需返回3×24行，对应3种BidType和24小时）。

之前尝试的方法及问题 #

方法1 #

.rolling('date', period='100d', group_by=['Node', 'Hour', 'BidType']).agg(pl.col('Price').mean())

问题：仅返回原始数据中存在的分组组合，每个日期每个Node仅24行，无法覆盖3种BidType的全组合。

方法2 #

.rolling('date', period='10d', group_by=['Node', 'Hour'])
    .agg(
        pl.col('Price').filter(pl.col('BidType') == 'Buy').mean().alias('Price_buy'),
        pl.col('Price').filter(pl.col('BidType') == 'Sell').mean().alias('Price_sell'),
        pl.col('Price').filter(pl.col('BidType') == 'None').mean().alias('Price_none')
    )
    .unpivot(index=['Node', 'Hour', 'date'])

问题：未生成预期的全分组覆盖结果。

正确实现方式 #

核心思路：先构建所有可能的Node、Hour、BidType、date笛卡尔积，再左连接原始数据，最后基于全组合表做滚动均值计算，确保无分组缺失。

步骤1：修正后的测试数据集生成代码 #

import pandas as pd
import numpy as np
import polars as pl

# 生成10个Node、100天、每天24小时的测试数据
df = pd.DataFrame(np.random.randint(0,10,size=(10*100*24, 1)), columns=['Price'])
df['Node'] = np.repeat(np.arange(0,10), 100*24)
df['date'] = np.tile(pd.date_range('2020-01-01', periods=100, freq='D'), 10*24)
df['Hour'] = np.tile(np.arange(0,24), 10*100)

# 转换为Polars DataFrame并生成BidType列
df = pl.DataFrame(df)
df = df.sort(['Node', 'date', 'Hour'])
df = df.with_columns(
    pl.when(pl.col('Price') >=5).then(pl.lit('Buy'))
       .when(pl.col('Price') <2).then(pl.lit('Sell'))
       .otherwise(pl.lit('None'))
       .alias('BidType')
)

步骤2：核心计算代码 #

# 提取所有维度的唯一值
unique_nodes = df['Node'].unique()
unique_hours = df['Hour'].unique()
unique_bidtypes = pl.Series(['Buy', 'Sell', 'None'])
unique_dates = df['date'].unique()

# 构建全组合基准表
full_combinations = pl.cartesian_product([
    unique_nodes.alias('Node'),
    unique_hours.alias('Hour'),
    unique_bidtypes.alias('BidType'),
    unique_dates.alias('date')
]).sort(['Node', 'date', 'Hour', 'BidType'])

# 左连接原始数据，保留所有组合
full_df = full_combinations.join(
    df.select(['Node', 'date', 'Hour', 'BidType', 'Price']),
    on=['Node', 'date', 'Hour', 'BidType'],
    how='left'
)

# 按分组计算过去100天的滚动均值
result = full_df.sort(['Node', 'BidType', 'Hour', 'date']).rolling(
    index_column='date',
    period='100d',
    group_by=['Node', 'Hour', 'BidType']
).agg(
    pl.col('Price').mean().alias('Rolling_100d_Price_Mean')
)

# 验证结果：每个Node每天有24*3行数据
print(result.head(24*3))

方法说明 #

1. 笛卡尔积全组合：确保所有Node、Hour、BidType、date的组合都被包含，不会因原始数据缺失某分组而遗漏。
2. 左连接填充：原始数据中无记录的分组会填充为null，Polars的mean函数会自动忽略null值，不影响滚动计算。
3. 滚动均值计算：基于全组合表分组计算，保证每个日期每个Node都有3×24行结果，完全覆盖需求。

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