Polars中嵌套over调用的优化实现方案咨询

要实现粗粒度组内排除当前细粒度子组后的均值，无需嵌套over调用，可通过分层聚合后做算术运算的方式实现，完全贴合Polars的向量化设计思路：

核心思路 #

1. 先计算粗粒度组的总总和与总计数
2. 再计算当前细粒度子组的总和与计数
3. 用粗粒度组的聚合结果减去子组的聚合结果，得到排除子组后的总和与计数，最后求均值
4. 额外处理粗粒度组仅包含当前子组的边界情况（避免除以0）

代码示例 #

假设你的数据包含粗粒度分组列group_col、细粒度子组列subgroup_col，以及目标计算列value：

import polars as pl

# 示例数据
df = pl.DataFrame({
    "group_col": ["A", "A", "A", "B", "B", "C"],
    "subgroup_col": ["X", "X", "Y", "Z", "Z", "W"],
    "value": [10, 20, 30, 40, 50, 60]
})

# 计算排除当前子组后的均值
result_df = df.with_columns(
    mean_excl_subgroup=pl.when(
        # 仅当粗粒度组包含多个子组时计算
        pl.col("value").count().over("group_col") > pl.col("value").count().over(["group_col", "subgroup_col"])
    ).then(
        (
            pl.col("value").sum().over("group_col") - pl.col("value").sum().over(["group_col", "subgroup_col"])
        ) / (
            pl.col("value").count().over("group_col") - pl.col("value").count().over(["group_col", "subgroup_col"])
        )
    ).otherwise(None)  # 粗粒度组仅一个子组时返回None
)

print(result_df)

输出结果 #

shape: (6, 4)
┌───────────┬──────────────┬───────┬────────────────────┐
│ group_col ┆ subgroup_col ┆ value ┆ mean_excl_subgroup │
│ ---       ┆ ---          ┆ ---   ┆ ---                │
│ str       ┆ str          ┆ i64   ┆ f64                │
╞═══════════╪══════════════╪═══════╪════════════════════╡
│ A         ┆ X            ┆ 10    ┆ 30.0               │
│ A         ┆ X            ┆ 20    ┆ 30.0               │
│ A         ┆ Y            ┆ 30    ┆ 15.0               │
│ B         ┆ Z            ┆ 40    ┆ null               │
│ B         ┆ Z            ┆ 50    ┆ null               │
│ C         ┆ W            ┆ 60    ┆ null               │
└───────────┴──────────────┴───────┴────────────────────┘

方案优势 #

• 完全使用Polars原生表达式，无需自定义函数或冗余参数传递
• 利用向量化计算，性能优于循环或嵌套分组逻辑
• 逻辑清晰，直接对应“排除子组后求均值”的业务需求

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