Spark-用户自定义聚合函数

╰+哭是因爲堅強的太久メ 2023-02-19 09:27 43阅读 0赞

\*\*`强类型`\*\*的Dataset和弱类型的DataFrame都提供了相关的聚合函数， 如 count()，countDistinct()，avg()，max()，min()。除此之外，用户可以设定自己的自定义聚合函数。  
\*\*`弱类型`\*\*用户自定义聚合函数：通过继承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 SparkSQL_UDAF { 
    
      def main(args: Array[String]) { 
    
       //创建配置对象
       val sparkConf = new SparkConf().setMaster("local[*]").setAppName("sql")
    
       //创建SparkSQL的环境
       val spark = SparkSession.builder().config(sparkConf).getOrCreate()
    
       //进行转换之前，需要引入隐士转换规则
       //这里spark不是包名的含义，是sparksession对象的名字
       import spark.implicits._
    
       //自定义集合函数
       //创建聚合函数对象
       val udaf = new MyAgeAvgFunction
       //注册聚合函数
       spark.udf.register("avgAge", udaf)
    
       //使用聚合函数
       val frame = spark.read.json("in/user.json")
    
       frame.createOrReplaceTempView("user")
    
       spark.sql("select avgAge(age) from user").show
    
       //释放资源
       spark.stop()
    
    
    
      }
    
    }
    //声明用户自定义聚合函数(弱类型)
    //1.继承UserDefinedAggregateFunction
    //2.实现方法
    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
     //计算前的缓冲区的初始化 0代表sum 1代表count
     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)
     }
    }

**`强类型`**

import org.apache.spark.SparkConf
    import org.apache.spark.sql.{ Encoder, Encoders, SparkSession}
    import org.apache.spark.sql.expressions.Aggregator
    
    
    object SparkSQL_UDAF1 { 
    
      def main(args: Array[String]) { 
    
        //创建配置对象
        val sparkConf = new SparkConf().setMaster("local[*]").setAppName("sql")
    
        //创建SparkSQL的环境
        val spark = SparkSession.builder().config(sparkConf).getOrCreate()
    
        //进行转换之前，需要引入隐士转换规则
        //这里spark不是包名的含义，是sparksession对象的名字
        import spark.implicits._
    
        //创建聚合函数对象
        val udaf = new MyAgeAvgFunction1
    
        //将聚合函数转换为查询列
        val avgCol = udaf.toColumn.name("avgAge")
    
        val frame = spark.read.json("in/user.json")
    
        val userDS = frame.as[UserBean]
    
        userDS.select(avgCol).show()
    
       //自定义集合函数
       //创建聚合函数对象
    
       //释放资源
       spark.stop()
    
      }
    
    }
    
    case class UserBean(name:String,age:BigInt)
    case class AvgBuffer(var sum:BigInt,var count:Int)
    //声明用户自定义聚合函数(强类型)
    //1.继承Aggregator
    //2.实现方法
    class MyAgeAvgFunction1 extends Aggregator[UserBean,AvgBuffer,Double] { 
    
      //初始化
      override def zero: AvgBuffer = { 
        AvgBuffer(0,0)
      }
      //聚合数据
      override def reduce(b: AvgBuffer, a: UserBean): AvgBuffer = { 
        b.sum=b.sum+a.age
        b.count=b.count+1
        b
      }
      //缓冲区的合并操作
      override def merge(b1: AvgBuffer, b2: AvgBuffer): AvgBuffer = { 
        b1.sum=b1.sum +b2.sum
        b1.count=b1.count+b2.count
        b1
      }
      //完成计算
      override def finish(reduction: AvgBuffer): Double = { 
        reduction.sum.toDouble / reduction.count
      }
    
      override def bufferEncoder: Encoder[AvgBuffer] = Encoders.product
    
      override def outputEncoder: Encoder[Double] = Encoders.scalaDouble
    }

user.json中的数据

{ "name":"zhangsan","age":20}
    { "name":"lisi","age":30}
    { "name":"wabgwu","age":40}

**`虽然会报错，但是spark是支持的，不用管`**  
![在这里插入图片描述][20200619221526311.png]

[20200619221526311.png]: https://img-blog.csdnimg.cn/20200619221526311.png