通过使用Spark DataFrame获得分组后所有分组的TopN [英] get TopN of all groups after group by using Spark DataFrame

问题描述

我有一个Spark SQL DataFrame:

I have a Spark SQL DataFrame:

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

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

How to group by user and then return TopN items from every group using Scala?

使用Python的相似代码:

Similarity code using Python:

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

推荐答案

您可以按以下方式使用rank窗口函数

You can use rank window function as follows

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

If you don't care about ties then you can replace rank with row_number

这篇关于通过使用Spark DataFrame获得分组后所有分组的TopN的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！