Polars中分组后组内排序取首行的最优实现方式

你想要的分组内单独排序再取第一个值的操作，Polars完全支持，不需要先对整个DataFrame排序。直接在group_by后的聚合操作里，针对每组的目标列做排序再取首值即可：

import polars as pl

df = pl.DataFrame({"foo": [1, 1, 1, 2, 2, 2, 3], "bar": [5, 7, 6, 4, 2, 3, 1]})
df_desired = pl.DataFrame({"foo": [1, 2, 3], "bar": [5, 2, 1]})

# 分组内排序取首值的实现
df_suggestion = df.group_by("foo", maintain_order=True).agg(
    pl.col("bar").sort().first()
)

assert df_desired.equals(df_suggestion)

说明 #

• pl.col("bar").sort()会对每个分组内的"bar"列单独排序，而非全表排序，完美契合你提到的效率优化思路（分组规模n更小，排序时间复杂度为n*log(n)）。
• 如果需要按降序排序取首值，只需给sort()加reverse=True参数：pl.col("bar").sort(reverse=True).first()。
• maintain_order=True保证分组结果的顺序和原DataFrame中"foo"首次出现的顺序一致，避免额外的排序操作。

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