1. 为什么需要自定义聚合函数
在日常的数据处理中，我们经常遇到内置函数无法满足需求的场景。比如计算加权平均值、统计特定条件下的百分位数，或者实现复杂的业务逻辑聚合。这时候就需要自己动手写用户自定义聚合函数（UDAF）。
我去年做过一个电商用户行为分析项目，需要计算用户浏览商品时的"深度浏览率"（浏览时长超过30秒的比例）。这个指标用内置的avg函数根本无法实现，最终就是靠自定义聚合函数搞定的。Spark SQL的UDAF机制给了我们极大的灵活性，但很多新手在实现时容易踩坑。
2. 环境准备与数据示例
2.1 基础环境搭建
首先确保你的开发环境已经配置好：
- JDK 1.8+
- Scala 2.12.x
- Spark 3.x
建议使用sbt构建项目，build.sbt中需要包含这些依赖：
libraryDependencies ++= Seq( "org.apache.spark" %% "spark-core" % "3.3.0", "org.apache.spark" %% "spark-sql" % "3.3.0" )2.2 准备测试数据
我们用这个JSON文件作为示例数据（保存为users.json）：
[ {"name":"张三", "age":18, "score":85.5}, {"name":"李四", "age":15, "score":92.0}, {"name":"王五", "age":20, "score":78.5} ]这个数据集包含了学生的姓名、年龄和考试成绩，后面我们会用它来演示如何实现自定义的平均值计算。
3. UDAF核心实现详解
3.1 继承UserDefinedAggregateFunction
所有自定义聚合函数都必须继承UserDefinedAggregateFunction抽象类。它要求实现6个关键方法：
class MyAvg extends UserDefinedAggregateFunction { // 输入数据类型 override def inputSchema: StructType = ??? // 缓冲区数据类型 override def bufferSchema: StructType = ??? // 最终输出类型 override def dataType: DataType = ??? // 是否幂等 override def deterministic: Boolean = ??? // 初始化缓冲区 override def initialize(buffer: MutableAggregationBuffer): Unit = ??? // 分区内聚合 override def update(buffer: MutableAggregationBuffer, input: Row): Unit = ??? // 分区间合并 override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = ??? // 计算结果 override def evaluate(buffer: Row): Any = ??? }3.2 实现inputSchema和bufferSchema
inputSchema定义输入数据的结构。对于计算平均年龄的场景：
override def inputSchema: StructType = StructType(StructField("age", DoubleType) :: Nil)bufferSchema定义中间计算结果的结构。计算平均值需要保存总和和计数：
override def bufferSchema: StructType = StructType( StructField("sum", DoubleType) :: StructField("count", LongType) :: Nil )3.3 initialize方法实现
这个方法初始化聚合缓冲区。对于平均值计算，我们需要初始化总和为0，计数为0：
override def initialize(buffer: MutableAggregationBuffer): Unit = { buffer(0) = 0.0 // sum初始化为0.0 buffer(1) = 0L // count初始化为0L }注意这里必须用0L而不是0，否则会被识别为Int类型导致类型不匹配错误。
4. 聚合逻辑实现
4.1 update方法：分区内聚合
update方法处理每条输入数据。对于我们的例子：
override def update(buffer: MutableAggregationBuffer, input: Row): Unit = { if (!input.isNullAt(0)) { // 处理null值 val age = input.getDouble(0) buffer(0) = buffer.getDouble(0) + age // 累加年龄 buffer(1) = buffer.getLong(1) + 1L // 计数加1 } }这里添加了null值检查，避免空指针异常。实际项目中这是必须考虑的边缘情况。
4.2 merge方法：分区间合并
当Spark合并不同分区的中间结果时调用merge方法：
override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = { buffer1(0) = buffer1.getDouble(0) + buffer2.getDouble(0) // 合并sum buffer1(1) = buffer1.getLong(1) + buffer2.getLong(1) // 合并count }这个方法把多个分区的中间结果合并成一个最终结果。
4.3 evaluate方法：最终计算
在所有数据聚合完成后，调用evaluate生成最终结果：
override def evaluate(buffer: Row): Any = { if (buffer.getLong(1) == 0L) null // 处理除零情况 else buffer.getDouble(0) / buffer.getLong(1) }这里添加了除零保护，当没有输入数据时返回null而不是抛出异常。
5. 注册和使用UDAF
5.1 完整代码实现
把前面的代码组合起来：
class MyAvg extends UserDefinedAggregateFunction { // 省略之前的方法实现... } object UDAFDemo { def main(args: Array[String]): Unit = { val spark = SparkSession.builder() .appName("UDAF Demo") .master("local[*]") .getOrCreate() import spark.implicits._ val df = spark.read.json("users.json") df.createOrReplaceTempView("students") spark.udf.register("myAvg", new MyAvg) spark.sql("SELECT myAvg(age) FROM students").show() spark.stop() } }5.2 执行结果分析
运行程序会输出类似这样的结果：
+------------------+ |myavg(CAST(age AS DOUBLE))| +------------------+ | 17.666666666666668| +------------------+这个结果是我们三个样本年龄(18,15,20)的平均值。注意Spark自动将整型的age转换为Double进行计算。
6. 高级技巧与性能优化
6.1 处理复杂数据类型
UDAF不仅支持基本类型，还可以处理Array、Map等复杂类型。比如计算数组列的平均值：
override def inputSchema: StructType = StructType(StructField("array_col", ArrayType(DoubleType)) :: Nil) override def update(buffer: MutableAggregationBuffer, input: Row): Unit = { val array = input.getAs[Seq[Double]](0) buffer(0) = buffer.getDouble(0) + array.sum buffer(1) = buffer.getLong(1) + array.length }6.2 性能优化建议
- 减少中间数据：bufferSchema尽量简单，避免存储不必要的数据
- 使用原生类型：优先使用Int/Long/Double等原生类型而非对象
- 避免装箱拆箱：注意Scala和Java类型转换的开销
- 合理设置分区数：数据量大的时候适当增加分区数
7. 常见问题排查
7.1 类型不匹配错误
最常见的错误是buffer数据类型与schema声明不一致。比如：
java.lang.ClassCastException: java.lang.Integer cannot be cast to java.lang.Long这是因为在initialize或update中使用了Int而不是Long。解决方法是在数字字面量后加L：
buffer(1) = 0L // 正确 buffer(1) = 0 // 错误7.2 空指针异常
没有正确处理null值会导致运行时异常。建议：
- 在update方法中添加null检查
- 在evaluate方法中添加除零保护
- 使用Option类型处理可能为null的值
8. 实际项目中的应用扩展
8.1 加权平均值实现
实际项目中经常需要计算加权平均值。只需要稍作修改：
class WeightedAvg extends UserDefinedAggregateFunction { override def inputSchema: StructType = StructType(StructField("value", DoubleType) :: StructField("weight", DoubleType) :: Nil) override def bufferSchema: StructType = StructType(StructField("sum", DoubleType) :: StructField("weightSum", DoubleType) :: Nil) override def evaluate(buffer: Row): Any = buffer.getDouble(0) / buffer.getDouble(1) // 其他方法实现类似... }8.2 条件聚合实现
有时需要对满足特定条件的数据进行聚合。例如只统计及格分数的平均值：
override def update(buffer: MutableAggregationBuffer, input: Row): Unit = { val score = input.getDouble(0) if (score >= 60.0) { buffer(0) = buffer.getDouble(0) + score buffer(1) = buffer.getLong(1) + 1L } }这种条件聚合用内置函数很难实现，而UDAF可以灵活处理。
