如何在Polars分组聚合结果中添加全量数据汇总行？

原始代码与需求 #

原始代码实现了按item_num分组，聚合满足consider == 1条件的rate字段均值、中位数和75分位数：

import polars as pl

df_test = pl.DataFrame({
    'item_num': [1, 2, 1, 2, 3, 1, 2, 3, 1, 2],
    'rate': [100, 200, 150, 250, 300, 120, 220, 350, 180, 270],
    'consider': [1, 1, 0, 1, 1, 0, 1, 1, 0, 1]
})

filter_for_agg = pl.col('consider') == 1

groups = (
    df_test.group_by('item_num')
    .agg(
        pl.col("rate").filter(filter_for_agg).mean().alias("mean"),
        pl.col("rate").filter(filter_for_agg).median().alias("median"),  
        pl.col("rate").filter(filter_for_agg).quantile(.75).alias("p75"),          
    )
)

需求：在groups数据框中新增一行，item_num值为'all'，对应所有满足filter_for_agg条件的数据的各项统计指标。

最优实现方案 #

核心思路：先一次性过滤出符合条件的数据，避免重复计算；再分别生成分组统计结果和全量统计行，最后拼接两个结果。这种方式计算效率最高，代码可读性也强。

完整代码 #

import polars as pl

df_test = pl.DataFrame({
    'item_num': [1, 2, 1, 2, 3, 1, 2, 3, 1, 2],
    'rate': [100, 200, 150, 250, 300, 120, 220, 350, 180, 270],
    'consider': [1, 1, 0, 1, 1, 0, 1, 1, 0, 1]
})

# 一次性过滤符合条件的数据，避免重复应用过滤逻辑
df_filtered = df_test.filter(pl.col('consider') == 1)

# 生成分组聚合结果
groups = df_filtered.group_by('item_num').agg(
    pl.col('rate').mean().alias('mean'),
    pl.col('rate').median().alias('median'),
    pl.col('rate').quantile(0.75).alias('p75')
)

# 生成全量统计行
total_row = df_filtered.select(
    pl.lit('all').alias('item_num'),
    pl.col('rate').mean().alias('mean'),
    pl.col('rate').median().alias('median'),
    pl.col('rate').quantile(0.75).alias('p75')
)

# 拼接分组结果与全量统计行
final_result = pl.concat([groups, total_row])

print(final_result)

输出结果 #

shape: (4, 4)
┌──────────┬────────────┬────────┬───────┐
│ item_num ┆ mean       ┆ median ┆ p75   │
│ ---      ┆ ---        ┆ ---    ┆ ---   │
│ str      ┆ f64        ┆ f64    ┆ f64   │
╞══════════╪════════════╪════════╪═══════╡
│ 2        ┆ 235.0      ┆ 235.0  ┆ 250.0 │
│ 1        ┆ 100.0      ┆ 100.0  ┆ 100.0 │
│ 3        ┆ 325.0      ┆ 325.0  ┆ 350.0 │
│ all      ┆ 241.428571 ┆ 250.0  ┆ 300.0 │
└──────────┴────────────┴────────┴───────┘

方案优势 #

1. 高效性：仅执行一次过滤操作，后续统计均基于过滤后的数据集，避免重复计算开销。
2. 可读性：代码逻辑拆分清晰，各步骤职责明确，便于后续维护和修改。
3. 灵活性：分组统计与全量统计的逻辑可独立调整，互不干扰。

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