Spark sql实现自定义函数
Spark sql实现自定义函数
文章目录
- 一、为什么要自定义function？
- 二、实现自定义的函数
- 三、测试效果
-
- 总结
一、为什么要自定义function？
有小伙伴可能会疑问：Spark Sql提供了编写UDF和UDAF的接口扩展，为什么还有开发自定义函数呢？
虽然Spark SQL 提供了UDF和UDAF，但是当我们想要实现 原生函数一样的功能比如:语义参数 ，可变参数等 功能时候，UDF和UDAF就无法满足。
例如 我们想要实现类似于substr这样的函数， udf就无法实现， 其中的参数 ‘Spark SQL’ FROM 5、还有后面两个参数中最后一个可有可无的情况下。
> SELECT substr('Spark SQL', 5);
k SQL
> SELECT substr('Spark SQL', -3);
SQL
> SELECT substr('Spark SQL', 5, 1);
k
> SELECT substr('Spark SQL' FROM 5);
k SQL
> SELECT substr('Spark SQL' FROM -3);
SQL
> SELECT substr('Spark SQL' FROM 5 FOR 1);
k
``
二、实现自定义的函数
spark 官网提供了 SparkSessionExtensions类 ，可以自定义的增强和扩展Spark的很多能力，例如： injectOptimizerRule、injectOptimizerRule等等。
举个例子吧。
为什么会有这样的需求呢？
原因是我想要解决Spark SQl 中的一些函数不完全满足我想要的功能。
比如：原生的spark Sql 函数to_timestamp 在执行有些参数的时候因为数据的格式和指定的parrten不匹配导致运行为null （严格模式下会报错）
我期望的结果应该为：2020-08-08 00:00:00,而不是为null， 简言之就是parrten只要是正确的时间格式，就应该解析出来。
这里是我们的需求，如果各位其他的需求 spark Sql 中的函数不是完全满足，通过UDF能实现，就用UDF实现，或者不完全满足 就跟我这个例子一样进行重写覆盖，如果完全没有 也可以按照这个逻辑自己定义一个全新的函数实现。
解决思路：
老套路，跟踪源码找到 报null和报错的代码逻辑，开发函数，重写逻辑，然后覆盖原函数。
问题代码如下：
1.ToTimestamp的eval方法
case StringType =>
val fmt = right.eval(input)
if (fmt == null) {
null
} else {
val formatter = formatterOption.getOrElse(getFormatter(fmt.toString))
try {
formatter.parse(t.asInstanceOf[UTF8String].toString) / downScaleFactor
} catch {
case e if isParseError(e) =>
if (failOnError) {
throw e
} else {
null
}
}
可以看出解析失败 直接catch,根据failOnError 是否为严格模式报错还是返回null
2.ToTimestamp的doGenCode方法
case StringType => formatterOption.map {
fmt =>
val df = classOf[TimestampFormatter].getName
val formatterName = ctx.addReferenceObj("formatter", fmt, df)
nullSafeCodeGen(ctx, ev, (datetimeStr, _) =>
s"""
|try {
| ${ev.value} = $formatterName.parse($datetimeStr.toString()) / $downScaleFactor;
|} catch (java.time.DateTimeException e) {
| $parseErrorBranch
|} catch (java.time.format.DateTimeParseException e) {
| $parseErrorBranch
|} catch (java.text.ParseException e) {
| $parseErrorBranch
|}
|""".stripMargin)
}
这里是拼接java代码的逻辑，逻辑和eval方法相同。
解决
1.开发逻辑
新建一个样例类继承ToTimestamp，重写上述的逻辑代码
解决思路： 当获取异常后，判断如果是应为格式问题解释失败，识别数据格式，将数据按照数据的格式解析成时间，然后再将时间类型的数据，解析成用户指定的字符串格式。详情看代码。
package v2.jdbc.spark.expressions.function
import java.text.ParseException
import java.time.format.DateTimeParseException
import java.time.{
DateTimeException, ZoneId}
import org.apache.spark.sql.catalyst.expressions.codegen.Block.BlockHelper
import org.apache.spark.sql.catalyst.expressions.codegen.{
CodeGenerator, CodegenContext, ExprCode}
import org.apache.spark.sql.catalyst.expressions.{
Cast, Expression, TimeZoneAwareExpression, ToTimestamp}
import org.apache.spark.sql.catalyst.util.DateTimeUtils.daysToMicros
import org.apache.spark.sql.catalyst.util.{
LegacyDateFormats, TimestampFormatter}
import org.apache.spark.sql.catalyst.{
FunctionIdentifier, InternalRow}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.{
DataType, DateType, StringType, TimestampType}
import org.apache.spark.unsafe.types.UTF8String
import v2.jdbc.spark.expressions.extra.{
ExpressionUtils, FunctionDescription}
import v2.jdbc.spark.expressions.function.DateTimeUtils.dateStrChangeFormat
case class BiGetTimestamp(
left: Expression,
right: Expression,
timeZoneId: Option[String] = None,
failOnError: Boolean = SQLConf.get.ansiEnabled)extends ToTimestamp {
override val downScaleFactor = 1
override def dataType: DataType = TimestampType
override def withTimeZone(timeZoneId: String): TimeZoneAwareExpression =
copy(timeZoneId = Option(timeZoneId))
private def isParseError(e: Throwable): Boolean = e match {
case _: DateTimeParseException |
_: DateTimeException |
_: RuntimeException |
_: ParseException => true
case _ => false
}
override def eval(input: InternalRow): Any = {
val t = left.eval(input)
if (t == null) {
null
} else {
left.dataType match {
case DateType =>
daysToMicros(t.asInstanceOf[Int], zoneId) / downScaleFactor
case TimestampType =>
t.asInstanceOf[Long] / downScaleFactor
case StringType =>
val fmt = right.eval(input)
if (fmt == null) {
null
} else {
val formatter = formatterOption.getOrElse(getFormatter(fmt.toString))
try {
formatter.parse(t.asInstanceOf[UTF8String].toString) / downScaleFactor
} catch {
case e if isParseError(e)=>
val dateStr =UTF8String.fromString(dateStrChangeFormat(t.toString, fmt.toString)).toString
formatter.parse(dateStr) / downScaleFactor
case other=>
if (failOnError) {
throw other
} else {
null
}
}
}
}
}
}
def doGenCodeErrorProcess(str1:String, datetimeStr: String,
ev: ExprCode,formatterName:String,pattern:String):String={
s"""
|
| boolean year = false;
| $str1
| if(pattern.matcher($datetimeStr.toString().substring(0, 4)).matches()) {
| year = true;
| }
| StringBuilder sb = new StringBuilder();
| int index = 0;
| if(!year) {
| if($datetimeStr.toString().contains("月") || $datetimeStr.toString().contains("-") || $datetimeStr.toString().contains("/")) {
| if(Character.isDigit($datetimeStr.toString().charAt(0))) {
| index = 1;
| }
| }else {
| index = 3;
| }
| }
| for (int i = 0; i < $datetimeStr.toString().length(); i++) {
| char chr = $datetimeStr.toString().charAt(i);
| if(Character.isDigit(chr)) {
| if(index==0) {
| sb.append("y");
| }
| if(index==1) {
| sb.append("M");
| }
| if(index==2) {
| sb.append("d");
| }
| if(index==3) {
| sb.append("H");
| }
| if(index==4) {
| sb.append("m");
| }
| if(index==5) {
| sb.append("s");
| }
| if(index==6) {
| sb.append("S");
| }
| }else {
| if(i>0) {
| char lastChar = $datetimeStr.toString().charAt(i-1);
| if(Character.isDigit(lastChar)) {
| index++;
| }
| }
| sb.append(chr);
| }
| }
| java.text.SimpleDateFormat simpleDateFormat = new java.text.SimpleDateFormat(sb.toString());
| java.util.Date date = simpleDateFormat.parse($datetimeStr.toString());
| java.text.SimpleDateFormat simpleDateFormat2 = new java.text.SimpleDateFormat("$pattern");
| ${ev.value} = $formatterName.parse(simpleDateFormat2.format(date)) / $downScaleFactor;
|""".stripMargin
}
override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
val javaType = CodeGenerator.javaType(dataType)
val parseErrorBranch = if (failOnError) "throw e;" else s"${ev.isNull} = true;"
val code = left.dataType match {
case StringType => formatterOption.map {
fmt =>
val df = classOf[TimestampFormatter].getName
val formatterName = ctx.addReferenceObj("formatter", fmt, df)
val patternField = fmt.getClass.getDeclaredField("pattern")
patternField.setAccessible(true)
val pattern = patternField.get(fmt).toString
val str1="java.util.regex.Pattern pattern = java.util.regex.Pattern.compile(\"^[-\\\\+]?[\\\\d]*$\");"
nullSafeCodeGen(ctx, ev, (datetimeStr, _) =>
s"""
|try {
| ${ev.value} = $formatterName.parse($datetimeStr.toString()) / $downScaleFactor;
|} catch (java.time.DateTimeException e) {
| ${doGenCodeErrorProcess(str1,datetimeStr,ev,formatterName,pattern)}
|} catch (java.time.format.DateTimeParseException e) {
| ${doGenCodeErrorProcess(str1,datetimeStr,ev,formatterName,pattern)}
|} catch (java.text.ParseException e) {
| ${doGenCodeErrorProcess(str1,datetimeStr,ev,formatterName,pattern)}
|} catch (java.lang.RuntimeException e) {
| ${doGenCodeErrorProcess(str1, datetimeStr, ev, formatterName, pattern)}
|} catch (java.lang.Exception e) {
| $parseErrorBranch
|}
|""".stripMargin)
}.getOrElse {
val zid = ctx.addReferenceObj("zoneId", zoneId, classOf[ZoneId].getName)
val tf = TimestampFormatter.getClass.getName.stripSuffix("$")
val ldf = LegacyDateFormats.getClass.getName.stripSuffix("$")
val timestampFormatter = ctx.freshName("timestampFormatter")
nullSafeCodeGen(ctx, ev, (string, format) =>
s"""
|$tf $timestampFormatter = $tf$$.MODULE$$.apply(
| $format.toString(),
| $zid,
| $ldf$$.MODULE$$.SIMPLE_DATE_FORMAT(),
| true);
|try {
| ${ev.value} = $timestampFormatter.parse($string.toString()) / $downScaleFactor;
|} catch (java.time.format.DateTimeParseException e) {
| $parseErrorBranch
|} catch (java.time.DateTimeException e) {
| $parseErrorBranch
|} catch (java.text.ParseException e) {
| $parseErrorBranch
|}
|""".stripMargin)
}
case TimestampType =>
val eval1 = left.genCode(ctx)
ev.copy(code =
code"""
${eval1.code}
boolean ${ev.isNull} = ${eval1.isNull};
$javaType ${ev.value} = ${CodeGenerator.defaultValue(dataType)};
if (!${ev.isNull}) {
${ev.value} = ${eval1.value} / $downScaleFactor;
}""")
case DateType =>
val zid = ctx.addReferenceObj("zoneId", zoneId, classOf[ZoneId].getName)
val dtu = DateTimeUtils.getClass.getName.stripSuffix("$")
val eval1 = left.genCode(ctx)
ev.copy(code =
code"""
${eval1.code}
boolean ${ev.isNull} = ${eval1.isNull};
$javaType ${ev.value} = ${CodeGenerator.defaultValue(dataType)};
if (!${ev.isNull}) {
${ev.value} = $dtu.daysToMicros(${eval1.value}, $zid) / $downScaleFactor;
}""")
}
code
}
}
object BiGetTimestamp {
val fd: FunctionDescription = (
new FunctionIdentifier("to_timestamp"),
ExpressionUtils.getExpressionInfo(classOf[BiGetTimestamp], "to_timestamp"),
(children: Seq[Expression]) =>
children.size match {
case 1=>
Cast(children.head, TimestampType,Some("Asia/Shanghai"))
case 2=>
BiGetTimestamp(children.head,children(1),Some("Asia/Shanghai"))
case _=>throw new Exception("参数异常")
}
)
val fd_toDte: FunctionDescription = (
new FunctionIdentifier("to_date"),
ExpressionUtils.getExpressionInfo(classOf[BiGetTimestamp], "to_date"),
(children: Seq[Expression]) =>
children.size match {
case 1=>
Cast(children.head, DateType,Some("Asia/Shanghai"))
case 2=>
Cast(BiGetTimestamp(children.head,children(1),Some("Asia/Shanghai")), DateType,Some("Asia/Shanghai"))
case _=>throw new Exception("参数异常")
}
)
}
package v2.jdbc.spark.expressions.function
import java.text.SimpleDateFormat
import java.util.regex.Pattern
object DateTimeUtils {
/**
* 识别日期字符串的日期格式
*/
def identifyDateType(str: String): String = {
var year = false
val pattern = Pattern.compile("^[-\\+]?[\\d]*$")
if (pattern.matcher(str.substring(0, 4)).matches) year = true
val sb = new StringBuilder
var index = 0
if (!year) if (str.contains("月") || str.contains("-") || str.contains("/")) if (Character.isDigit(str.charAt(0))) index = 1
else index = 3
for (i <- 0 until str.length) {
val chr = str.charAt(i)
if (Character.isDigit(chr)) {
if (index == 0) sb.append("y")
if (index == 1) sb.append("M")
if (index == 2) sb.append("d")
if (index == 3) sb.append("H")
if (index == 4) sb.append("m")
if (index == 5) sb.append("s")
if (index == 6) sb.append("S")
}
else {
if (i > 0) {
val lastChar = str.charAt(i - 1)
if (Character.isDigit(lastChar)) index += 1
}
sb.append(chr)
}
}
sb.toString
}
def dateStrChangeFormat(dateStr: String, targetFormat: String): String = {
val sourceFormat = new SimpleDateFormat(identifyDateType(dateStr))
val date = sourceFormat.parse(dateStr)
val sourceFormat2 = new SimpleDateFormat(targetFormat)
sourceFormat2.format(date)
}
}
2.注册函数
函数注册 有两种方式：
1.直接在构建Spark Session时候通过withExtensions直接使用。
2.不直接使用，通过SparkConf配置。
在配置参数中配置：
spark.sql.extensions=v2.jdbc.spark.expressions.extra.FunctionExtensions
三、测试效果
总结
这里是我们的需求改写并覆盖原有的函数，如果各位其他的需求 spark Sql 中的函数不是完全满足，通过UDF能实现，就用UDF实现，或者不完全满足 就跟我这个例子一样进行重写覆盖，如果完全没有 也可以按照这个逻辑 根据要实现函数的类型进行继承对应的Expression,编写eval和doGenCode方法， 自己定义一个全新的函数。