如何用Scala扩展Spark SQL内置聚合函数？（实现dollarSum）

环境说明 #

• Databricks Runtime 10.4 LTS ML
• Spark 3.2.1
• Scala 2.12

可行方案一：基于DeclarativeAggregate实现 #

针对你提到的需求适配DeclarativeAggregate的场景，可基于Spark内置Sum的核心逻辑扩展，仅在最终结果阶段添加ROUND操作，完整实现代码如下：

import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.expressions.aggregate.DeclarativeAggregate
import org.apache.spark.sql.types._
import org.apache.spark.sql.catalyst.util.TypeUtils

case class DollarSum(child: Expression) extends DeclarativeAggregate {
  // 输入数据类型
  override def inputTypes: Seq[DataType] = Seq(child.dataType)
  // 输出数据类型
  override def dataType: DataType = DoubleType
  // 确定性标识
  override def deterministic: Boolean = true

  // 定义聚合缓冲区的类型与引用
  private lazy val sumDataType = TypeUtils.getNumericType(child.dataType)
  private lazy val sum = AttributeReference("sum", sumDataType)()

  override lazy val aggBufferAttributes: Seq[AttributeReference] = sum :: Nil

  // 缓冲区初始化表达式
  override lazy val initialValues: Seq[Expression] = Seq(
    Literal.default(sumDataType)
  )

  // 缓冲区更新逻辑：复用内置Sum的累加逻辑
  override lazy val updateExpressions: Seq[Expression] = Seq(
    Add(sum, Coalesce(child :: Literal.default(sumDataType) :: Nil, sumDataType))
  )

  // 缓冲区合并逻辑：复用内置Sum的合并逻辑
  override lazy val mergeExpressions: Seq[Expression] = Seq(
    Add(sum.left, sum.right)
  )

  // 最终结果计算：对sum结果执行ROUND(...,2)
  override lazy val evaluateExpression: Expression = {
    Round(Cast(sum, DoubleType), Literal(2))
  }
}

// 注册为Spark SQL可调用的函数
spark.udf.register("dollarSum", (col: Column) => Column(DollarSum(col.expr)))

代码说明 #

• 完全复用内置Sum的初始化、更新、合并逻辑，仅在最终计算阶段添加ROUND操作，避免重复造轮子
• 自动适配所有数值类型的输入列，无需额外处理类型兼容问题
• 最终返回保留两位小数的Double类型结果

可行方案二：组合内置函数注册自定义聚合函数 #

你之前尝试用functions.sum和round组合失败，核心原因是注册方式错误。以下是无需编写UDAF源码的简化方案：

import org.apache.spark.sql.functions._

// 注册自定义聚合函数
spark.sqlContext.udf.register("dollarSum", (col: Column) => round(sum(col), 2))

// SQL中使用示例
spark.sql("SELECT dollarSum(amount) FROM sales_table").show()

补充使用方式（DataFrame API） #

如果需要在DataFrame链式调用中使用，可直接定义工具方法：

def dollarSum(col: Column): Column = round(sum(col), 2)

// DataFrame中使用示例
df.select(dollarSum($"amount")).show()

之前四种方法失败的原因 #

• 继承Aggregator类：Aggregator是强类型UDAF实现，需要严格匹配输入、缓冲区、输出的类型映射，你可能在类型转换或finish方法逻辑上出现错误
• 复制修改Sum类源码：Spark内置Sum类依赖大量私有API与内部实现细节，直接复制会因访问权限限制失败
• 模仿try_sum的ExpressionBuilder：ExpressionBuilder是Spark内部私有类，不对外暴露，无法直接复用
• 组合函数注册失败：之前的注册方式错误套用了普通UDF的注册逻辑，未适配聚合函数的特性，上述方案二修正了该问题

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