Spark Scala实现按item_id1取Top N最高similarity_score记录

嘿，你的思路完全正确！要实现这个需求，Spark的**窗口函数（Window Functions）**就是最佳方案，咱们一步步来落地实现：

核心步骤拆解 #

1. 导入必要依赖：先引入窗口函数相关的类和工具方法
2. 定义窗口规则：按item_id1分组，再按score降序排序
3. 标记排名并过滤：给每条记录添加排名标签，筛选出排名≤N的结果

完整代码示例 #

下面是包含模拟数据的可运行示例，直接对应你的需求场景：

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions._
import org.apache.spark.sql.expressions.Window

// 初始化SparkSession（生产环境请移除master配置）
val spark = SparkSession.builder()
  .appName("TopNSimilarItems")
  .master("local[*]")
  .getOrCreate()

// 模拟你的DataFrame结构：(item_id1, item_id2, score)
val sampleData = Seq(
  ("itemA", "itemX", 0.95),
  ("itemA", "itemY", 0.92),
  ("itemA", "itemZ", 0.88),
  ("itemB", "itemP", 0.98),
  ("itemB", "itemQ", 0.91),
  ("itemB", "itemR", 0.90)
)
val df = spark.createDataFrame(sampleData).toDF("item_id1", "item_id2", "score")

// 定义要获取的Top N数量
val topN = 2

// 1. 构建窗口：按item_id1分组，按score降序排序
val windowSpec = Window
  .partitionBy("item_id1")
  .orderBy(col("score").desc)

// 2. 添加排名列、过滤Top N、移除多余的排名列
val topNDF = df
  .withColumn("rank", row_number().over(windowSpec)) // 给每组内记录分配唯一排名
  .filter(col("rank") <= topN)
  .drop("rank") // 不需要排名列时可以删除

// 查看最终结果
topNDF.show()

运行后会输出每个item_id1的前2条最高得分记录：

+--------+--------+-----+
|item_id1|item_id2|score|
+--------+--------+-----+
|   itemA|   itemX| 0.95|
|   itemA|   itemY| 0.92|
|   itemB|   itemP| 0.98|
|   itemB|   itemQ| 0.91|
+--------+--------+-----+

关键细节说明 #

• 选择合适的排名函数：
  • row_number()：给每组内记录分配唯一序号，即使分数相同也会有不同排名，严格返回Top N条（不会超量）
  • rank()：如果有相同分数的记录，会分配相同排名，比如两个记录分数都是0.95，都会排第1，此时可能返回超过N条的结果（适合允许同分数项都保留的场景）
• 性能优化提示：如果处理超大数据集，可以提前对数据进行过滤，或者给partitionBy的字段设置合理的分区数，减少窗口函数的计算压力

内容的提问来源于stack exchange，提问作者Daniil Andreyevich Baunov