如何在Polars中计算分组SKU/标识组合占总值的百分比？

原始DataFrame #

import polars as pl

df = pl.from_repr("""
┌──────────────┬──────┬──────┐
│ id           ┆ sku1 ┆ sku2 │
│ ---          ┆ ---  ┆ ---  │
│ str          ┆ i64  ┆ i64  │
╞══════════════╪══════╪══════╡
│ identifier 1 ┆ 4    ┆ null │
│ identifier 1 ┆ 2    ┆ 3    │
│ identifier 2 ┆ 3    ┆ null │
│ identifier 2 ┆ null ┆ 3    │
│ identifier 2 ┆ 1    ┆ null │
└──────────────┴──────┴──────┘
""")

需求说明 #

计算每个id分组下，sku1、sku2列的分组求和值占所有非null数值全局总和的百分比，得到如下结果：

shape: (2, 3)
┌──────────────┬───────┬────────┐
│ id           ┆ sku1  ┆ sku2   │
│ ---          ┆ ---   ┆ ---    │
│ str          ┆ f64   ┆ f64    │
╞══════════════╪═══════╪════════╡
│ identifier 1 ┆ 0.375 ┆ 0.1875 │
│ identifier 2 ┆ 0.25  ┆ 0.1875 │
└──────────────┴───────┴────────┘

实现代码 #

可以通过链式调用完成单步骤计算：

result = (
    df
    # 计算全局所有非null数值总和并添加为临时列
    .with_columns(pl.sum_horizontal(pl.all().exclude("id")).alias("global_total"))
    # 按id分组
    .group_by("id")
    # 聚合计算各sku列的分组占比
    .agg(
        (pl.col("sku1").sum() / pl.col("global_total").first()).alias("sku1"),
        (pl.col("sku2").sum() / pl.col("global_total").first()).alias("sku2")
    )
)

print(result)

也可以先提取全局总和再分组计算，逻辑更直观：

# 先计算所有非null数值的全局总和
global_total = df.select(pl.sum_horizontal(pl.all().exclude("id"))).item()

# 分组求和并计算占比
result = df.group_by("id").agg(
    (pl.col("sku1").sum() / global_total).alias("sku1"),
    (pl.col("sku2").sum() / global_total).alias("sku2")
)

代码说明 #

1. 全局总和计算：用pl.sum_horizontal(pl.all().exclude("id"))对id外所有列的非null值求和，得到所有数值的总合（示例中为16）。
2. 分组占比计算：按id分组后，对每个分组的sku1、sku2列求和，再除以全局总和，最终得到各分组数值占全局总值的百分比。

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