How to apply Window function to multiple columns in DataFrame - scala

Question

I have the following DataFrame df:
Id label field1 field2
1 xxx 2 3
1 yyy 1 5
2 aaa 0 10
1 zzz 2 6
For each unique Id I want to know the label with highest field1 and field2.
Expected result:
Id labelField1 lableLield2
1 xxx zzz
2 aaa aaa
I know how to do it if I would only have labelField1 or labelField2.
But I am not sure what is the best way to deal with both labels.
val w1 = Window.partitionBy($"Id").orderBy($"field1".desc)
val w2 = Window.partitionBy($"Id").orderBy($"field2".desc)
val myLabels = df.select("Id", "label", "field1", "field2")
.withColumn("rn", row_number.over(w1)).where($"rn" === 1)
.drop("rn")
.drop("field1")

Answer 1

You can combine struct and max inbuild functions to achieve your requirement as
import org.apache.spark.sql.functions._
df.groupBy("Id")
.agg(max(struct("field1", "label")).as("temp1"), max(struct("field2", "label")).as("temp2"))
.select(col("Id"), col("temp1.label").as("labelField1"), col("temp2.label").as("labelField2"))
.show(false)
which should give you
+---+-----------+-----------+
|Id |labelField1|labelField2|
+---+-----------+-----------+
|1 |xxx |zzz |
|2 |aaa |aaa |
+---+-----------+-----------+
Note: In case of tie as in field1 for Id=1 there is tie between xxx and zzz so random will be chosen