如何在Polars中计算排除null值的分组加权平均值？

问题说明 #

当分组数据里存在null/None值时，常规加权平均计算会出现偏差。比如示例中分组1的a列有一个None，正确的加权平均应该只计算a列非空行的权重，但常规写法会把分组内所有权重都纳入分母，导致结果错误。

错误示例回顾 #

import polars as pl

data = {"id":[1, 1, 2, 2], "a" : [2, None, 1, 3], "b":[0,1,2,3], "weights":[0.5, 1, 0.2, 3]}
df = pl.DataFrame(data)
weighted_average = (pl.col("a", "b") * pl.col("weights")).sum() / pl.col("weights").sum()
df.group_by("id").agg(weighted_average)

执行结果里分组1的a列得到0.666667，但正确值应为2（仅计算a非空行：(2*0.5)/0.5）。

解决方案 #

核心是针对每个待计算列，用该列的非空掩码过滤权重后再求和，确保分子分母的计算范围一致。

方法1：逐列定义加权平均逻辑 #

import polars as pl

data = {"id":[1, 1, 2, 2], "a" : [2, None, 1, 3], "b":[0,1,2,3], "weights":[0.5, 1, 0.2, 3]}
df = pl.DataFrame(data)

def weighted_avg(col_name):
    # 分子：仅计算非空行的加权和（null参与运算后仍为null，sum会自动跳过）
    numerator = (pl.col(col_name) * pl.col("weights")).sum()
    # 分母：仅保留当前列非空对应的权重求和
    denominator = pl.col("weights").filter(pl.col(col_name).is_not_null()).sum()
    return numerator / denominator

result = df.group_by("id").agg(
    weighted_avg("a").alias("a"),
    weighted_avg("b").alias("b")
)

print(result)

方法2：批量处理多列 #

如果需要计算的列较多，可通过循环批量生成聚合表达式：

import polars as pl

data = {"id":[1, 1, 2, 2], "a" : [2, None, 1, 3], "b":[0,1,2,3], "weights":[0.5, 1, 0.2, 3]}
df = pl.DataFrame(data)

cols_to_calc = ["a", "b"]
agg_exprs = [
    ((pl.col(c) * pl.col("weights")).sum() / pl.col("weights").filter(pl.col(c).is_not_null()).sum()).alias(c)
    for c in cols_to_calc
]

result = df.group_by("id").agg(agg_exprs)
print(result)

正确执行结果 #

shape: (2, 3)
┌─────┬─────┬──────────┐
│ id  ┆ a   ┆ b        │
│ --- ┆ --- ┆ ---      │
│ i64 ┆ f64 ┆ f64      │
╞═════╪═════╪══════════╡
│ 1   ┆ 2.0 ┆ 0.666667 │
│ 2   ┆ 2.875 ┆ 2.9375 │
└─────┴─────┴──────────┘

分组1的a列得到正确值2.0，b列因无null值，结果与常规计算一致。

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