使用 Spark DataFrame 获取所有组的 TopN

【问题标题】：get TopN of all groups after group by using Spark DataFrame使用 Spark DataFrame 获取所有组的 TopN
【发布时间】：2016-02-12 20:01:30
【问题描述】：

我有一个 Spark SQL 数据框：

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

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

使用 Python 的相似性代码：

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

【问题讨论】：

标签： sql scala apache-spark apache-spark-sql

【解决方案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)

如果您不关心关系，那么您可以将rank 替换为row_number

【讨论】：