UDAF Spark中的多列输出

import org.apache.spark.sql.expressions.{UserDefinedAggregateFunction, MutableAggregationBuffer}
import org.apache.spark.sql.types._
import org.apache.spark.sql.Row

class MultiColumnUDAF extends UserDefinedAggregateFunction {
  // 定义输入数据的结构
  def inputSchema: StructType = StructType(StructField("input", DoubleType) :: Nil)

  // 定义聚合缓冲区的结构
  def bufferSchema: StructType = StructType(
    StructField("sum", DoubleType) ::
    StructField("count", LongType) ::
    StructField("avg", DoubleType) :: Nil
  )

  // 定义聚合函数的返回类型
  def dataType: DataType = StructType(
    StructField("sum", DoubleType) ::
    StructField("count", LongType) ::
    StructField("avg", DoubleType) :: Nil
  )

  // 初始化聚合缓冲区
  def initialize(buffer: MutableAggregationBuffer): Unit = {
    buffer(0) = 0.0  // sum
    buffer(1) = 0L   // count
    buffer(2) = 0.0  // avg
  }

  // 根据输入数据更新聚合缓冲区
  def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
    if (!input.isNullAt(0)) {
      val inputValue = input.getDouble(0)
      buffer(0) = buffer.getDouble(0) + inputValue  // sum
      buffer(1) = buffer.getLong(1) + 1L            // count
      buffer(2) = buffer.getDouble(0) / buffer.getLong(1)  // avg
    }
  }

  // 合并两个聚合缓冲区
  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
    buffer1(2) = buffer1.getDouble(0) / buffer1.getLong(1)    // avg
  }

  // 计算最终的聚合结果
  def evaluate(buffer: Row): Any = {
    Row(buffer.getDouble(0), buffer.getLong(1), buffer.getDouble(2))
  }
}

// 注册自定义聚合函数
val multiColumnUDAF = new MultiColumnUDAF
spark.udf.register("multiColumnUDAF", multiColumnUDAF)

// 使用自定义聚合函数进行聚合操作
val result = spark.sql("SELECT multiColumnUDAF(value) FROM table")