嘿，我太懂这种对着简单问题卡壳的感觉了！你一开始用的group_by().agg()其实是做聚合压缩——它会把每个分组合并成一行，但你要的是保留所有原始行，同时给每行添加上对应分组的求和结果，这时候用窗口函数就对了！

举个具体的例子，就用你给出的数据集：

import polars as pl

df = pl.DataFrame(
    {
        "a": ["a", "b", "a", "b", "c"],
        "b": [1, 2, 1, 3, 3],
    }
)

如果你用传统的group_by().agg()，得到的是每个分组压缩后的结果：

df.group_by("a").agg(pl.col("b").sum())
# 输出：
# shape: (3, 2)
# ┌─────┬─────┐
# │ a   ┆ b   │
# │ --- ┆ --- │
# │ str ┆ i64 │
# ╞═════╪═════╡
# │ a   ┆ 2   │
# │ b   ┆ 5   │
# │ c   ┆ 3   │
# └─────┴─────┘

但你要的是保留原始的5行，每行都带上对应分组的总和，这时候用over()指定窗口分区就行：

# 新增分组求和列，保留所有原始行
result_df = df.with_columns(
    group_sum_b=pl.col("b").sum().over("a")
)

print(result_df)

运行后会得到你想要的结果：

shape: (5, 3)
┌─────┬─────┬───────────┐
│ a   ┆ b   ┆ group_sum_b │
│ --- ┆ --- ┆ ---       │
│ str ┆ i64 ┆ i64       │
╞═════╪═════╪═══════════╡
│ a   ┆ 1   ┆ 2         │
│ b   ┆ 2   ┆ 5         │
│ a   ┆ 1   ┆ 2         │
│ b   ┆ 3   ┆ 5         │
│ c   ┆ 3   ┆ 3         │
└─────┴─────┴───────────┘

简单解释下：

• pl.col("b").sum()是计算求和，但加上.over("a")后，它会按"a"列的分组分别计算求和
• 这个求和结果会被自动广播到该分组的每一行，所以原始的所有行都保留了，同时每行都能看到自己所在分组的总和

如果需要按多列分组计算，也很简单，把over()里的参数改成列表就行，比如over(["a", "another_column"])。

备注：内容来源于stack exchange，提问作者gernophil