Блог

Главная - Вопросы по программированию - Scala DataFrame — как печатать только строки с наибольшими значениями

Scala DataFrame — как печатать только строки с наибольшими значениями

#scala #dataframe #apache-spark #apache-spark-sql

#scala #фрейм данных #apache-spark #apache-spark-sql

Вопрос:

У меня есть следующий фрейм данных df :

  ---------- -------- --------- 
|        ID| text   |    count|
 ---------- -------- --------- 
|         3|    word|      316|
|         3|    work|      385|
|         3|    want|      205|
|         3|     cat|      251|
|         1|  office|      343|
|         1|     sky|      643|
|         1|   going|      126|
|         2|    home|      124|
|         2|  school|       23|
|         2|   sleep|      103|
//and so on

Теперь для каждого ID я хочу отображать только строки с наибольшими 2 counts и удалять / скрывать остальные:

  ---------- -------- --------- 
|        ID| text   |    count|
 ---------- -------- --------- 
|         3|    word|      316|
|         3|    work|      385|
|         1|  office|      343|
|         1|     sky|      643|
|         2|    home|      124|
|         2|   sleep|      103|
//and so on

Как нам добиться этого наиболее эффективно?

Ответ №1:

Используйте функцию window в spark и partitionBy ID orderBy далее count .

Пример:

 val df=Seq((3,"word",316),(3,"work",385),(3,"want",205),(3,"cat",251),(1,"office",343),(1,"sky",643),(1,"going",126),(2,"home",124),(2,"school",23),(2,"sleep",103)).toDF("ID","text","count")

import org.apache.spark.sql.functions._
import org.apache.spark.sql.expressions._
val w=Window.partitionBy(col("ID")).orderBy(desc("count"))
df.withColumn("rn",row_number().over(w)).filter(col("rn") <=2).drop("rn").show()
// --- ------ ----- 
//| ID|  text|count|
// --- ------ ----- 
//|  1|   sky|  643|
//|  1|office|  343|
//|  3|  work|  385|
//|  3|  word|  316|
//|  2|  home|  124|
//|  2| sleep|  103|
// --- ------ -----

Комментарии:

1. всегда ваше решение простое и понятное .. 🙂

Ответ №2:

Проверьте приведенный ниже код.

 scala> df.show(false)
 --- ------ ----- 
|ID |text  |count|
 --- ------ ----- 
|3  |word  |316  |
|3  |work  |385  |
|3  |want  |205  |
|3  |cat   |251  |
|1  |office|343  |
|1  |sky   |643  |
|1  |going |126  |
|2  |home  |124  |
|2  |school|23   |
|2  |sleep |103  |
 --- ------ ----- 
  
 scala> import org.apache.spark.sql.expressions._

scala> val windowSpec = Window
.partitionBy($"id")
.orderBy($"ID".asc,$"count".desc)
  
 scala> df
.withColumn("rno",row_number().over(windowSpec))
.filter($"rno" <= 2)
.show(false)

 --- ------ ----- --- 
|ID |text  |count|rno|
 --- ------ ----- --- 
|3  |work  |385  |1  |
|3  |word  |316  |2  |
|1  |sky   |643  |1  |
|1  |office|343  |2  |
|2  |school|23   |1  |
|2  |home  |124  |2  |
 --- ------ ----- ---