Как взять x наиболее частое место по часам с pyspark?

Question

У меня есть фрейм данных искры с часом дня, locationID и частотой. Частота - это количество идентификаторов locationID в этот час дня.

+----+----------+---------+
|hour|locationID|frequency|
+----+----------+---------+
| 0  |     1    |    20   |
| 0  |     2    |    11   |   
| 0  |     9    |     6   |
| 1  |     3    |    32   |
| 1  |     1    |    22   |
| 1  |     5    |     4   |

Я хочу взять 2 самых частых locationID в час.

Answer 1

Это можно сделать с помощью оконной функции row_number.Окно группируется по hour и упорядочивает frequency в порядке убывания.После этого отфильтруйте две верхние строки.

from pyspark.sql import Window
from pyspark.sql.functions import row_number,desc
w = Window.partitionBy(df.hour).orderBy(df.frequency.desc())
rnum_df = df.withColumn('rnum',row_number().over(w))
rnum_df.filter(rnum_df.rnum <= 2).show()