使用Spark DataFrame获取一组之后所有组的TopN

Question

我有一个Spark SQL DataFrame:

user1 item1 rating1
user1 item2 rating2
user1 item3 rating3
user2 item1 rating4
...

如何使用Scala按用户分组，然后从每个组返回TopN项？

使用Python的相似代码：

df.groupby("user").apply(the_func_get_TopN)

Answer 1

您可以使用rank窗口函数，如下所示

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions.{rank, desc}
val n: Int = ???
// Window definition
val w = Window.partitionBy($"user").orderBy(desc("rating"))
// Filter
df.withColumn("rank", rank.over(w)).where($"rank" <= n)

如果您不关心关系，则可以用row_number 替换rank