Polars get grouped rows where column value is maximum

user_4126652 0
Viewed 130

So consider this snippet

import polars as pl

df = pl.DataFrame({'class': ['a', 'a', 'b', 'b'], 'name': ['Ron', 'Jon', 'Don', 'Von'], 'score': [0.2, 0.5, 0.3, 0.4]})
df.groupby('class').agg([pl.col('score').max()])

This gives me:

class score
  b     0.4
  a     0.5

But I want the entire row of the group that corresponded to the maximum score. I can do a join with the original dataframe like

sdf = df.groupby('class').agg([pl.col('score').max()])
sdf.join(df, on=['class', 'score'])

To get

class  score  name
  a     0.5    Jon
  b     0.4    Von

Is there any way to avoid the join and include the name column as part of the groupby aggregation?

python-polars python
1 Answers

You can use a sort_by expression to sort your observations in each group by score, and then use the last expression to take the last observation.

For example, to take all columns:

df.groupby('class').agg([
    pl.all().sort_by('score').last(),
])

shape: (2, 3)
┌───────┬──────┬───────┐
│ class ┆ name ┆ score │
│ ---   ┆ ---  ┆ ---   │
│ str   ┆ str  ┆ f64   │
╞═══════╪══════╪═══════╡
│ a     ┆ Jon  ┆ 0.5   │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ b     ┆ Von  ┆ 0.4   │
└───────┴──────┴───────┘

Edit: using over

If you have more than one observation that is the max, another easy way to get all rows is to use over.

For example, if your data has two students in class b ('Von' and 'Yvonne') who tied for highest score:

df = pl.DataFrame(
    {
        "class": ["a", "a", "b", "b", "b"],
        "name": ["Ron", "Jon", "Don", "Von", "Yvonne"],
        "score": [0.2, 0.5, 0.3, 0.4, 0.4],
    }
)
df

shape: (5, 3)
┌───────┬────────┬───────┐
│ class ┆ name   ┆ score │
│ ---   ┆ ---    ┆ ---   │
│ str   ┆ str    ┆ f64   │
╞═══════╪════════╪═══════╡
│ a     ┆ Ron    ┆ 0.2   │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ a     ┆ Jon    ┆ 0.5   │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ b     ┆ Don    ┆ 0.3   │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ b     ┆ Von    ┆ 0.4   │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ b     ┆ Yvonne ┆ 0.4   │
└───────┴────────┴───────┘

df.filter(pl.col('score') == pl.col('score').max().over('class'))

shape: (3, 3)
┌───────┬────────┬───────┐
│ class ┆ name   ┆ score │
│ ---   ┆ ---    ┆ ---   │
│ str   ┆ str    ┆ f64   │
╞═══════╪════════╪═══════╡
│ a     ┆ Jon    ┆ 0.5   │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ b     ┆ Von    ┆ 0.4   │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ b     ┆ Yvonne ┆ 0.4   │
└───────┴────────┴───────┘
user_185598750
Related