Возьмите первые п записей из dataframe сгруппированных по уникальному идентификатору

Question

У меня есть набор данных, как этот

Как вы видите, упорядочено по рейтингу и USERID мне нужно, чтобы получить новый Dataframe только с топ 2 результатов из каждой группы по уникальной user_id Я попытался

dataframe.groupBy("user_id").agg(someUdfFuntion) 

Я пытался использовать функцию ранга, но это, кажется, не работает, я пытался фильтровать dataframe, но никакого результата, как я мог это сделать?

Answer 1

Try:

import org.apache.spark.sql.expressions.Window 
import org.apache.spark.sql.functions.row_number 

val window = Window.partitionBy("userId").orderBy($"rating".desc) 

dataframe.withColumn("r", row_number.over(window)).where($"r" <= n)