spark自定义聚合器>=2.0 (scala)

import org.apache.spark.sql.expressions.{Aggregator, MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.{Encoder, Encoders, Row, SparkSession}
import org.apache.spark.sql.types.{DataType, StructType}

// 自定义聚合器
class CustomAggregator extends Aggregator[Row, Int, Double] {
  // 初始化聚合缓冲区
  def zero: Int = 0

  // 更新聚合缓冲区
  def reduce(buffer: Int, input: Row): Int = {
    buffer + input.getInt(0)
  }

  // 合并聚合缓冲区
  def merge(buffer1: Int, buffer2: Int): Int = {
    buffer1 + buffer2
  }

  // 完成聚合操作
  def finish(reduction: Int): Double = {
    reduction.toDouble
  }

  // 定义编码器
  def bufferEncoder: Encoder[Int] = Encoders.scalaInt

  def outputEncoder: Encoder[Double] = Encoders.scalaDouble
}

// 创建SparkSession
val spark = SparkSession.builder()
  .appName("CustomAggregatorExample")
  .getOrCreate()

// 创建测试数据
val data = Seq(Row(1), Row(2), Row(3), Row(4), Row(5))

// 定义数据结构
val schema = new StructType().add("value", "integer")

// 将数据转换为DataFrame
val df = spark.createDataFrame(spark.sparkContext.parallelize(data), schema)

// 注册自定义聚合器
val customAggregator = new CustomAggregator
spark.udf.register("customAggregator", customAggregator)

// 使用自定义聚合器进行聚合操作
val result = df.selectExpr("customAggregator(value)").collect()

// 打印结果
result.foreach(println)