Polars滚动函数在NaN与Null共存时的异常行为问询

环境信息 #

• Polars 1.31.0
• NumPy 2.3.1
• Python 3.12.11

问题复现 #

正常场景（仅含NaN） #

当数据仅包含np.nan时，rolling_sum行为符合预期：

import polars as pl
import numpy as np

data_dict_1 = {"x":[1., 1., 1., np.nan, 1., 1., 1., 1., 1.]}
data_1 = pl.DataFrame(data_dict_1)
with pl.Config(tbl_rows=20):
    print(data_1.with_columns(pl.col("x").rolling_sum(2).alias("rolling")))

输出：

shape: (9, 2)
┌─────┬─────────┐
│ x   ┆ rolling │
│ --- ┆ ---     │
│ f64 ┆ f64     │
╞═════╪═════════╡
│ 1.0 ┆ null    │
│ 1.0 ┆ 2.0     │
│ 1.0 ┆ 2.0     │
│ NaN ┆ NaN     │
│ 1.0 ┆ NaN     │
│ 1.0 ┆ 2.0     │
│ 1.0 ┆ 2.0     │
│ 1.0 ┆ 2.0     │
│ 1.0 ┆ 2.0     │
└─────┴─────────┘

预期：索引0为null，索引3、4为NaN，其余正常求和，结果符合预期。

异常场景（同时含NaN与null） #

当数据同时包含np.nan和None（Polars解析为null）时，rolling_sum出现异常的NaN传播：

data_dict_2 = {"x":[1., 1., 1., np.nan, 1., 1., 1., 1., 1., None, 1., 1., 1.]}
data_2 = pl.DataFrame(data_dict_2)
with pl.Config(tbl_rows=20):
    print(data_2.with_columns(pl.col("x").rolling_sum(2).alias("rolling")))

输出：

shape: (13, 2)
┌──────┬─────────┐
│ x    ┆ rolling │
│ ---  ┆ ---     │
│ f64  ┆ f64     │
╞══════╪═════════╡
│ 1.0  ┆ null    │
│ 1.0  ┆ 2.0     │
│ 1.0  ┆ 2.0     │
│ NaN  ┆ NaN     │
│ 1.0  ┆ NaN     │
│ 1.0  ┆ 2.0     │
│ 1.0  ┆ NaN     │
│ 1.0  ┆ NaN     │
│ 1.0  ┆ NaN     │
│ null ┆ null    │
│ 1.0  ┆ null    │
│ 1.0  ┆ NaN     │
│ 1.0  ┆ NaN     │
└──────┴─────────┘

预期：索引0为null，索引3、4为NaN，索引9、10为null，其余正常求和。但实际在索引6开始出现无理由的NaN，且后续非null/NaN的行也持续输出NaN。

对比：使用rolling_map(sum)的正确结果 #

用rolling_map结合Python内置sum函数可得到符合预期的结果：

with pl.Config(tbl_rows=20):
    print(data_2.with_columns(pl.col("x").rolling_map(sum, 2).alias("rolling")))

输出：

shape: (13, 2)
┌──────┬─────────┐
│ x    ┆ rolling │
│ ---  ┆ ---     │
│ f64  ┆ f64     │
╞══════╪═════════╡
│ 1.0  ┆ null    │
│ 1.0  ┆ 2.0     │
│ 1.0  ┆ 2.0     │
│ NaN  ┆ NaN     │
│ 1.0  ┆ NaN     │
│ 1.0  ┆ 2.0     │
│ 1.0  ┆ 2.0     │
│ 1.0  ┆ 2.0     │
│ 1.0  ┆ 2.0     │
│ null ┆ null    │
│ 1.0  ┆ null    │
│ 1.0  ┆ 2.0     │
│ 1.0  ┆ 2.0     │
└──────┴─────────┘

原因分析 #

Polars内置滚动函数（如rolling_sum）的底层实现对NaN（浮点特殊值）和null（Polars缺失值）的状态跟踪存在逻辑错误。当数据中同时存在两者时，内置函数在处理完null后未能正确重置NaN的传播状态，导致后续正常数值的窗口求和错误地输出NaN。而rolling_map使用Python的sum函数，会正确区分null（被视为缺失值，求和时忽略）和NaN（求和时直接返回NaN），因此行为符合预期。

临时解决方案 #

1. 统一缺失值类型：将数据中的np.nan转换为Polars的null，避免混合类型：

   data_2 = data_2.with_columns(pl.col("x").cast(pl.Float64).replace(np.nan, None))
   

   若原复杂数据集转换后仍有问题，需检查是否存在其他隐性的混合缺失值场景。
2. 使用rolling_map替代：虽然性能低于内置函数，但能保证结果正确性，适合对精度要求优先的场景。
3. 升级Polars版本：该问题大概率是版本bug，后续版本可能已修复，可尝试升级至最新稳定版测试。

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