Computing and retrieving operations at the group level without collapsing data frame in polars?

Viewed 28

I am trying to compute a stat (or more) at the group level without having to create a second data frame. The current way I do it is by relying on the generation of a second data frame with the desired aggregation that I then merge back to the original one.

A silly example:

import polars as pl

df = pl. DataFrame( {'name' : ['Steve', 'Larry', 'Tom', 'Steve', 'Tom', 'Steve'],
                     'points': range(6)})
print(df)
shape: (6, 2)
┌───────┬────────┐
│ name  ┆ points │
│ ---   ┆ ---    │
│ str   ┆ i64    │
╞═══════╪════════╡
│ Steve ┆ 0      │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ Larry ┆ 1      │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ Tom   ┆ 2      │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ Steve ┆ 3      │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ Tom   ┆ 4      │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ Steve ┆ 5      │
└───────┴────────┘

We created a simple data frame below in which some groups have more entries than others. In a second step we compute an additional data frame to keep track of the size of each group.

entries= df.groupby('name').agg(pl.count().alias('entries'))
print(entries)
shape: (3, 2)
┌───────┬─────────┐
│ name  ┆ entries │
│ ---   ┆ ---     │
│ str   ┆ u32     │
╞═══════╪═════════╡
│ Steve ┆ 3       │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ Tom   ┆ 2       │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ Larry ┆ 1       │
└───────┴─────────┘

Now we bring back this information to the original data frame in a third step.

print(df.join(entries, left_on='name', right_on='name', how='left'))
shape: (6, 3)
┌───────┬────────┬─────────┐
│ name  ┆ points ┆ entries │
│ ---   ┆ ---    ┆ ---     │
│ str   ┆ i64    ┆ u32     │
╞═══════╪════════╪═════════╡
│ Steve ┆ 0      ┆ 3       │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ Larry ┆ 1      ┆ 1       │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ Tom   ┆ 2      ┆ 2       │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ Steve ┆ 3      ┆ 3       │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ Tom   ┆ 4      ┆ 2       │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ Steve ┆ 5      ┆ 3       │
└───────┴────────┴─────────┘

Is there a way to avoid this triangulation? I have the feeling that using over might be a solution but I can't figure it out yet.

1 Answers

Well ... I managed. Posting the question helped me organize my thoughts and indeed, over was the solution.

df.with_column(pl.col('name').count().over('name').alias('entries'))

shape: (6, 3)
┌───────┬────────┬─────────┐
│ name  ┆ points ┆ entries │
│ ---   ┆ ---    ┆ ---     │
│ str   ┆ i64    ┆ u32     │
╞═══════╪════════╪═════════╡
│ Steve ┆ 0      ┆ 3       │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ Larry ┆ 1      ┆ 1       │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ Tom   ┆ 2      ┆ 2       │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ Steve ┆ 3      ┆ 3       │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ Tom   ┆ 4      ┆ 2       │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ Steve ┆ 5      ┆ 3       │
└───────┴────────┴─────────┘

Related