Spark SQL实现自定义聚合函数

超、凢脫俗 2022-01-26 16:29 377阅读 0赞

## 概述 ##

spark自定义聚合函数需要继承一个抽象类UserDefinedAggregateFunction，并需要重写属性和方法:

*  inputSchema：函数的参数列表，不过需要写成StructType的格式
 *  bufferSchema：中间结果的类型，比如求和时，a、b、c相加，需要先计算a+b并保存结果ab,然后计算ab+c，这个ab就是中间结果。
 *  dataType：返回值结果类型，显示是DataType，也就是org.apache.spark.sql.types包下的那些类
 *  deterministic: 结果是否是确定性的，即相同的输入，是否一定会有相同的输出
 *  initialize：初始化中间结果，例如求和函数,开始计算前需要先将中间结果赋值为0
 *  update(buffer: MutableAggregationBuffer, input: Row)：更新中间结果，input是dataframe的一行，buffer是整个分片遍历过来的中间结果。
 *  merge(buffer1:MutableAggregationBuffer,buffer2:Row)：分片的合并，buffer2一个分片的中间结果，buffer1是整个合并过程的中间结果
 *  evaluate(buffer:Row)：返回函数结果，buffer是7的合并过程的中间结果buffer1遍历所有分片结束后的结果。

实现自定义聚合函数代码

package cn.demo.udf
    
    import org.apache.spark.sql.Row
    import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
    import org.apache.spark.sql.types.{DataType, StringType, StructField, StructType}
    
    /**
      * 自定义聚合函数
      */
    class GroupConcatDistinctUDAF extends UserDefinedAggregateFunction {
    
      override def inputSchema: StructType = StructType(StructField("cityInfo", StringType) :: Nil)
    
      override def bufferSchema: StructType = StructType(StructField("bufferCityInfo", StringType) :: Nil)
    
    // dataType，指的是，函数返回值的类型
      override def dataType: DataType = StringType
    
      override def deterministic: Boolean = true
      
      // 为每个分组的数据执行初始化操作
      override def initialize(buffer: MutableAggregationBuffer): Unit = {
        buffer(0)= ""
      }
    
      /**
        * 更新
        * 可以认为是，一个一个地将组内的字段值传递进来
        * 实现拼接的逻辑
        */
      override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
        // 缓冲中的已经拼接过的城市信息串
        var bufferCityInfo = buffer.getString(0)
        // 刚刚传递进来的某个城市信息
        val cityInfo = input.getString(0)
    
        // 在这里要实现去重的逻辑
        // 判断：之前没有拼接过某个城市信息，那么这里才可以接下去拼接新的城市信息
        if(!bufferCityInfo.contains(cityInfo)) {
          if("".equals(bufferCityInfo))
            bufferCityInfo += cityInfo
          else {
            // 比如1:北京
            // 1:北京,2:上海
            bufferCityInfo += "," + cityInfo
          }
    
          buffer.update(0, bufferCityInfo)
        }
      }
       
      // 由于Spark是分布式的，所以一个分组的数据，可能会在不同的节点上进行局部聚合，就是update
      // 但是，最后一个分组，在各个节点上的聚合值，要进行merge，也就是合并
      override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
        var bufferCityInfo1 = buffer1.getString(0);
        val bufferCityInfo2 = buffer2.getString(0);
    
        for(cityInfo <- bufferCityInfo2.split(",")) {
          if(!bufferCityInfo1.contains(cityInfo)) {
            if("".equals(bufferCityInfo1)) {
              bufferCityInfo1 += cityInfo;
            } else {
              bufferCityInfo1 += "," + cityInfo;
            }
          }
        }
    
        buffer1.update(0, bufferCityInfo1);
      }
      
       // 最后，指的是，一个分组的聚合值，如何通过中间的缓存聚合值，最后返回一个最终的聚合值
      override def evaluate(buffer: Row): Any = {
        buffer.getString(0)
      }
    
    }

实现注册自定义聚合函数

package cn.demo.udf
    
    import org.apache.spark.SparkConf
    import org.apache.spark.sql.SparkSession
    
    object GroupConcatMain {
      def main(args: Array[String]): Unit = {
        // 构建Spark上下文
        val sparkConf = new SparkConf().setAppName("SessionAnalyzer").setMaster("local[*]")
    
        // 创建Spark客户端
        val spark = SparkSession.builder().config(sparkConf).enableHiveSupport().getOrCreate()
        val sc = spark.sparkContext
    
        // 注册自定义函数
        spark.udf.register("concat_long_string", new GroupConcatDistinctUDAF)
      }
    }

这样就可以在spark中使用自定义聚合函数了