SparkSQL自定义udf

自定义udf求平均值：

测试数据：

{"name":"Michael", "salary":3000}
{"name":"Andy", "salary":4500}
{"name":"Justin", "salary":3500}
{"name":"Berta", "salary":4000}

第一种方法：

def main(args: Array[String]): Unit = {
val spark = SparkSession.builder().appName("myUdf").master("local[*]").getOrCreate()
        spark.udf.register("myAverage", MyAverage)
        val df: DataFrame = spark.read.json("employees.json")
        df.createOrReplaceTempView("employees")
        val result = spark.sql("select myAverage(salary) as avg from employees")
        result.show()
}
object MyAverage extends UserDefinedAggregateFunction {
  override def inputSchema: StructType = StructType(StructField("salary", LongType) :: Nil)

  override def bufferSchema: StructType = StructType(StructField("sum", LongType) :: StructField("count", LongType) :: Nil)
//返回值类型
  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 = {
    if (!input.isNullAt(0)) {
      buffer(0) = buffer.getLong(0) + input.getLong(0)//将每一条salary进行累加
      buffer(1) = buffer.getLong(1) + 1//数量的累加
    }
  }
//合并
  override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
    buffer1(0) = buffer1.getLong(0) + buffer2.getLong(0)//sum合并
    buffer1(1) = buffer1.getLong(1) + buffer2.getLong(1)//count合并
  }
//最后的处理
  override def evaluate(buffer: Row): Any = {
    buffer.getLong(0).toDouble / buffer.getLong(1)//求平均值
  }
}

第二种方法：

  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder().appName("myUdf").master("local[*]").getOrCreate()
        spark.udf.register("myAverage", MyAverage)
        val df: DataFrame = spark.read.json("employees.json")    
import spark.implicits._
        val ds = df.as[Employee]
        ds.show()
        val avg: TypedColumn[Employee, Double] = MyAverage1.toColumn.name("avg")
        ds.select(avg).show()
}
case class Employee(name: String, salary: Long)

case class Average(var sum: Long, var count: Long)

object MyAverage1 extends Aggregator[Employee, Average, Double] {
  override def zero: Average = Average(0L, 0L)

  override def reduce(b: Average, a: Employee): Average = {
    b.sum += a.salary
    b.count += 1
    b
  }

  override def merge(b1: Average, b2: Average): Average = {
    b1.sum += b2.sum
    b1.count += b2.count
    b1
  }

  override def finish(reduction: Average): Double = {
    reduction.sum.toDouble / reduction.count
  }

  override def bufferEncoder: Encoder[Average] = Encoders.product

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

代码来自官网