Есть ли лучший способ выполнить этот процесс надлежащей обрезки моего искрового DataFrame?

Question

В следующем примере я хочу иметь возможность использовать только х идентификаторы с наибольшим количеством.x - это число, которое я хочу, которое определяется переменной с именем howMany.

Для следующего примера с данным кадром данных:

+------+--+-----+
|query |Id|count|
+------+--+-----+
|query1|11|2    |
|query1|12|1    |
|query2|13|2    |
|query2|14|1    |
|query3|13|2    |
|query4|12|1    |
|query4|11|1    |
|query5|12|1    |
|query5|11|2    |
|query5|14|1    |
|query5|13|3    |
|query6|15|2    |
|query6|16|1    |
|query7|17|1    |
|query8|18|2    |
|query8|13|3    |
|query8|12|1    |
+------+--+-----+

Я хотел бы получить следующий кадр данных, если номер переменной равен 2.

+------+-------+-----+
|query |Ids    |count|
+------+-------+-----+
|query1|[11,12]|2    |
|query2|[13,14]|2    |
|query3|[13]   |2    |
|query4|[12,11]|1    |
|query5|[11,13]|2    |
|query6|[15,16]|2    |
|query7|[17]   |1    |
|query8|[18,13]|2    |
+------+-------+-----+

I тогдаЯ хочу удалить столбец подсчета, но это тривиально.

У меня есть способ сделать это, но я думаю, что это полностью противоречит цели scala и полностью тратит много времени выполнения.Будучи новичком, я не уверен в том, как лучше всего это сделать

Мой текущий метод заключается в том, чтобы сначала получить отдельный список столбца запроса и создать итератор.Во-вторых, я перебираю список с помощью итератора и обрезаю фрейм данных только до текущего запроса в списке, используя df.select ($ "eachColumnName" ...). Where ("query" .equalTo (iter.next ())),Затем я .limit (howMany), а затем groupBy ($ "query"). Agg (collect_list ($ "Id"). As ("Ids")).Наконец, у меня есть пустой фрейм данных, и я добавляю каждый из них по одному в пустой фрейм данных и возвращаю этот вновь созданный фрейм данных.

df.select($"query").distinct().rdd.map(r => r(0).asInstanceOf[String]).collect().toList
val iter = queries.toIterator
while (iter.hasNext) {
    middleDF = df.select($"query", $"Id", $"count").where($"query".equalTo(iter.next()))
    queryDF = middleDF.sort(col("count").desc).limit(howMany).select(col("query"), col("Ids")).groupBy(col("query")).agg(collect_list("Id").as("Ids"))
    emptyDF.union(queryDF) // Assuming emptyDF is made
}
emptyDF

Answer 1

Я бы сделал это, используя Window-Functions для получения ранга, затем groupBy для агрегирования:

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions._

val howMany = 2

val newDF = df
.withColumn("rank",row_number().over(Window.partitionBy($"query").orderBy($"count".desc)))
.where($"rank"<=howMany)
.groupBy($"query")
.agg(
 collect_list($"Id").as("Ids"),
 max($"count").as("count") 
)