如何将特定列的相同值分配给窗口分区中的所有行

Posted 2023-03-31

【中文标题】如何将特定列的相同值分配给窗口分区中的所有行

【英文标题】：How to assign the same value for a specific column to all rows in a Window Partition

【发布时间】：2021-06-07 11:29:47

【问题描述】：

我有一个数据框，其名称和分数按年份如下：

val originalDF = Seq(
     |     (1, "Arthur", 5, "2010"),
     |     (2, "Arthur", 15, "2012"),
     |     (3, "Arthur", 10, "2017"),
     |     (4, "Trevor", 30, "2014"),
     |     (5, "Trevor", 12, "2015"),
     |     (6, "Franklin", 15, "2016"),
     |     (7, "Franklin", 20, "2018")
     |   ).toDF("id", "user", "score", "year")

返回：

+---+--------+-----+----+
| id|    user|score|year|
+---+--------+-----+----+
|  1|  Arthur|    5|2010|
|  2|  Arthur|   15|2012|
|  3|  Arthur|   10|2017|
|  4|  Trevor|   30|2014|
|  5|  Trevor|   12|2015|
|  6|Franklin|   15|2016|
|  7|Franklin|   20|2018|
+---+--------+-----+----+

我想要为每一行保留id、user 和score，但根据他的最高历史分数，年份列应显示user 的best_year。那将是：

+---+--------+-----+---------+
| id|    name|score|best_year|
+---+--------+-----+---------+
|  1|  Arthur|    5|     2012|
|  2|  Arthur|   15|     2012|
|  3|  Arthur|   10|     2012|
|  4|  Trevor|   30|     2014|
|  5|  Trevor|   12|     2014|
|  6|Franklin|   15|     2018|
|  7|Franklin|   20|     2018|
+---+--------+-----+---------+

到目前为止，我想出的解决这个问题的方法需要生成一个新的 DataFrame，其中包含按用户分组的最高分数，最后将这个新的 DataFrame 与原来的 DataFrame 连接起来：

val wind = Window.partitionBy("user").orderBy(col("score").desc)

val bestScores = originalDF.withColumn("rank", row_number over wind).where($"rank" === 1)    
     
val solution = originalDF.alias("o").join(bestScores.alias("b"), $"o.user" === $"b.user").select($"o.id", $"o.user", $"o.score", $"b.year")

我的问题：是否可以用更少的步骤实现相同的目标？也许只使用窗口函数但没有额外的连接？

【参考方案1】：

您可以使用结构的最大窗口函数根据最高分数找到年份：

val df2 = originalDF.withColumn(
    "year", 
    max(struct(col("score"),col("year"))).over(Window.partitionBy("user"))("year")
)

df2.show
+---+--------+-----+----+
| id|    user|score|year|
+---+--------+-----+----+
|  1|  Arthur|    5|2012|
|  2|  Arthur|   15|2012|
|  3|  Arthur|   10|2012|
|  4|  Trevor|   30|2014|
|  5|  Trevor|   12|2014|
|  6|Franklin|   15|2018|
|  7|Franklin|   20|2018|
+---+--------+-----+----+