如何在Python Polars DataFrame中按分组选取前k行？

Polars DataFrame的top_k方法可选取某列排序后前k个最大值对应的完整行。例如以下代码会选出val列最大的两行：

df = pl.DataFrame({'grp':['a','a','a','b','b','b'], 'val':[1,2,3,10,20,30], 'etc':[0,1,2,3,4,5]})
# 原DataFrame内容：
# grp val etc
# str i64 i64
# "a" 1   0
# "a" 2   1
# "a" 3   2
# "b" 10  3
# "b" 20  4
# "b" 30  5

df.top_k(2, by='val')
# 输出结果：
# grp val etc
# str i64 i64
# "b" 30  5
# "b" 20  4

当前需求是：按分组获取每组中val列前k大值对应的完整行。尝试直接调用df.groupby('grp').top_k(2, by='val')无法运行，因为Polars的GroupBy对象没有top_k方法。理想输出应为：

grp val etc
str i64 i64
"b" 30  5
"b" 20  4
"a" 3   2
"a" 2   1

目前已找到两种可行但性能不佳的实现方式：

• 使用map_groups：性能表现通常较差
• 先排序再取前k：算法效率低于top_k（小k大样本场景下，排序复杂度为O(n log n)，而top_k为O(n)）

两种性能不佳的实现代码 #

方案1：依赖map_groups

# 可行但性能差
df.group_by('grp').map_groups(lambda df: df.top_k(2, by='val'))
# 输出结果：
# grp val etc
# str i64 i64
# "b" 30  5
# "b" 20  4
# "a" 3   2
# "a" 2   1

方案2：整组排序后取前k

# 可行但需要对整组排序
df.group_by('grp').agg(pl.all().sort_by('val', descending=True).head(2)).explode('val','etc')
# 输出结果：
# grp val etc
# str i64 i64
# "a" 3   2
# "a" 2   1
# "b" 30  5
# "b" 20  4

现有方案总结 #

• df.group_by('grp').top_k(2, by='val')：Polars中无法运行
• df.group_by('grp').map_groups(lambda df: df.top_k(2, by='val'))：可行但依赖map_groups，性能差
• df.group_by('grp').agg(pl.all().sort_by('val', descending=True).head(2)).explode('val','etc')：可行但需对整组排序，效率低于top_k

需要一种能直接结合GroupBy与top_k的高效方案。

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