SparkSQL之自定义函数

目录

一 UDF

二 UDAF（弱类型）

三 Aggregator（强类型）

四 Spark早期强类型UDAF操作

---

用户可以通过 spark.udf 功能添加自定义函数，实现自定义功能。

一 UDF

val df: DataFrame = spark.read.json("data/user.json")
df.createOrReplaceTempView("user")

// SparkSQL自定义函数
spark.udf.register("prefixName", (name: String) => {
	"Name:" + name
})

spark.udf.register("tailSui", (age: Int) => {
	age + "岁"
})

spark.sql(
	"""
		|select
		|tailSui(age),
		|prefixName(username)
		|from
		|user
	""".stripMargin).show()

+------------+--------------------+
|tailSui(age)|prefixName(username)|
+------------+--------------------+
|        20岁|           Name:小王|
|        30岁|           Name:小李|
|        40岁|           Name:小黑|
+------------+--------------------+

二 UDAF（弱类型）

        强类型的Dataset 和弱类型的 DataFrame 都提供了相关的聚合函数， 如 count(),countDistinct()，avg()，max()，min()。除此之外，用户可以设定自己的自定义聚合函数。通过继承 UserDefinedAggregateFunction 来实现用户自定义弱类型聚合函数。从Spark3.0 版本后，UserDefinedAggregateFunction 已经不推荐使用了。可以统一采用强类型聚合函数Aggregator

 在Spark3.0之后不推荐使用

def main(args: Array[String]): Unit = {

val sparkConf: SparkConf = new SparkConf().setMaster("local[*]").setAppName("udf")
val spark: SparkSession = SparkSession.builder().config(sparkConf).getOrCreate()
import spark.implicits._

val df: DataFrame = spark.read.json("data/user.json")
df.createTempView("user")

spark.udf.register("MyAvg", new MyAggregateBuffer)

spark.sql("select MyAvg(age) from user").show()


spark.stop()
}

class MyAggregateBuffer extends UserDefinedAggregateFunction {
// 输入的数据结构
override def inputSchema: StructType = StructType(
	Array(
		StructField("age", LongType)
	)
)

// 缓冲区,临时计算
override def bufferSchema: StructType = StructType(
	Array(
		StructField("total", LongType),
		StructField("count", LongType)
	)
)

// 函数计算结果输出类型
override def dataType: DataType = LongType

// 函数稳定性
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.update(0, buffer.getLong(0) + input.getLong(0)) // 年龄总和在第一个位置
	buffer.update(1, buffer.getLong(1) + 1) // 年龄数量
}

// 分区合并
override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
	buffer1.update(0, buffer1.getLong(0) + buffer2.getLong(0))
	buffer1.update(1, buffer1.getLong(1) + buffer2.getLong(1))
}

// 计算平均值
override def evaluate(buffer: Row): Any = {
	buffer.getLong(0) / buffer.getLong(1)
}
}

+----------------------+
|myaggregatebuffer(age)|
+----------------------+
|                    30|
+----------------------+

 

三 Aggregator（强类型）

在Spark3.0之后可以使用

Aggregator是一种类型安全的聚合器，可以用于对数据进行聚合操作。它将输入数据类型和输出数据类型分别定义为泛型，并提供了两个方法，可以对输入数据进行聚合操作并返回最终的输出数据。下面是使用Aggregator的基本步骤：

def main(args: Array[String]): Unit = {
	
	val sparkConf: SparkConf = new SparkConf().setMaster("local[*]").setAppName("udf")
	val spark: SparkSession = SparkSession.builder().config(sparkConf).getOrCreate()
	
	val df: DataFrame = spark.read.json("data/user.json")
	df.createTempView("user")
	
	spark.udf.register("MyAvg", functions.udaf(new MyAvgUDAF))
	
	spark.sql("select MyAvg(age) from user").show()
	
	
	spark.stop()
}

// org.apache.spark.sql.expressions.Aggregator
// 输入类型：Long
// 缓冲：Buf
// 输出数据类型：Double
case class Buf(var total: Long, var count: Long)

class MyAvgUDAF extends Aggregator[Long, Buf, Double] {
	
	// 初始值，缓冲区初始化
	override def zero: Buf = {
		Buf(0L, 0L)
	}
	
	// 根据输入数据更新缓冲区
	override def reduce(b: Buf, in: Long): Buf = {
		b.total += in
		b.count += 1
		b
	}
	
	override def merge(b1: Buf, b2: Buf): Buf = {
		b1.total += b2.total
		b1.count += b2.count
		b1
	}
	
	// 计算结果
	override def finish(reduction: Buf): Double = {
		reduction.total.toDouble / reduction.count
	}
	
	// 分布式计算，需要序列化,自定义的类叫Encoders.product
	override def bufferEncoder: Encoder[Buf] = Encoders.product
	
	// 分布式计算，需要序列化,scala的类叫Encoders.scalaxxx
	override def outputEncoder: Encoder[Double] = Encoders.scalaDouble
}

+--------------+
|myavgudaf(age)|
+--------------+
|          30.0|
+--------------+

四 Spark早期强类型UDAF操作

def main(args: Array[String]): Unit = {
	
	val sparkConf: SparkConf = new SparkConf().setMaster("local[*]").setAppName("udf")
	val spark: SparkSession = SparkSession.builder().config(sparkConf).getOrCreate()
	import spark.implicits._
	val df: DataFrame = spark.read.json("data/user.json")
	
	// 早期UDAF强类型聚合函数使用DSL语法
	val ds: Dataset[User] = df.as[User]
	
	// 将UDAF函数转换为查询的列对象
	val udafCol: TypedColumn[User, Double] = new MyAvgUDAF().toColumn
	
	ds.select(udafCol).show()
	spark.stop()
}

case class User(var username: String, var age: Long)

// org.apache.spark.sql.expressions.Aggregator
// 输入类型：Long
// 缓冲：Buf
// 输出数据类型：Double
case class Buf(var total: Long, var count: Long)

class MyAvgUDAF extends Aggregator[User, Buf, Double] {
	
	// 初始值，缓冲区初始化
	override def zero: Buf = {
		Buf(0L, 0L)
	}
	
	// 根据输入数据更新缓冲区
	override def reduce(b: Buf, in: User): Buf = {
		b.total += in.age
		b.count += 1
		b
	}
	
	override def merge(b1: Buf, b2: Buf): Buf = {
		b1.total += b2.total
		b1.count += b2.count
		b1
	}
	
	// 计算结果
	override def finish(reduction: Buf): Double = {
		reduction.total.toDouble / reduction.count
	}
	
	// 分布式计算，需要序列化,自定义的类叫Encoders.product
	override def bufferEncoder: Encoder[Buf] = Encoders.product
	
	// 分布式计算，需要序列化,scala的类叫Encoders.scalaxxx
	override def outputEncoder: Encoder[Double] = Encoders.scalaDouble
}

+-----------------------------------------------------------------+
|MyAvgUDAF(com.mingyu.spark.sql.Spark02_SparkSQL_aggregator2$User)|
+-----------------------------------------------------------------+
|                                                             30.0|
+-----------------------------------------------------------------+