如何用Python Polars实现整数滚动窗口GroupBy及自定义聚合

1. 基于整数列的滚动窗口配置 #

Polars的rolling_apply支持直接基于整数列定义窗口，无需依赖datetime。首先确保你的整数列（如i）是有序的：

import polars as pl
df = df.sort("i")

• 数值范围窗口：如果窗口是基于i的数值区间（比如包含i ∈ [当前i-窗口大小, 当前i]的行），直接指定by="i"：

  df = df.with_columns(
      pl.col("data").rolling_apply(
          window_size=10,  # 窗口的数值范围，比如i的差值不超过10
          function=your_custom_func,
          by="i"
      ).alias("rolling_result")
  )
  

• 固定行数窗口：如果需要固定行数的回溯窗口（比如取当前行之前的4行+当前行，共5行），用行索引作为by参数：

  df = df.with_columns(
      pl.col("data").rolling_apply(
          window_size=5,  # 固定5行窗口
          function=your_custom_func,
          by=pl.int_range(0, pl.count())  # 用行索引替代整数列
      ).alias("fixed_row_result")
  )
  

2. 应用自定义聚合函数 #

针对你的f_jparams这类自定义函数，直接传入rolling_apply的function参数即可。注意函数需要接收一个Polars Series并返回标量：

def f_jparams(window_series: pl.Series) -> float:
    # 自定义聚合逻辑，比如加权移动平均、特殊过滤后的均值等
    filtered = window_series.filter(window_series > window_series.mean())
    return filtered.mean()

# 应用到单列
df = df.with_columns(
    pl.col("price").rolling_apply(
        window_size=5,
        function=f_jparams,
        by="i"
    ).alias("custom_ma")
)

# 应用到多列（打包成结构体处理）
def multi_col_custom(window_structs: pl.Series) -> float:
    window_df = window_structs.to_frame()
    # 基于多列的自定义逻辑，比如用volume加权price
    return (window_df["price"] * window_df["volume"]).sum() / window_df["volume"].sum()

df = df.with_columns(
    pl.struct(["price", "volume"]).rolling_apply(
        window_size=5,
        function=multi_col_custom,
        by="i"
    ).alias("volume_weighted_ma")
)

3. 简化列选择（无需手动遍历） #

Polars原生支持通过exclude或列模式匹配批量选择列，完全替代手动遍历列表：

• 排除指定列：选择除i和timestamp外的所有列：

  df = df.with_columns(
      pl.col(exclude=["i", "timestamp"]).rolling_apply(
          window_size=5,
          function=f_jparams,
          by="i"
      ).suffix("_rolling")
  )
  

• 按前缀/后缀选择：选择所有以data_开头的列：

  df = df.with_columns(
      pl.col("data_*").rolling_apply(
          window_size=5,
          function=f_jparams,
          by="i"
      ).suffix("_ma")
  )
  

4. 性能优化：优先用Polars原生表达式 #

如果自定义逻辑可以用Polars内置表达式组合实现，绝对不要用Python自定义函数——原生表达式的性能是Python函数的几十倍。比如把f_jparams的过滤+均值逻辑改成原生表达式：

df = df.with_columns(
    pl.col("price")
    .rolling_mean(window_size=5, by="i")
    .alias("normal_ma"),
    # 过滤掉窗口内低于均值的数再取平均
    (pl.col("price")
     .filter(pl.col("price") > pl.col("price").rolling_mean(window_size=5, by="i"))
     .rolling_mean(window_size=5, by="i"))
    .alias("filtered_ma")
)

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