Кумулятивное среднее и стандартное отклонение по группам в Polars

Кумулятивное среднее и стандартное отклонение по группам в Polars: используем rolling_mean и rolling_std, а также короткую формулу cum_sum/cum_count. Подробно.

Вычисление кумулятивного среднего и кумулятивного стандартного отклонения в Polars DataFrame кажется делом простым, но самый очевидный подход быстро становится многословным, особенно когда вычисления разбиваются по категориям. Ниже — краткое руководство, которое держит логику прозрачной и снижает риск ошибок.

Базовый пример: кумулятивное среднее по группе

Предположим, у нас есть небольшой DataFrame с числовым столбцом и ключом группировки. Один из способов получить кумулятивное среднее — делить кумулятивную сумму на кумулятивный счётчик внутри каждой группы.

import polars as pl
tbl = pl.DataFrame({
    "val": [4, 6, 8, 11, 5, 6, 8, 15],
    "grp": ["A", "A", "B", "A", "B", "A", "B", "B"]
})
res = tbl.with_columns(
    running_avg = pl.col("val").cum_sum().over("grp") 
                  / pl.int_range(pl.len()).add(1).over("grp")
)

Что происходит и почему это выглядит громоздко

Выражение строит по-групповое бегущее среднее, беря разбитую по группам кумулятивную сумму и деля её на по-групповый бегущий индекс плюс один. Это работает, но использование глобального диапазона, который затем разбивается по группам, не самое читаемое решение и становится неудобным, когда вы пытаетесь распространить идею на кумулятивное стандартное отклонение.

Более аккуратный подход со скользящими окнами

Кумулятивное среднее и кумулятивное стандартное отклонение можно получить через rolling-функции. Идея проста: применяйте rolling_mean и rolling_std по каждой группе с окном размером во всю таблицу и параметром min_samples=1. В пределах группы, пока её длина меньше окна, такое окно фактически ведёт себя как кумулятивное. Для кумулятивного среднего можно также оставить компактный вариант на базе cum_sum и cum_count.

clean = tbl.with_columns(
    avg_cum = pl.col("val").cum_sum().over("grp") 
              / pl.col("val").cum_count().over("grp"),
    avg_cum_roll = pl.col("val").rolling_mean(
        window_size=tbl.shape[0],
        min_samples=1
    ).over("grp"),
    std_cum_roll = pl.col("val").rolling_std(
        window_size=tbl.shape[0],
        min_samples=1
    ).over("grp")
)

Rolling-среднее совпадает с кумулятивным средним. Для стандартного отклонения по rolling первая запись каждой группы даёт null, что соответствует обычному поведению при единственном наблюдении.

Зачем это важно

Бегущие агрегаты — основа аналитических пайплайнов, мониторинга, расчёта признаков и разведочного анализа. Читаемое выражение легче поддерживать и проще осмыслять. Подход со скользящими окнами избавляет от ручной арифметики с индексами, а компактная формула для кумулятивного среднего остаётся точной без лишних деталей. Более специализированной встроенной функции для этого паттерна нет, но семейство rolling хорошо закрывает потребность.

Вывод

Для кумулятивного среднего по группам предпочитайте короткую формулу: cum_sum, делённая на cum_count. Для кумулятивного стандартного отклонения используйте rolling_std с окном, равным числу строк, и min_samples=1; обе операции применяйте через over, чтобы разделить вычисления по ключу группировки. Это держит код компактным, выразительным и менее склонным к ошибкам по мере усложнения логики.