如何用Polars窗口函数计算分组中排除当前行的标准差？

Polars没有直接对应SQL中EXCLUDE CURRENT ROW窗口子句的语法，但可以通过分组统计量的数学推导实现相同效果——利用分组的总计数、总和、平方和，减去当前行的值后计算标准差，这种方法效率远高于逐行过滤分组的暴力实现。

先看DuckDB的参考实现（对比用） #

SELECT
  answer,
  value,
  STDDEV(value) OVER (
    PARTITION BY answer
    EXCLUDE CURRENT ROW
  ) AS stddev_exclude_current
FROM df;

Polars实现步骤 #

1. 构造示例数据

import polars as pl

df = pl.DataFrame({
    "answer": ["A", "A", "A", "B", "B"],
    "value": [1, 2, 3, 4, 5]
})

2. 计算分组统计量

先统计每个分组的总条数、数值总和、数值平方和，再关联回原表：

group_stats = df.group_by("answer").agg(
    pl.col("value").count().alias("count_total"),
    pl.col("value").sum().alias("sum_total"),
    pl.col("value").pow(2).sum().alias("sum_sq_total")
)

df_with_stats = df.join(group_stats, on="answer")

3. 推导排除当前行后的标准差

根据样本标准差公式，排除当前行后的统计量可通过分组总量减去当前行值得到，再代入公式计算：

result = df_with_stats.with_columns(
    pl.when(pl.col("count_total") - 1 >= 2)  # 至少需要2个样本才能计算标准差
    .then(
        (
            (pl.col("sum_sq_total") - pl.col("value").pow(2)) 
            - ((pl.col("sum_total") - pl.col("value")) ** 2) / (pl.col("count_total") - 1)
        ).div(pl.col("count_total") - 2)
        .sqrt()
    )
    .alias("stddev_exclude_current")
).drop(["count_total", "sum_total", "sum_sq_total"])  # 移除中间计算列

print(result)

期望输出 #

shape: (5, 3)
┌────────┬───────┬──────────────────────┐
│ answer ┆ value ┆ stddev_exclude_current │
│ ---    ┆ ---   ┆ ---                  │
│ str    ┆ i64   ┆ f64                  │
╞════════╪═══════╪══════════════════════╡
│ A      ┆ 1     ┆ 1.0                  │
│ A      ┆ 2     ┆ 1.0                  │
│ A      ┆ 3     ┆ 1.0                  │
│ B      ┆ 4     ┆ null                 │
│ B      ┆ 5     ┆ null                 │
└────────┴───────┴──────────────────────┘

原理说明 #

样本标准差公式：$\sigma = \sqrt{\frac{\sum x^2 - \frac{(\sum x)^2}{n}}{n-1}}$

• 排除当前行后，样本量 $n$ 变为 $n-1$
• 总和 $\sum x$ 变为 $\sum_{total}x - x_{current}$
• 平方和 $\sum x^2$ 变为 $\sum_{total}x^2 - x_{current}^2$

当分组剩余样本量小于2时，无法计算有效标准差，返回null。这种方法只需一次分组聚合和关联操作，适合处理大数据集。

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