在多个列上应用自定义Spark聚合器(Spark 2.0)

import org.apache.spark.sql.expressions.Aggregator
import org.apache.spark.sql.{Encoder, Encoders, SparkSession}
import org.apache.spark.sql.functions._

// 创建自定义聚合器类
class MyAggregator extends Aggregator[(Int, Int), (Int, Int), Double] {
  // 初始化中间状态
  def zero: (Int, Int) = (0, 0)

  // 在每个分区上进行聚合操作
  def reduce(buffer: (Int, Int), data: (Int, Int)): (Int, Int) = {
    (buffer._1 + data._1, buffer._2 + data._2)
  }

  // 合并各个分区的聚合结果
  def merge(buffer1: (Int, Int), buffer2: (Int, Int)): (Int, Int) = {
    (buffer1._1 + buffer2._1, buffer1._2 + buffer2._2)
  }

  // 定义最终的聚合操作
  def finish(buffer: (Int, Int)): Double = {
    buffer._1.toDouble / buffer._2.toDouble
  }

  // 定义编码器
  def bufferEncoder: Encoder[(Int, Int)] = Encoders.product[(Int, Int)]
  def outputEncoder: Encoder[Double] = Encoders.scalaDouble
}

// 创建SparkSession
val spark = SparkSession.builder()
  .appName("Spark Aggregator Example")
  .getOrCreate()

import spark.implicits._

// 创建测试数据
val data = Seq((1, 2), (3, 4), (5, 6))
val df = data.toDF("col1", "col2")

// 创建聚合器实例
val myAggregator = new MyAggregator

// 注册聚合器为用户定义的聚合函数
val myAggregatorUDF = udf(myAggregator.toColumn)

// 在多个列上应用自定义聚合器
val result = df.select(myAggregatorUDF($"col1", $"col2").as("avg"))

result.show()