Используйте Iterator, чтобы получить лучшие k ключевых слов

Question

Я пишу алгоритм Spark, чтобы получить топ-k ключевых слов для каждой страны, теперь у меня уже есть Dataframe, содержащий все записи, и я планирую сделать

df.repartition($"country_id").mapPartition()

, чтобы получить топ-k ключевых слов, но я не совсем понимаю, какЯ мог бы написать итератор, чтобы получить его.

Если я могу написать метод или вызвать собственный метод, я могу отсортировать по каждому разделу и получить top k, который, кажется, не является правильным подходом, если вводитератор.

У кого-нибудь есть идеи?

Answer 1

вы можете добиться этого с помощью оконных функций, давайте предположим, что столбец _1 является вашим ключевым словом, а _2 является количеством ключевых слов. В этом случае k = 2

scala> df.show()
+---+---+
| _1| _2|
+---+---+
|  1|  3|
|  2|  2|
|  1|  4|
|  1|  1|
|  2|  0|
|  1| 10|
|  2|  5|
+---+---+

scala> df.select('*,row_number().over(Window.orderBy('_2.desc).partitionBy('_1)).as("rn")).where('rn < 3).show()
+---+---+---+
| _1| _2| rn|
+---+---+---+
|  1| 10|  1|
|  1|  4|  2|
|  2|  5|  1|
|  2|  2|  2|
+---+---+---+