Получите первые n в каждой группе DataFrame в pyspark

В pyspark есть DataFrame с данными, как показано ниже:

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

Я ожидаю, что в каждой группе будет возвращено 2 записи с одинаковым идентификатором user_id, которые должны иметь самый высокий балл. Следовательно, результат должен выглядеть следующим образом:

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

Я действительно новичок в pyspark. Может ли кто-нибудь дать мне фрагмент кода или портал с соответствующей документацией по этой проблеме? Большое спасибо!

Ответы на вопрос(2)

Решение Вопроса

что вам нужно использоватьоконные функции достичь ранга каждой строки на основеuser_id а такжеscoreи впоследствии отфильтруйте результаты, чтобы сохранить только первые два значения.

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|
#+-------+---------+-----+----+

В общем, официальныйруководство по программированию это хорошее место, чтобы начать изучать Spark.

Данные
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"])
 mtoto05 окт. 2016 г., 10:37
@ Уилмертон, который не работает.
 Wilmerton05 окт. 2016 г., 09:56
Вы можете использовать оконную функцию в фильтре:df.filter(rank().over(window) <= 2)
 Wilmerton05 окт. 2016 г., 20:17
Я ошеломлен ... Я был уверен, что раньше использовал функцию окна в фильтре. Но я действительно не мог воспроизвести его (ни в 2, ни в 1.6). Я использовал это экзотическим способом, но я не могу вспомнить, когда и как. Сожалею.
 Tomer Ben David19 дек. 2018 г., 13:49
Вы можете рассмотреть возможность использованияrow_number вместоrank в случае получения того же ранга, и вы все еще хотите топ н
 KAs15 июл. 2016 г., 16:33
Я думаю, что-то нужно настроить. object_id не влияет ни наgroupby или жеtop процедура. И я хочу, чтобыgroup by user_id, и в каждой группе получить первыйдва записи с наибольшим количеством баллов отдельно, а не только первые записи. Большое спасибо!

Топ-н более точен при использованииrow_number вместоrank при получении ранга равенства:

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

Заметкаlimit(20).toPandas() трюк вместоshow() для ноутбуков Jupyter для лучшего форматирования.

Ваш ответ на вопрос