不使用with_row_index(),如何通过单次with_columns()调用实现Polars指定窗口的Rolling滚动聚合
不使用with_row_index(),如何通过单次with_columns()调用实现Polars指定窗口的Rolling滚动聚合
好问题!其实Polars早就支持直接在rolling()里用表达式生成临时索引,完全不用额外创建和删除索引列——咱们只需要在index_column参数里直接传入一个动态生成的整数序列表达式就行,一步到位完成需求。
直接上代码:
import polars as pl df = pl.DataFrame({"x": [1.2, 1.3, 3.4, 3.5]}) # 单次with_columns调用完成滚动聚合 result = df.with_columns( y=pl.col("x").rolling( index_column=pl.int_range(0, pl.count()), period="4i", offset="-3i" ) ) print(result)
这段代码的核心逻辑:
pl.int_range(0, pl.count())会动态生成一个从0到DataFrame总行数-1的整数序列,相当于临时的行索引,直接传给index_column参数就行,不需要把它作为实际列加到DataFrame里。- 你之前用到的
period="4i"和offset="-3i"参数完全保留,逻辑和原来一致:period定义窗口覆盖4个整数索引单位,offset调整窗口的位置,刚好对应你需要的[-2:1]窗口范围。
运行后得到的结果和你期望的完全一致:
shape: (4, 2) ┌─────┬───────────────────┐ │ x ┆ y │ │ --- ┆ --- │ │ f64 ┆ list[f64] │ ╞═════╪═══════════════════╡ │ 1.2 ┆ [1.2, 1.3] │ │ 1.3 ┆ [1.2, 1.3, 3.4] │ │ 3.4 ┆ [1.2, 1.3, 3.4, 3.5] │ │ 3.5 ┆ [1.3, 3.4, 3.5] │ └─────┴───────────────────┘
如果之后需要调整窗口范围,只需要修改period和offset的参数值就好,全程不用碰索引列的创建和删除,代码更简洁也更易维护。
备注:内容来源于stack exchange,提问作者bretauv
相关产品推荐
相关产品推荐

