[SPARK-18467][SQL] Extracts method for preparing arguments from StaticInvoke, Invoke and NewInstance and modify to short circuit if arguments have null when needNullCheck == true.
#15901
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -32,6 +32,78 @@ import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCo | |
| import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, GenericArrayData} | ||
| import org.apache.spark.sql.types._ | ||
|
|
||
| /** | ||
| * Common base class for [[StaticInvoke]], [[Invoke]], and [[NewInstance]]. | ||
| */ | ||
| trait InvokeLike extends Expression with NonSQLExpression { | ||
|
|
||
| def arguments: Seq[Expression] | ||
|
|
||
| def propagateNull: Boolean | ||
|
|
||
| protected lazy val needNullCheck: Boolean = propagateNull && arguments.exists(_.nullable) | ||
|
|
||
| /** | ||
| * Prepares codes for arguments. | ||
| * | ||
| * - generate codes for argument. | ||
| * - use ctx.splitExpressions() to not exceed 64kb JVM limit while preparing arguments. | ||
| * - avoid some of nullabilty checking which are not needed because the expression is not | ||
| * nullable. | ||
| * - when needNullCheck == true, short circuit if we found one of arguments is null because | ||
| * preparing rest of arguments can be skipped in the case. | ||
| * | ||
| * @param ctx a [[CodegenContext]] | ||
| * @return (code to prepare arguments, argument string, result of argument null check) | ||
| */ | ||
| def prepareArguments(ctx: CodegenContext): (String, String, String) = { | ||
|
|
||
| val resultIsNull = if (needNullCheck) { | ||
| val resultIsNull = ctx.freshName("resultIsNull") | ||
| ctx.addMutableState("boolean", resultIsNull, "") | ||
| resultIsNull | ||
| } else { | ||
| "false" | ||
| } | ||
| val argValues = arguments.map { e => | ||
| val argValue = ctx.freshName("argValue") | ||
| ctx.addMutableState(ctx.javaType(e.dataType), argValue, "") | ||
| argValue | ||
| } | ||
|
|
||
| val argCodes = if (needNullCheck) { | ||
| val reset = s"$resultIsNull = false;" | ||
| val argCodes = arguments.zipWithIndex.map { case (e, i) => | ||
| val expr = e.genCode(ctx) | ||
| val updateResultIsNull = if (e.nullable) { | ||
| s"$resultIsNull = ${expr.isNull};" | ||
| } else { | ||
| "" | ||
| } | ||
| s""" | ||
| if (!$resultIsNull) { | ||
| ${expr.code} | ||
| $updateResultIsNull | ||
| ${argValues(i)} = ${expr.value}; | ||
| } | ||
| """ | ||
| } | ||
| reset +: argCodes | ||
| } else { | ||
| arguments.zipWithIndex.map { case (e, i) => | ||
| val expr = e.genCode(ctx) | ||
| s""" | ||
| ${expr.code} | ||
| ${argValues(i)} = ${expr.value}; | ||
| """ | ||
| } | ||
| } | ||
| val argCode = ctx.splitExpressions(ctx.INPUT_ROW, argCodes) | ||
|
|
||
| (argCode, argValues.mkString(", "), resultIsNull) | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * Invokes a static function, returning the result. By default, any of the arguments being null | ||
| * will result in returning null instead of calling the function. | ||
|
|
@@ -50,7 +122,7 @@ case class StaticInvoke( | |
| dataType: DataType, | ||
| functionName: String, | ||
| arguments: Seq[Expression] = Nil, | ||
| propagateNull: Boolean = true) extends InvokeLike { | ||
|
|
||
| val objectName = staticObject.getName.stripSuffix("$") | ||
|
|
||
|
|
@@ -62,16 +134,10 @@ case class StaticInvoke( | |
|
|
||
| override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { | ||
| val javaType = ctx.javaType(dataType) | ||
|
|
||
| val (argCode, argString, resultIsNull) = prepareArguments(ctx) | ||
|
|
||
| val callFunc = s"$objectName.$functionName($argString)" | ||
|
|
||
| // If the function can return null, we do an extra check to make sure our null bit is still set | ||
| // correctly. | ||
|
|
@@ -82,9 +148,9 @@ case class StaticInvoke( | |
| } | ||
|
|
||
| val code = s""" | ||
| $argCode | ||
| boolean ${ev.isNull} = $resultIsNull; | ||
| final $javaType ${ev.value} = $resultIsNull ? ${ctx.defaultValue(dataType)} : $callFunc; | ||
| $postNullCheck | ||
| """ | ||
| ev.copy(code = code) | ||
|
|
@@ -103,13 +169,15 @@ case class StaticInvoke( | |
| * @param functionName The name of the method to call. | ||
| * @param dataType The expected return type of the function. | ||
| * @param arguments An optional list of expressions, whos evaluation will be passed to the function. | ||
| * @param propagateNull When true, and any of the arguments is null, null will be returned instead | ||
| * of calling the function. | ||
| */ | ||
| case class Invoke( | ||
| targetObject: Expression, | ||
| functionName: String, | ||
| dataType: DataType, | ||
| arguments: Seq[Expression] = Nil, | ||
| propagateNull: Boolean = true) extends InvokeLike { | ||
|
|
||
| override def nullable: Boolean = true | ||
| override def children: Seq[Expression] = targetObject +: arguments | ||
|
|
@@ -131,8 +199,8 @@ case class Invoke( | |
| override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { | ||
| val javaType = ctx.javaType(dataType) | ||
| val obj = targetObject.genCode(ctx) | ||
|
|
||
| val (argCode, argString, resultIsNull) = prepareArguments(ctx) | ||
|
|
||
| val returnPrimitive = method.isDefined && method.get.getReturnType.isPrimitive | ||
| val needTryCatch = method.isDefined && method.get.getExceptionTypes.nonEmpty | ||
|
|
@@ -164,28 +232,26 @@ case class Invoke( | |
| """ | ||
| } | ||
|
|
||
| // If the function can return null, we do an extra check to make sure our null bit is still set | ||
| // correctly. | ||
| val postNullCheck = if (ctx.defaultValue(dataType) == "null") { | ||
| s"${ev.isNull} = ${ev.value} == null;" | ||
| } else { | ||
| "" | ||
| } | ||
|
|
||
| val code = s""" | ||
| ${obj.code} | ||
| boolean ${ev.isNull} = true; | ||
| $javaType ${ev.value} = ${ctx.defaultValue(dataType)}; | ||
| if (!${obj.isNull}) { | ||
|
Contributor
There was a problem hiding this comment. Anyway, I think this new code looks clearer than before, what do you think? @ueshin
Member
Author
There was a problem hiding this comment. Yes, I agree with you. Thank you for your suggestion. |
||
| $argCode | ||
| ${ev.isNull} = $resultIsNull; | ||
| if (!${ev.isNull}) { | ||
| $evaluate | ||
| } | ||
| $postNullCheck | ||
| } | ||
| """ | ||
| ev.copy(code = code) | ||
| } | ||
|
|
@@ -223,10 +289,10 @@ case class NewInstance( | |
| arguments: Seq[Expression], | ||
| propagateNull: Boolean, | ||
| dataType: DataType, | ||
| outerPointer: Option[() => AnyRef]) extends InvokeLike { | ||
| private val className = cls.getName | ||
|
|
||
| override def nullable: Boolean = needNullCheck | ||
|
|
||
| override def children: Seq[Expression] = arguments | ||
|
|
||
|
|
@@ -245,52 +311,25 @@ case class NewInstance( | |
|
|
||
| override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { | ||
| val javaType = ctx.javaType(dataType) | ||
|
|
||
| val (argCode, argString, resultIsNull) = prepareArguments(ctx) | ||
|
|
||
| val outer = outerPointer.map(func => Literal.fromObject(func()).genCode(ctx)) | ||
|
|
||
| ev.isNull = resultIsNull | ||
|
|
||
| val constructorCall = outer.map { gen => | ||
| s"${gen.value}.new ${cls.getSimpleName}($argString)" | ||
| }.getOrElse { | ||
| s"new $className($argString)" | ||
| } | ||
|
|
||
| val code = s""" | ||
| $argCode | ||
| ${outer.map(_.code).getOrElse("")} | ||
| final $javaType ${ev.value} = ${ev.isNull} ? ${ctx.defaultValue(javaType)} : $constructorCall; | ||
| """ | ||
| ev.copy(code = code) | ||
| } | ||
|
|
||
| override def toString: String = s"newInstance($cls)" | ||
|
|
||
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
how about:
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks, I'll use it.