Spark 获得每个(项目1、项目2、分数)的前 N 个最高分结果

Question

我有一个以下格式的数据帧：

item_id1: Long, item_id2: Long, similarity_score: Double

我想做的是为每个item_id1获得前 N 条最高 similarity_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对记录进行分组，然后按分数反向排序，然后限制结果来完成。但是我坚持如何在Spark Scala中实现它。

谢谢。

Answer 1

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

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

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