Récupérer top n dans chaque groupe d'un DataFrame dans pyspark

Il y a un DataFrame dans pyspark avec les données ci-dessous:

user_id object_id score
user_1  object_1  3
user_1  object_1  1
user_1  object_2  2
user_2  object_1  5
user_2  object_2  2
user_2  object_2  6

Ce que j'attends, c'est de revenir 2 enregistrements dans chaque groupe avec le même user_id, qui doivent avoir le score le plus élevé. En conséquence, le résultat devrait ressembler à la suivante:

user_id object_id score
user_1  object_1  3
user_1  object_2  2
user_2  object_2  6
user_2  object_1  5

Je suis vraiment nouveau à pyspark, quelqu'un pourrait-il me donner un extrait de code ou de portail de la documentation liée à ce problème? Un grand merci!

source d'informationauteur KAs

  1. 27

    Je crois que vous devez utiliser les fonctions de la fenêtre atteindre le rang de chaque ligne en fonction de user_id et scoreet ensuite filtrer vos résultats afin de garder uniquement les deux premières valeurs.

    from pyspark.sql.window import Window
from pyspark.sql.functions import rank, col

window = Window.partitionBy(df['user_id']).orderBy(df['score'].desc())

df.select('*', rank().over(window).alias('rank')) 
  .filter(col('rank') <= 2) 
  .show() 
#+-------+---------+-----+----+
#|user_id|object_id|score|rank|
#+-------+---------+-----+----+
#| user_1| object_1|    3|   1|
#| user_1| object_2|    2|   2|
#| user_2| object_2|    6|   1|
#| user_2| object_1|    5|   2|
#+-------+---------+-----+----+

    En général, le fonctionnaire guide de programmation est un bon endroit pour commencer à apprendre Étincelle.

    De données

    rdd = sc.parallelize([("user_1",  "object_1",  3), 
                      ("user_1",  "object_2",  2), 
                      ("user_2",  "object_1",  5), 
                      ("user_2",  "object_2",  2), 
                      ("user_2",  "object_2",  6)])
df = sqlContext.createDataFrame(rdd, ["user_id", "object_id", "score"])
  2. 7

    Haut n est plus précise si l'aide row_number au lieu de rank lors de l'obtention de grade de l'égalité:

    val n = 5
df.select(col('*'), row_number().over(window).alias('row_number')) \
  .where(col('row_number') <= n) \
  .limit(20) \
  .toPandas()

    Note limit(20).toPandas() truc au lieu de show() pour Jupyter cahiers pour la belle mise en forme.