Polars中对应Pandas idxmin的分组聚合实现方案

要实现按[a, b]分组后，获取d列最小值对应的c列值，Polars有几种高效的实现方式，完全适配百万级数据量的处理需求：

方法一：分组聚合时直接取排序后的首个值 #

这是最简洁高效的方式，利用Polars分组聚合中的sort_by方法对每组内的d列排序，再取对应c列的第一个值：

import polars as pl

# 构造示例DataFrame
df = pl.DataFrame({
    'a': [1, 1, 2, 2],
    'b': ['x', 'x', 'y', 'y'],
    'c': ['foo', 'bar', 'baz', 'qux'],
    'd': [3, 1, 4, 2]
})

# 分组聚合获取结果
result = df.group_by(['a', 'b'], maintain_order=True).agg(
    pl.col('c').sort_by('d').first().alias('c_min_d')
)

print(result)

输出结果：

shape: (2, 3)
┌─────┬─────┬─────────┐
│ a   ┆ b   ┆ c_min_d │
│ --- ┆ --- ┆ ---     │
│ i64 ┆ str ┆ str     │
╞═════╪═════╪═════════╡
│ 1   ┆ x   ┆ bar     │
│ 2   ┆ y   ┆ qux     │
└─────┴─────┴─────────┘

方法二：窗口函数筛选每组最小d对应的行 #

如果需要保留原DataFrame的更多列，或者需要更灵活的筛选逻辑，可以用窗口函数标记每组内d最小的行：

result = df.with_columns(
    # 在[a,b]分组内，按d升序排序后标记行号
    pl.row_number().over(['a', 'b']).sort_by(pl.col('d')).alias('rn')
).filter(pl.col('rn') == 1).select(['a', 'b', 'c'])

为什么Polars的arg_min不适用？ #

Polars的pl.arg_min('d')是全局函数，返回的是整个DataFrame中d值最小的行索引，而非分组内的索引。如果要在分组场景下使用，需要结合窗口函数，但相比上面两种方法，代码复杂度更高，效率也没有优势。

以上两种方法都是Polars矢量化引擎优化后的实现，处理百万级数据的速度远快于Pandas，完全满足你的性能需求。

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