基于Polars DataFrame计算各ID分组SKU列占总计的百分比

问题背景 #

给定如下Polars DataFrame：

import polars as pl

df = pl.DataFrame({
    "id": [1, 1, 2, 2, 2],
    "sku1": [4, 2, 3, None, 1],
    "sku2": [None, 3, None, 3, None],
})

需要实现每个id对应的各SKU列总和占该SKU列全局总计的百分比，且必须适配SKU列数量动态变化的场景。

实现思路 #

1. 动态匹配SKU列：通过列名前缀筛选所有SKU相关列，避免硬编码列名，兼容列数量增减的情况。
2. 分组聚合SKU总和：按id分组，计算每个分组下各SKU列的求和值（自动忽略None）。
3. 计算占比百分比：将分组后的SKU总和除以对应列的全局总计，转换为百分比格式，可选保留指定小数位数。

代码实现 #

方法一：分步计算（直观清晰） #

import polars as pl

df = pl.DataFrame({
    "id": [1, 1, 2, 2, 2],
    "sku1": [4, 2, 3, None, 1],
    "sku2": [None, 3, None, 3, None],
})

# 1. 动态获取所有SKU列
sku_cols = [col for col in df.columns if col.startswith("sku")]

# 2. 按id分组计算各SKU列的总和
grouped_df = df.group_by("id").agg(pl.sum(col).alias(col) for col in sku_cols)

# 3. 计算百分比（使用窗口函数获取全局总计）
result_df = grouped_df.with_columns(
    [(pl.col(col) / pl.sum(col).over() * 100).round(2).alias(f"{col}_pct") for col in sku_cols]
)

print(result_df)

输出结果 #

shape: (2, 5)
┌─────┬──────┬──────┬──────────┬──────────┐
│ id  ┆ sku1 ┆ sku2 ┆ sku1_pct ┆ sku2_pct │
│ --- ┆ ---  ┆ ---  ┆ ---      ┆ ---      │
│ i64 ┆ i64  ┆ i64  ┆ f64      ┆ f64      │
╞═════╪══════╪══════╪══════════╪══════════╡
│ 1   ┆ 6    ┆ 3    ┆ 50.0     ┆ 50.0     │
│ 2   ┆ 4    ┆ 3    ┆ 33.33    ┆ 50.0     │
└─────┴──────┴──────┴──────────┴──────────┘

方法二：精简链式写法 #

如果追求代码简洁，可将聚合与百分比计算合并为链式调用：

result_df = (
    df.group_by("id")
    .agg(pl.sum(col).alias(col) for col in sku_cols)
    .with_columns(
        [(pl.col(col) / pl.sum(col).over() * 100).round(2).alias(f"{col}_pct") for col in sku_cols]
    )
)

关键说明 #

• 动态列适配：通过col.startswith("sku")自动识别SKU列，新增或删除SKU列时无需修改代码。
• 窗口函数over()：pl.sum(col).over()会计算该列的全局总和，替代单独提取总计的步骤，代码更高效。
• 小数位数控制：使用.round(n)可指定百分比保留的小数位数，根据需求调整即可。

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