弱类型用户自定义聚合函数:通过继承UserDefinedAggregateFunction来实现用户自定义聚合函数。
import org.apache.spark.SparkConf
import org.apache.spark.sql.{Row, SparkSession}
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, DoubleType, LongType, StructType}
object SparkSQL06_UDAF {
def main(args: Array[String]): Unit = {
// 创建SparkConf
val conf = new SparkConf().setMaster("local[*]").setAppName("Test")
//创建SparkSession
val spark = SparkSession.builder().config(conf).getOrCreate()
//创建自定义函数对象
val udaf = new MyAgeAvgFunction
//注册udaf
spark.udf.register("avgAge",udaf)
//读取数据
val frame = spark.read.json("input/user.json")
frame.createOrReplaceTempView("users")
spark.sql("select avgAge(age) from users").show
spark.stop()
}
}
/**
* 用户自定义聚合函数
*/
class MyAgeAvgFunction extends UserDefinedAggregateFunction{
//输入的数据结构
override def inputSchema: StructType = {
new StructType().add("age",LongType)
}
//计算时的数据结构
override def bufferSchema: StructType = {
new StructType().add("sum",LongType).add("count",LongType)
}
//函数返回的数据类型
override def dataType: DataType = DoubleType
//函数是否稳定
override def deterministic: Boolean = true
//计算前缓冲区的初始化
override def initialize(buffer: MutableAggregationBuffer): Unit = {
buffer(0) = 0L
buffer(1) = 0L
}
//根据查询结果更新缓冲区数据
override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
buffer(0) = buffer.getLong(0) + input.getLong(0)
buffer(1) = buffer.getLong(1) + 1
}
//将多个节点的缓冲区合并
override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
//sum
buffer1(0) = buffer1.getLong(0) + buffer2.getLong(0)
//count
buffer1(1) = buffer1.getLong(1) + buffer2.getLong(1)
}
//计算
override def evaluate(buffer: Row): Any = {
buffer.getLong(0).toDouble / buffer.getLong(1)
}
}