How to select the first row of each group?

Question

I have a DataFrame generated as follow:

df.groupBy($"Hour", $"Category")
  .agg(sum($"value") as "TotalValue")
  .sort($"Hour".asc, $"TotalValue".desc))

The results look like:

+----+--------+----------+
|Hour|Category|TotalValue|
+----+--------+----------+
|   0|   cat26|      30.9|
|   0|   cat13|      22.1|
|   0|   cat95|      19.6|
|   0|  cat105|       1.3|
|   1|   cat67|      28.5|
|   1|    cat4|      26.8|
|   1|   cat13|      12.6|
|   1|   cat23|       5.3|
|   2|   cat56|      39.6|
|   2|   cat40|      29.7|
|   2|  cat187|      27.9|
|   2|   cat68|       9.8|
|   3|    cat8|      35.6|
| ...|    ....|      ....|
+----+--------+----------+

As you can see, the DataFrame is ordered by Hour in an increasing order, then by TotalValue in a descending order.

I would like to select the top row of each group, i.e.

from the group of Hour==0 select (0,cat26,30.9)
from the group of Hour==1 select (1,cat67,28.5)
from the group of Hour==2 select (2,cat56,39.6)
 

and so on
So the desired output would be:

+----+--------+----------+
|Hour|Category|TotalValue|
+----+--------+----------+
|   0|   cat26|      30.9|
|   1|   cat67|      28.5|
|   2|   cat56|      39.6|
|   3|    cat8|      35.6|
| ...|     ...|       ...|
+----+--------+----------+

Answer 1

According to your question, we can select the first row of each group using various methods below:

import org.apache.spark.sql.functions.{row_number, max, broadcast}

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

val df = sc.parallelize(Seq(

  (0,"cat26",30.9), (0,"cat13",22.1), (0,"cat95",19.6), (0,"cat105",1.3),

  (1,"cat67",28.5), (1,"cat4",26.8), (1,"cat13",12.6), (1,"cat23",5.3),

  (2,"cat56",39.6), (2,"cat40",29.7), (2,"cat187",27.9), (2,"cat68",9.8),

  (3,"cat8",35.6))).toDF("Hour", "Category", "TotalValue")

Method1: 

val w = Window.partitionBy($"hour").orderBy($"TotalValue".desc)

val dfTop = df.withColumn("rn", row_number.over(w)).where($"rn" === 1).drop("rn")

dfTop.show

// +----+--------+----------+

// |Hour|Category|TotalValue|

// +----+--------+----------+

// |   0|   cat26|      30.9|

// |   1|   cat67|      28.5|

// |   2|   cat56|      39.6|

// |   3|    cat8|      35.6|

// +----+--------+----------+

Method2: Using dataset API 

case class Record(Hour: Integer, Category: String, TotalValue: Double)

df.as[Record]

  .groupByKey(_.Hour)

  .reduceGroups((x, y) => if (x.TotalValue > y.TotalValue) x else y)

  .show

If you want to know more about Spark, then do check out this awesome video tutorial: