viirya commented on a change in pull request #24735: [SPARK-27871][SQL] LambdaVariable should use per-query unique IDs instead of globally unique IDs URL: https://github.com/apache/spark/pull/24735#discussion_r288875034
########## File path: sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects/objects.scala ########## @@ -575,15 +575,43 @@ case class WrapOption(child: Expression, optType: DataType) } } +object LambdaVariable { + private val curId = new java.util.concurrent.atomic.AtomicLong() + + // Returns the codegen-ed `LambdaVariable` and add it to mutable states, so that it can be + // accessed anywhere in the generated code. + def prepareLambdaVariable(ctx: CodegenContext, variable: LambdaVariable): ExprCode = { + val variableCode = variable.genCode(ctx) + assert(variableCode.code.isEmpty) + + ctx.addMutableState( + CodeGenerator.javaType(variable.dataType), + variableCode.value, + forceInline = true, + useFreshName = false) + + if (variable.nullable) { + ctx.addMutableState( + CodeGenerator.JAVA_BOOLEAN, + variableCode.isNull, + forceInline = true, + useFreshName = false) + } + + variableCode + } +} + /** - * A placeholder for the loop variable used in [[MapObjects]]. This should never be constructed + * A placeholder for the loop variable used in [[MapObjects]]. This should never be constructed * manually, but will instead be passed into the provided lambda function. */ +// TODO: Merge this and `NamedLambdaVariable`. case class LambdaVariable( - value: String, - isNull: String, + name: String, dataType: DataType, - nullable: Boolean = true) extends LeafExpression with NonSQLExpression { + nullable: Boolean, + id: Long = LambdaVariable.curId.incrementAndGet) extends LeafExpression with NonSQLExpression { Review comment: Probably add comment about the `id`? The allocated `id`s here can't be directly used. It needs going through `ReassignLambdaVariableID` to reassign. If the expressions including `LambdaVariable` skip the normal query processing steps, we still make sure `ReassignLambdaVariableID` is applied. ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org With regards, Apache Git Services --------------------------------------------------------------------- To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org