如何在Polars中实现分组内按指定列排序并选取目标列Top K

确实可以用Polars的top_k方法配合分组/窗口函数实现，逻辑更紧凑，还能避免不必要的全局排序。

写法一：窗口函数+去重 #

(df
    # 第一步：计算每个门店-客户的总消费
    .group_by('storeId', 'customerId')
    .agg(pl.col('saleValue').sum().alias('totalSales'))
    # 第二步：按门店分组，取总消费最高的前5个客户
    .select(
        'storeId',
        pl.col('customerId').top_k(5, by='totalSales').over('storeId').alias('customerIds')
    )
    # 去重：窗口函数会给每个门店的每一行生成相同的客户列表，所以保留唯一值即可
    .unique('storeId')
)

写法二：嵌套分组聚合（更推荐） #

这种写法不需要额外去重，逻辑更直接：

(df
    # 先聚合出每个门店-客户的总消费
    .group_by('storeId', 'customerId')
    .agg(total_sales=pl.col('saleValue').sum())
    # 按门店分组，直接用top_k取总消费最高的前5个客户
    .group_by('storeId')
    .agg(
        customerIds=pl.col('customerId').top_k(5, by='total_sales')
    )
)

关键说明 #

1. top_k的参数作用：top_k(n=5, by='total_sales') 表示根据total_sales降序排列，取前5个customerId；如果要升序取前N，可以加descending=False参数。
2. 性能优势：两种写法都避免了全局排序，而是在每个门店的分组内做局部排序取前N，对于大数量级的数据（比如你的1000万行），性能会优于先全局排序再取head的方案。
3. 结果一致性：两种写法的输出结果和你原来的方案完全一致，都是每个门店对应总消费最高的前5个客户列表。

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