Spark -添加一列，该列对另一列进行求和，并按第三列分组，而不会丢失其他列

Question

我有一个包含5列的DataFrame：col11、col2、col3、col4和col5

我想添加col6，它将是col5的总和，按col1分组。但我不想失去其他栏目。

如果我这样做了：

df
  .groupBy("col1")
  .agg(sum("col5") as "col6")

然后我失去了第2-4列。

我可以通过运行以下命令来执行连接：

val sumValues = df
  .groupBy("col1")
  .agg(sum("col5") as "col6")

df
  .join(sumValues, Seq("col1"))

但这感觉有点过头了。

我希望做一些像这样的事情：

df
  .withGroupedColumn("col6", "col1", sum("col5") as "col6")

在Spark中有没有简单的方法可以做到这一点？

Answer 1

您可以使用窗口函数：

val df2 = df.withColumn("col6", expr("sum(col5) over (partition by col1)"))

或等同于

import org.apache.spark.sql.expressions.Window

val df2 = df.withColumn("col6", sum("col5").over(Window.partitionBy("col1")))