Spark Scala聚合中如何实现表达式展开 解决:_*注解报错

问题场景 #

我正在尝试将一段简单的聚合代码从PySpark转换为Scala版本。
两种语言下创建测试DataFrame的代码如下：

PySpark 测试代码 #

# PySpark
from pyspark.sql import functions as F
df = spark.createDataFrame(
    [([10, 100],),
     ([20, 200],)],
    ['vals'])

Scala 测试代码 #

// Scala
val df = Seq(
    (Seq(10, 100)),
    (Seq(20, 200)),
).toDF("vals")

PySpark中通过*解包列表推导式生成的聚合表达式传入agg方法可正常运行，代码及运行结果如下：

df2 = df.agg(
    *[F.sum(F.col("vals")[i]).alias(f"col{i}") for i in range(2)]
)
df2.show()
# +----+----+
# |col0|col1|
# +----+----+
# |  30| 300|
# +----+----+

编写等价Scala代码时触发编译报错：

val df2 = df.agg(
  (0 until 2).map(i => sum($"vals"(i)).alias(s"col$i")): _*
)

报错信息：

(0 until 2).map(i => sum($"vals"(i)).alias(s"col$i")): _*
                                                              ^
On line 2: error: no `: _*` annotation allowed here
       (such annotations are only allowed in arguments to *-parameters)

相同的: _*展开语法在select方法中可以正常运行：

val df2 = df.select(
  (0 until 2).map(i => $"vals"(i).alias(s"col$i")): _*
)

原因说明 #

Scala Spark的agg完全支持动态生成的表达式展开，报错的核心原因是两个方法的参数签名存在差异：

• select方法的签名为def select(cols: Column*): DataFrame，直接接收可变长度的Column参数，因此可以直接对Column序列使用:_*展开传参
• agg方法的Column类型重载签名为def agg(expr: Column, exprs: Column*): DataFrame，要求第一个聚合表达式单独传入，剩余表达式才作为可变参数传入，直接对整个序列使用:_*展开无法匹配参数规则，编译器就会抛出错误

正确实现方式 #

只需要将生成的聚合表达式序列拆分，把第一个元素作为单独参数传入，剩余元素展开传入即可匹配方法签名：

import org.apache.spark.sql.functions._
// 动态生成聚合表达式序列
val aggExprs = (0 until 2).map(i => sum($"vals"(i)).alias(s"col$i"))
// 拆分传参：第一个元素单独传，剩余元素用:_*展开
val df2 = df.agg(aggExprs.head, aggExprs.tail: _*)
df2.show()
/*
+----+----+
|col0|col1|
+----+----+
|  30| 300|
+----+----+
*/

如果需要兼容表达式序列为空的边缘场景，可以加一层判断：

val df2 = aggExprs match {
  case head :: tail => df.agg(head, tail: _*)
  case Nil => df
}

内容的提问来源于stack exchange，提问作者ZygD