Spark中自定义聚合函数实现百分位数

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

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

  // 定义聚合缓冲区的结构
  def bufferSchema: StructType = StructType(StructField("values", ArrayType(DoubleType)) :: Nil)

  // 定义返回结果的数据类型
  def dataType: DataType = DoubleType

  // 初始化聚合缓冲区
  def initialize(buffer: MutableAggregationBuffer): Unit = {
    buffer(0) = Seq.empty[Double]
  }

  // 根据输入数据更新聚合缓冲区
  def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
    val values = buffer.getAs[Seq[Double]](0)
    buffer(0) = values :+ input.getDouble(0)
  }

  // 合并两个聚合缓冲区
  def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
    val values1 = buffer1.getAs[Seq[Double]](0)
    val values2 = buffer2.getAs[Seq[Double]](0)
    buffer1(0) = values1 ++ values2
  }

  // 计算最终结果
  def evaluate(buffer: Row): Any = {
    val values = buffer.getAs[Seq[Double]](0)
    val sortedValues = values.sorted
    val percentile = 0.5 // 50th percentile
    val index = (percentile * sortedValues.size).toInt
    sortedValues(index)
  }
}

// 注册自定义聚合函数
spark.udf.register("percentile", new PercentileUDAF)

// 使用自定义聚合函数
val result = spark.sql("SELECT category, percentile(value) FROM table GROUP BY category")