如何用Polars按group1分组计算group2占比？求更优实现

初始数据与需求 #

初始代码与DataFrame #

import polars as pl

df = pl.DataFrame({
    'group1': ['a', 'a', 'b', 'c', 'a', 'b'], 
    'group2': [0, 1, 1, 0, 1, 1]
})

对应的DataFrame输出：

shape: (6, 2)
┌────────┬────────┐
│ group1 ┆ group2 │
│ ---    ┆ ---    │
│ str    ┆ i64    │
╞════════╪════════╡
│ a      ┆ 0      │
│ a      ┆ 1      │
│ b      ┆ 1      │
│ c      ┆ 0      │
│ a      ┆ 1      │
│ b      ┆ 1      │
└────────┴────────┘

需求 #

针对每个group1分组，计算group2的分布情况，包含计数与占比，期望输出如下：

shape: (4, 4)
┌────────┬────────┬───────┬────────────┐
│ group1 ┆ group2 ┆ count ┆ percentage │
│ ---    ┆ ---    ┆ ---   ┆ ---        │
│ str    ┆ i64    ┆ u32   ┆ f64        │
╞════════╪════════╪═══════╪════════════╡
│ a      ┆ 0      ┆ 1     ┆ 0.333333   │
│ a      ┆ 1      ┆ 2     ┆ 0.666667   │
│ b      ┆ 1      ┆ 2     ┆ 1.0        │
│ c      ┆ 0      ┆ 1     ┆ 1.0        │
└────────┴────────┴───────┴────────────┘

原实现代码 #

counts = df.group_by('group1', 'group2').len()
counts.with_columns(
    (
         counts['len']
         / counts.select(pl.col('len').sum().over('group1'))['len']
    ).alias('percentage')
).sort('group1', 'group2')

更符合Polars风格的优化实现 #

(df
 .group_by('group1', 'group2')
 .len().alias('count')
 .with_columns(
     percentage=pl.col('count') / pl.col('count').sum().over('group1')
 )
 .sort('group1', 'group2')
)

优化点说明 #

• 链式调用：将所有操作连贯起来，无需中间变量，完全贴合Polars推崇的流畅代码风格。
• 简化窗口函数使用：直接在with_columns中通过pl.col('count').sum().over('group1')计算分组总和，避免额外的select取值操作，代码更简洁直观。
• 直接命名计数列：用.alias('count')将分组计数结果直接命名为期望的列名，省去后续改名步骤。
• 可选顺序保留：如果需要保留原始数据的分组顺序，可以在group_by中添加maintain_order=True参数，更灵活贴合需求。

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