如何用PySpark窗口函数获取每组前N个带计数的名称列表？

步骤1：定义窗口并计算排名 #

先通过窗口函数按variation分组、count降序排序，给组内每行计算排名。这里推荐用row_number()，它会给同组内的记录分配唯一排名，避免因并列count导致最终列表超出N个；如果需要保留并列排名的记录，改用dense_rank()即可。

from pyspark.sql import Window
import pyspark.sql.functions as F

# 定义窗口规则：按商户变体分组，按count降序排列
window_spec = Window.partitionBy("variation").orderBy(F.desc("count"))

# 为每条记录添加排名列
ranked_df = df.withColumn("rank", F.row_number().over(window_spec))

步骤2：筛选前N个记录并聚合为列表 #

筛选排名≤N的行后，按variation分组，用collect_list()聚合name生成目标列表。Spark 2.0+版本中，collect_list()会保留窗口排序后的顺序，确保列表内的name是按count降序排列的。

N = 10
result_df = ranked_df.filter(F.col("rank") <= N) \
    .groupBy("variation") \
    .agg(F.collect_list("name").alias(f"top_{N}_names"))

关键注意点 #

• 你之前出现重复行的核心原因是只做了排名筛选，没有对variation进行分组聚合，所以同一变体对应多行记录。
• 若业务允许并列排名（比如相同count的name都要保留，即使最终列表长度超过N），将row_number()替换为dense_rank()即可。

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