Polars滚动窗口中Slice切片权重序列的报错问题及解决咨询

问题背景 #

需求为基于滚动窗口长度对权重序列进行切片：窗口最大长度为5，前4行窗口长度分别为1-4，需取权重序列[0.1,0.2,0.3,0.4,0.5]的后N个元素（如第1行取[0.5]，第2行取[0.4,0.5]，第4行取[0.2,0.3,0.4,0.5]）；从第5行开始，窗口长度固定为5，权重取完整序列。

最初实现代码针对5行数据能运行，但结果不符合预期（取了权重序列的前N个而非后N个）；扩展到6行时触发索引越界错误：

初始错误代码 #

pl.DataFrame({
    't_idx': [1 ,2 ,3, 4, 5],
}).rolling('t_idx', period='5i').agg(
    -pl.col('t_idx').len().cast(pl.Int64).alias('start'),
    pl.col('t_idx').count().alias('end'),
    pl.lit(pl.Series([.1, .2, .3, .4, .5])).slice(-pl.col('t_idx').len().cast(pl.Int64), pl.col('t_idx').count()).alias('weights')
)

错误结果（5行数据） #

shape: (5, 4)
┌───────┬───────┬─────┬───────────────────┐
│ t_idx ┆ start ┆ end ┆ new               │
│ ---   ┆ ---   ┆ --- ┆ ---               │
│ i64   ┆ i64   ┆ u32 ┆ list[f64]         │
╞═══════╪═══════╪═════╪═══════════════════╡
│ 1     ┆ -1    ┆ 1   ┆ [0.1]             │
│ 2     ┆ -2    ┆ 2   ┆ [0.1, 0.2]        │
│ 3     ┆ -3    ┆ 3   ┆ [0.1, 0.2, 0.3]   │
│ 4     ┆ -4    ┆ 4   ┆ [0.1, 0.2, … 0.4] │
│ 5     ┆ -5    ┆ 5   ┆ [0.1, 0.2, … 0.5] │
└───────┴───────┴─────┴───────────────────┘

6行数据报错信息 #

PanicException: range end index 6 out of range for slice of length 5

单独Slice操作对比（无报错） #

pl.select(
    pl.lit(pl.Series([.1, .2, .3, .4, .5])).slice(-5, 6)
)

返回：

shape: (5, 1)
┌─────┐
│     │
│ --- │
│ f64 │
╞═════╡
│ 0.1 │
│ 0.2 │
│ 0.3 │
│ 0.4 │
│ 0.5 │
└─────┘

报错原因 #

1. 窗口范围定义错误：使用period='5i'会将t_idx视为时间值，窗口范围为当前t_idx减去5个单位到当前值。第6行t_idx=6时，窗口包含t_idx=1到6共6个元素，导致pl.col('t_idx').len()返回6。
2. 滚动Agg与单独Slice行为差异：单独执行slice时，Polars会自动截断超出序列长度的索引/参数；但在滚动窗口的Agg上下文里，动态参数的slice会严格计算索引，当起始索引-6超出权重序列的有效负索引范围（最小为-5）时，触发越界panic。
3. 权重切片方向错误：初始代码取权重序列的前N个元素，与需求的后N个完全相反。

正确实现方法 #

方案一：基于行数的滚动窗口（推荐） #

使用window_size定义固定行数的滚动窗口，直接计算权重序列的后N个元素：

import polars as pl

weights_full = pl.Series([0.1, 0.2, 0.3, 0.4, 0.5])
window_size = 5

df = pl.DataFrame({'t_idx': [1, 2, 3, 4, 5, 6]})

result = df.rolling('t_idx', window_size=window_size).agg(
    window_length=pl.col('t_idx').count(),
    weights=pl.lit(weights_full).slice(
        pl.lit(len(weights_full)) - pl.col('t_idx').count(),
        pl.col('t_idx').count()
    )
)

print(result)

方案二：无滚动窗口的直接计算 #

通过计算每行对应的窗口长度，直接生成权重列表：

import polars as pl

weights_full = [0.1, 0.2, 0.3, 0.4, 0.5]
window_size = 5

df = pl.DataFrame({'t_idx': [1, 2, 3, 4, 5, 6]})

result = df.with_columns(
    window_length=pl.int_range(0, pl.count()).add(1).min(window_size),
    weights=pl.int_range(0, pl.count()).add(1).min(window_size)
            .map_elements(lambda x: weights_full[-x:], return_dtype=pl.List(pl.Float64))
)

print(result)

预期结果 #

两种方案均返回符合需求的结果：

shape: (6, 3)
┌───────┬──────────────┬───────────────────┐
│ t_idx ┆ window_length ┆ weights           │
│ ---   ┆ ---          ┆ ---               │
│ i64   ┆ u32          ┆ list[f64]         │
╞═══════╪══════════════╪═══════════════════╡
│ 1     ┆ 1            ┆ [0.5]             │
│ 2     ┆ 2            ┆ [0.4, 0.5]        │
│ 3     ┆ 3            ┆ [0.3, 0.4, 0.5]   │
│ 4     ┆ 4            ┆ [0.2, 0.3, … 0.5] │
│ 5     ┆ 5            ┆ [0.1, 0.2, … 0.5] │
│ 6     ┆ 5            ┆ [0.1, 0.2, … 0.5] │
└───────┴──────────────┴───────────────────┘

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