【Spark精讲】一文讲透SparkSQL聚合过程以及UDAF开发

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.expressions.MutableAggregationBuffer
import org.apache.spark.sql.expressions.UserDefinedAggregateFunction
import org.apache.spark.sql.types.{DataType, StructType}
import org.apache.spark.sql.Row
import org.apache.spark.sql.types.DataTypes
import org.apache.spark.sql.types.StructField
 
// 定义一个简单的UDAF，用于计算一列数字的平均值
class Average extends UserDefinedAggregateFunction {
  // 输入数据的数据结构定义
  override def inputSchema: StructType = StructType(StructField("input", DataTypes.DoubleType) :: Nil)
 
  // 缓冲区的数据结构定义，用于累计中间结果
  override def bufferSchema: StructType = StructType(StructField("sum", DataTypes.DoubleType) :: StructField("count", DataTypes.LongType) :: Nil)
 
  // 返回结果的数据类型
  override def dataType: DataType = DataTypes.DoubleType
 
  // 是否是确定性的函数
  override def deterministic: Boolean = true
 
  // 初始化缓冲区
  override def initialize(buffer: MutableAggregationBuffer): Unit = {
    buffer(0) = 0.0 // 初始化sum为0.0
    buffer(1) = 0L  // 初始化count为0
  }
 
  // 更新缓冲区
  override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
    buffer(0) = buffer.getDouble(0) + input.getDouble(0) // 累加数字
    buffer(1) = buffer.getLong(1) + 1L // 累加计数
  }
 
  // 合并缓冲区
  override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
    buffer1(0) = buffer1.getDouble(0) + buffer2.getDouble(0)
    buffer1(1) = buffer1.getLong(1) + buffer2.getLong(1)
  }
 
  // 计算最终结果
  override def evaluate(buffer: Row): Any = {
    buffer.getDouble(0) / buffer.getLong(1)
  }
}
 
// 使用示例
object AverageExample {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder().appName("AverageExample").getOrCreate()
    import spark.implicits._
 
    val data = Seq(1, 2, 3, 4, 5).map(Row(_))
    val df = spark.createDataFrame(data, StructType(StructField("input", DataTypes.DoubleType) :: Nil))
 
    // 注册自定义的UDAF
    spark.udf.register("average", new Average)
 
    // 使用UDAF
    df.selectExpr("average(input) as average").show()
 
    spark.stop()
  }
}

这段代码定义了一个简单的UDAF，用于计算输入数字的平均值。它展示了如何使用Spark SQL的UserDefinedAggregateFunction接口来创建自定义的聚合函数。代码中包含了初始化、更新缓冲区、合并缓冲区以及计算最终结果的方法。最后，提供了一个使用该UDAF的示例，展示了如何注册该UDAF并在DataFrame上使用它。