Баг rolling_sum в Polars при смешении NaN и null: MRE, причины и фикс

Разбираем баг Polars 1.31.0: rolling_sum возвращает лишние NaN при смешении NaN и null. MRE, сравнение с rolling_map, ссылка на фикс в следующем релизе.

Аналитика со скользящим окном — привычный инструмент в дата-пайплайнах, но при появлении пропусков поведение может оказаться неожиданным. Если вы используете Polars 1.31.0 с Python 3.12.11 и NumPy 2.3.1, то при смешении в одном столбце значений NaN (np.nan) и null (None) rolling_sum способен вернуть нетипичные результаты. Ниже — минимальный пример воспроизведения, пояснение причины и ссылка на исправление, которое войдёт в следующий релиз.

Ожидаемое поведение при наличии только NaN

Начнём со столбца с одним NaN и остальными корректными числами с плавающей точкой. При размере окна 2 первая позиция — null (недостаточно данных), два окна, включающие NaN, дают NaN, а прочие окна суммируются до 2.0, как и ожидается.

import polars as pl
import numpy as np

src_one = {"x": [1., 1., 1., np.nan, 1., 1., 1., 1., 1.]}
df_one = pl.DataFrame(src_one)

with pl.Config(tbl_rows=20):
    print(
        df_one.with_columns(
            pl.col("x").rolling_sum(2).alias("roll_val")
        )
    )

shape: (9, 2)
┌─────┬──────────┐
│ x   ┆ roll_val │
│ --- ┆ ---      │
│ 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      │
└─────┴──────────┘

Когда NaN и null встречаются вместе

Добавим к NaN ещё и null. Интуитивно хочется увидеть прежний рисунок вокруг NaN и, вдобавок, разрывы null в окнах, содержащих null. Однако rolling_sum после первого появления начинает распространять NaN дальше, чем ожидается.

src_two = {"x": [1., 1., 1., np.nan, 1., 1., 1., 1., 1., None, 1., 1., 1.]}
df_two = pl.DataFrame(src_two)

with pl.Config(tbl_rows=20):
    print(
        df_two.with_columns(
            pl.col("x").rolling_sum(2).alias("roll_val")
        )
    )

shape: (13, 2)
┌──────┬──────────┐
│ x    ┆ roll_val │
│ ---  ┆ ---      │
│ 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      │
└──────┴──────────┘

После NaN появляется только одно нормальное 2.0; последующие окна неожиданно возвращают NaN, хотя содержат конечные значения. Это расходится с обычной семантикой, наблюдаемой при входе только с NaN.

Что происходит

Описанное поведение — ошибка в rolling_sum при совместном присутствии NaN и null в одном столбце. Она уже исправлена в main и будет включена в следующий выпуск Polars. Изменение отслеживается здесь: https://github.com/pola-rs/polars/pull/23482.

Базовая проверка корректности с помощью rolling_map

Если посчитать ту же скользящую сумму через rolling_map(sum, 2), результат совпадает с интуитивными ожиданиями и для окон с NaN, и для окон с null. Поэтому это удобная точка отсчёта для валидации.

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

shape: (13, 2)
┌──────┬──────────┐
│ x    ┆ roll_val │
│ ---  ┆ ---      │
│ 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      │
└──────┴──────────┘

Но rolling_map выполняет Python UDF и материализует объекты Series, что добавляет существенные накладные расходы. Документация прямо предупреждает не применять его в продакшене без крайней необходимости.

Вычисление пользовательских функций крайне медленно. По возможности используйте специализированные скользящие функции, такие как Expr.rolling_sum().

Исправление и корректный результат

После слияния фикса в main (см. PR выше) rolling_sum снова выдаёт ожидаемые результаты для сочетаний NaN и null. Повторная оценка того же выражения возвращает правильные суммы и локальные эффекты NaN/null.

with pl.Config(tbl_rows=20):
    print(
        df_two.with_columns(
            pl.col("x").rolling_sum(2).alias("roll_val")
        )
    )

┌──────┬──────────┐
│ x    ┆ roll_val │
│ ---  ┆ ---      │
│ 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      │
└──────┴──────────┘

Почему это важно

Скользящие агрегаты лежат в основе поиска аномалий, KPI и признаков временных рядов. Незаметные расхождения, возникающие при совместном появлении NaN и null, могут исказить метрики или сделать входы моделей некорректными. Понимание этого краевого случая и предстоящего исправления помогает выбрать верный подход — либо проверять корректность через rolling_map, либо полагаться на оптимизированный rolling_sum после выхода фикса.

Практические выводы

Если вы работаете на Polars 1.31.0 и смешиваете NaN с null в скользящих окнах, учитывайте показанное выше неожиданное распространение NaN. Для проверки корректности rolling_map(sum, 2) воспроизводит задуманное поведение, но из-за выполнения Python UDF и материализации Series влечёт серьёзные накладные расходы, о чём предупреждает документация. Проблема уже решена и войдёт в следующий релиз, трекинг: https://github.com/pola-rs/polars/pull/23482. Соотнесите выбор между производительностью и немедленной корректностью с тем, насколько вы готовы перейти на версию с исправлением.