Как применить функцию Window к нескольким столбцам в DataFrame

Question

У меня есть следующий DataFrame df:

Id   label   field1   field2
1    xxx     2        3
1    yyy     1        5
2    aaa     0        10
1    zzz     2        6

Для каждого уникального Id Я хочу знать label с самыми высокими field1 и field2.

Ожидаемый результат:

Id   labelField1   lableLield2
1    xxx           zzz
2    aaa           aaa

Я знаю, как это сделать, если бы у меня было только labelField1 или labelField2.Но я не уверен, как лучше всего работать с обоими лейблами.

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

Вы можете комбинировать встроенные функции struct и max для достижения ваших требований как

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)

, что должно дать вам

+---+-----------+-----------+
|Id |labelField1|labelField2|
+---+-----------+-----------+
|1  |xxx        |zzz        |
|2  |aaa        |aaa        |
+---+-----------+-----------+

Примечание: в случае связь как и в field1 для Id = 1 существует связь между xxx и zzz, поэтому будет выбран случайный