如何在Polars DataFrame中按列分组后选取最小值对应整行

需求说明 #

给定Polars DataFrame，按a列分组，提取每组中b列值最小的完整行数据。

原始数据 #

import polars as pl

a = pl.DataFrame({
    'a': [1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 3],
    'b': [10, 5, 10, 6, 5, 4, 5, 30, 25, 24, 30, 30],
    'c': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
})

解决方案 #

以下三种方法均可实现需求，根据场景选择即可：

方法1：排序后去重

先按a列（分组维度）、b列（目标列）排序，再按a列去重并保留每组第一行（即b最小的行），逻辑简洁高效：

result = a.sort(['a', 'b']).unique(subset='a', keep='first')
print(result)

方法2：分组聚合取首行

通过group_by明确分组，组内对所有列按b排序后取第一行，更直观体现分组逻辑：

result = a.group_by('a', maintain_order=True).agg(
    pl.all().sort_by('b').first()
)
print(result)

方法3：窗口函数过滤

用窗口函数标记每组中b为最小值的行，再过滤出标记行，适合需要保留中间判断逻辑的场景：

result = a.with_columns(
    is_min_b=pl.col('b') == pl.col('b').min().over('a')
).filter(pl.col('is_min_b')).drop('is_min_b')
print(result)

输出结果 #

三种方法均会得到如下目标DataFrame：

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

内容的提问来源于stack exchange，提问作者Kaster