Spark 获取每个（item1，item2，score）的前 N ​​个最高分结果

Posted 2023-04-15

【中文标题】Spark 获取每个（item1，item2，score）的前 N ​​个最高分结果

【英文标题】：Spark get top N highest score results for each (item1, item2, score)

【发布时间】：2018-03-09 19:06:10

【问题描述】：

我有一个以下格式的 DataFrame：

item_id1: Long, item_id2: Long, similarity_score: Double

我想要做的是为每个 item_id1 获取前 N 个最高相似性_score 记录。 所以，例如：

1 2 0.5
1 3 0.4
1 4 0.3
2 1 0.5
2 3 0.4
2 4 0.3

前 2 个相似的项目会给出：

1 2 0.5
1 3 0.4
2 1 0.5
2 3 0.4

我隐约猜测可以先按item_id1对记录进行分组，再按score逆序排序，再限制结果。但我不知道如何在 Spark Scala 中实现它。

谢谢。

【参考方案1】：

我建议为此使用窗口函数：

 df
  .withColumn("rnk",row_number().over(Window.partitionBy($"item_id1").orderBy($"similarity_score")))
  .where($"rank"<=2)

或者，您可以使用dense_rank/rank 而不是row_number，具体取决于如何处理相似度分数相等的情况。

【讨论】：

谢谢！这正是我一直在寻找的。需要在 Window 上阅读更多内容。