在下面的示例中,我希望只能采用计数最高的x个ID。 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 |
+------+-------+-----+
然后我要删除count列,但这很简单。
我有办法做到这一点,但是我认为这完全违背了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
答案 0 :(得分: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")
)