Get a grouped sum in polars, but keep all individual rows

Question

I am breaking my head over this probably pretty simply question and I just can't find the answer anywhere. I want to create a new column with a grouped sum of another column, but I want to keep all individual rows. So, this is what the docs say:

import polars as pl

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

df.group_by("a").agg(pl.col("b").sum())

The output of this would be:

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

However, what I need would be this:

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

I could create the sum in a separate df and then join it with the original one, but I am pretty sure, there is an easier solution.

It is also covered afer the "Aggregation" section in the User Guide: docs.pola.rs/user-guide/expressions/window-functions — jqurious
– jqurious, Commented Apr 16 at 9:58

Jonathan · Accepted Answer · 2025-04-16 09:37:41Z

3

All you need is a window function:

df.with_columns(
    b_sum=pl.col("b").sum().over(pl.col("a"))
)


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

answered Apr 16 at 9:37

Jonathan

2,3332 gold badges12 silver badges31 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

gernophil Apr 16 at 9:41

Jeez, sometimes it's really that easy. Was totally unaware of the over(). This helps a lots. I'll accept it as soon as the 5 minutes have passed :D.

Jonathan Apr 16 at 9:58

@gernophil Exactly. It's similar to SQL.

Collectives™ on Stack Overflow

Get a grouped sum in polars, but keep all individual rows

1 Answer 1

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related