如何使用pyspark按小时获取x个最频繁的位置？

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。之后，对前2行进行过滤。

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()