使用 Spark DataFrame 获取一组后所有组的 TopN [英] get TopN of all groups after group by using Spark DataFrame

问题描述

我有一个 Spark SQL 数据帧:

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)

如果你不关心平局，那么你可以用 row_number

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

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