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

Question

我有一个数据框架，按年份列出名称和分数，如下所示：

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")

我的问题是：有可能用更少的步骤实现同样的目标吗？也许只使用窗口函数，但没有附加的联接？

Answer 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|
+---+--------+-----+----+