Polars分组加权平均：权重和为0时改用等权均值

需求明确：对Polars DataFrame做分组加权平均时，需满足：

1. 自动忽略待平均列（vr）和权重列（wt）中的空值
2. 当组内有效权重之和为0时，返回该组的等权均值而非NaN

符号定义 #

• df: 输入的Polars DataFrame
• gr: 分组列名的列表
• vr: 需要计算均值的列名列表
• wt: 权重列的名称

解决方案：用Polars表达式实现高效计算 #

直接用Polars的向量化表达式实现逻辑，避免Python自定义函数的性能损耗，同时满足需求：

import polars as pl

def weighted_avg_with_fallback(gr: list[str], vr: list[str], wt: str):
    # 标记有效权重：仅保留待平均列非空时的权重
    valid_wt = pl.col(wt) * pl.col(vr).is_not_null()
    
    # 计算各分组的三个核心值：加权和、有效权重和、等权均值
    weighted_sum = (pl.col(vr) * pl.col(wt)).sum()
    valid_wt_total = valid_wt.sum()
    equal_weight_mean = pl.col(vr).mean()
    
    # 构造条件表达式：权重和为0时用等权均值，否则用加权平均
    wavg_expr = pl.when(valid_wt_total == 0) \
                  .then(equal_weight_mean) \
                  .otherwise(weighted_sum / valid_wt_total)
    
    # 执行分组聚合，返回结果
    return df.group_by(gr, maintain_order=True).agg(
        {col: wavg_expr.alias(col) for col in vr}
    )

示例测试 #

构造示例数据验证逻辑：

# 测试数据
df = pl.DataFrame({
    "id": ["x", "x", "y", "y", "y"],
    "a": [2, None, 3, 4, None],
    "wt": [0, 0, 1, 2, 0]
})

# 计算结果
result = weighted_avg_with_fallback(gr=["id"], vr=["a"], wt="wt")
print(result)

输出结果：

shape: (2, 2)
┌─────┬───────────┐
│ id  ┆ a         │
│ --- ┆ ---       │
│ str ┆ f64       │
╞═════╪═══════════╡
│ x   ┆ 2.0       │  # 组内有效权重和为0，返回等权均值2.0
│ y   ┆ 3.666667  │  # 加权平均：(3*1 +4*2)/(1+2) = 11/3 ≈3.6667
└─────┴───────────┘

关键逻辑说明 #

1. 空值过滤：通过pl.col(vr).is_not_null()生成布尔掩码，与权重列相乘后，自动将待平均列为空的行的权重置为0，实现空值忽略。
2. 条件 fallback：利用pl.when/then/otherwise实现分支判断，当有效权重和为0时，替换为Polars原生的mean()（自动忽略空值）。
3. 性能优势：全程使用Polars的向量化表达式，比Pandas中apply调用Python自定义函数的方式效率更高，尤其适合大数据量场景。

对比Pandas实现 #

如果习惯Pandas的自定义函数方式，逻辑如下（供参考）：

import pandas as pd

def _wavg_py(gr, vr_col, wt_col):
    # 过滤掉待平均列和权重列的空值
    filtered = gr[[vr_col, wt_col]].dropna()
    weight_sum = filtered[wt_col].sum()
    # 权重和为0时返回等权均值，否则返回加权平均
    if weight_sum == 0:
        return gr[vr_col].mean()
    else:
        return (filtered[vr_col] * filtered[wt_col]).sum() / weight_sum

# 调用示例
pd_df = df.to_pandas()
pd_result = pd_df.groupby(by=["id"]).apply(_wavg_py, vr_col="a", wt_col="wt")
print(pd_result)

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