Scala DataFrame - How to only print rows with largest values

Question

I have following Dataframe 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

Now, for each ID, I want to only display the rows with the largest 2 counts and drop/hide the rest:

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

How do we achieve this most efficiently?

Answer 1

Use window function in spark and partitionBy on ID orderBy on count.

Example:

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|
//+---+------+-----+

Answer 2

Check below code.

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  |
+---+------+-----+---+