Polars 獲得列值最大的分組行

Question

所以考慮這個片段

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()])

這給了我：

class score
  b     0.4
  a     0.5

但我想要與最高分相對應的組的整行。 我可以與原始數據框進行連接，例如

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

要得到

class  score  name
  a     0.5    Jon
  b     0.4    Von

有什么方法可以避免連接並將 name 列作為 groupby 聚合的一部分？

Answer 1

您可以使用sort_by表達式按分數對每個組中的觀察結果進行排序，然后使用last表達式獲取最后一個觀察值。

例如，要獲取所有列：

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   │
└───────┴──────┴───────┘

編輯：使用over

如果您有多個觀察值是max ，另一種獲取所有行的簡單方法是使用over 。

例如，如果您的數據中b班有兩個學生（“Von”和“Yvonne”）並列最高分：

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   │
└───────┴────────┴───────┘