时间戳

首页 > 解决方案 > 如何将 Window 函数应用于 DataFrame 中的多个列

scala - 如何将 Window 函数应用于 DataFrame 中的多个列

问题描述

我有以下 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

如果我只有labelField1or ，我知道该怎么做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")

标签： scalaapache-sparkapache-spark-sql

解决方案

您可以组合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        |
+---+-----------+-----------+

注意：如果Id=1中的平局，则两者之间存在平局，因此将选择随机field1xxxzzz

推荐阅读