Skip to content

[SPARK-37466][SQL] Support subexpression elimination in higher order functions #51272

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 8 commits into
base: master
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -146,9 +146,13 @@ class EquivalentExpressions(
// There are some special expressions that we should not recurse into all of its children.
// 1. CodegenFallback: it's children will not be used to generate code (call eval() instead)
// 2. ConditionalExpression: use its children that will always be evaluated.
// 3. HigherOrderFunction: lambda functions operate in the context of local lambdas and can't
// be called outside of that scope, only the arguments can be evaluated ahead of
// time.
private def childrenToRecurse(expr: Expression): Seq[Expression] = expr match {
case _: CodegenFallback => Nil
case c: ConditionalExpression => c.alwaysEvaluatedInputs.map(skipForShortcut)
case h: HigherOrderFunction => h.arguments
case other => skipForShortcut(other).children
}

Original file line number Diff line number Diff line change
@@ -174,6 +174,41 @@ class CodegenContext extends Logging {
*/
var currentVars: Seq[ExprCode] = null

/**
* Holding a map of current lambda variables.
*/
var currentLambdaVars: mutable.Map[Long, ExprCode] = mutable.HashMap.empty

def withLambdaVars(
namedLambdas: Seq[NamedLambdaVariable],
f: Seq[ExprCode] => ExprCode): ExprCode = {
val lambdaVars = namedLambdas.map { lambda =>
val id = lambda.exprId.id
if (currentLambdaVars.get(id).nonEmpty) {
throw QueryExecutionErrors.lambdaVariableAlreadyDefinedError(id)
}
val isNull = if (lambda.nullable) {
JavaCode.isNullGlobal(addMutableState(JAVA_BOOLEAN, "lambdaIsNull"))
} else {
FalseLiteral
}
val value = addMutableState(javaType(lambda.dataType), "lambdaValue")
val lambdaVar = ExprCode(isNull, JavaCode.global(value, lambda.dataType))
currentLambdaVars.put(id, lambdaVar)
lambdaVar
}

val result = f(lambdaVars)
namedLambdas.map(_.exprId.id).foreach(currentLambdaVars.remove)
result
}

def getLambdaVar(id: Long): ExprCode = {
currentLambdaVars.getOrElse(
id,
throw QueryExecutionErrors.lambdaVariableNotDefinedError(id))
}

/**
* Holding expressions' inlined mutable states like `MonotonicallyIncreasingID.count` as a
* 2-tuple: java type, variable name.
@@ -411,29 +446,11 @@ class CodegenContext extends Logging {
partitionInitializationStatements.mkString("\n")
}

/**
* Holds expressions that are equivalent. Used to perform subexpression elimination
* during codegen.
*
* For expressions that appear more than once, generate additional code to prevent
* recomputing the value.
*
* For example, consider two expression generated from this SQL statement:
* SELECT (col1 + col2), (col1 + col2) / col3.
*
* equivalentExpressions will match the tree containing `col1 + col2` and it will only
* be evaluated once.
*/
private val equivalentExpressions: EquivalentExpressions = new EquivalentExpressions

// Foreach expression that is participating in subexpression elimination, the state to use.
// Visible for testing.
private[expressions] var subExprEliminationExprs =
Map.empty[ExpressionEquals, SubExprEliminationState]

// The collection of sub-expression result resetting methods that need to be called on each row.
private val subexprFunctions = mutable.ArrayBuffer.empty[String]

val outerClassName = "OuterClass"

/**
@@ -1064,24 +1081,15 @@ class CodegenContext extends Logging {
}
}

/**
* Returns the code for subexpression elimination after splitting it if necessary.
*/
def subexprFunctionsCode: String = {
// Whole-stage codegen's subexpression elimination is handled in another code path
assert(currentVars == null || subexprFunctions.isEmpty)
splitExpressions(subexprFunctions.toSeq, "subexprFunc_split", Seq("InternalRow" -> INPUT_ROW))
}

/**
* Perform a function which generates a sequence of ExprCodes with a given mapping between
* expressions and common expressions, instead of using the mapping in current context.
* expressions and common expressions. Restores previous mapping after execution.
*/
def withSubExprEliminationExprs(
newSubExprEliminationExprs: Map[ExpressionEquals, SubExprEliminationState])(
f: => Seq[ExprCode]): Seq[ExprCode] = {
val oldsubExprEliminationExprs = subExprEliminationExprs
subExprEliminationExprs = newSubExprEliminationExprs
subExprEliminationExprs = oldsubExprEliminationExprs ++ newSubExprEliminationExprs

val genCodes = f

@@ -1090,25 +1098,26 @@ class CodegenContext extends Logging {
genCodes
}

private def collectSubExprCodes(subExprStates: Seq[SubExprEliminationState]): Seq[String] = {
subExprStates.flatMap { state =>
val codes = collectSubExprCodes(state.children) :+ state.eval.code.toString()
state.eval.code = EmptyBlock
codes
}
}

/**
* Evaluates a sequence of `SubExprEliminationState` which represent subexpressions. After
* evaluating a subexpression, this method will clean up the code block to avoid duplicate
* evaluation.
*/
def evaluateSubExprEliminationState(subExprStates: Iterable[SubExprEliminationState]): String = {
val code = new StringBuilder()

subExprStates.foreach { state =>
val currentCode = evaluateSubExprEliminationState(state.children) + "\n" + state.eval.code
code.append(currentCode + "\n")
state.eval.code = EmptyBlock
}

code.toString()
val codes = collectSubExprCodes(subExprStates.toSeq)
splitExpressionsWithCurrentInputs(codes, "subexprFunc_split")
}

/**
* Checks and sets up the state and codegen for subexpression elimination in whole-stage codegen.
* Checks and sets up the state and codegen for subexpression elimination.
*
* This finds the common subexpressions, generates the code snippets that evaluate those
* expressions and populates the mapping of common subexpressions to the generated code snippets.
@@ -1141,34 +1150,64 @@ class CodegenContext extends Logging {
* (subexpression -> `SubExprEliminationState`) into the map. So in next subexpression
* evaluation, we can look for generated subexpressions and do replacement.
*/
def subexpressionEliminationForWholeStageCodegen(expressions: Seq[Expression]): SubExprCodes = {
def subexpressionElimination(
expressions: Seq[Expression],
variablePrefix: String = ""): SubExprCodes = {
// Create a clear EquivalentExpressions and SubExprEliminationState mapping
val equivalentExpressions: EquivalentExpressions = new EquivalentExpressions
val localSubExprEliminationExprsForNonSplit =
val localSubExprEliminationExprs =
mutable.HashMap.empty[ExpressionEquals, SubExprEliminationState]

// Add each expression tree and compute the common subexpressions.
expressions.foreach(equivalentExpressions.addExprTree(_))

// Get all the expressions that appear at least twice and set up the state for subexpression
// elimination.
//
// Filter out any expressions that are already existing subexpressions. This can happen
// when finding common subexpressions inside a lambda function, and the common expression
// does not reference the lambda variables for that function, but top level attributes or
// outer lambda variables.
val commonExprs = equivalentExpressions.getCommonSubexpressions
.filter(e => !subExprEliminationExprs.contains(ExpressionEquals(e)))

val nonSplitCode = {
val allStates = mutable.ArrayBuffer.empty[SubExprEliminationState]
commonExprs.map { expr =>
withSubExprEliminationExprs(localSubExprEliminationExprsForNonSplit.toMap) {
withSubExprEliminationExprs(localSubExprEliminationExprs.toMap) {
val eval = expr.genCode(this)

val value = addMutableState(javaType(expr.dataType), s"${variablePrefix}subExprValue")

val isNullLiteral = eval.isNull match {
case TrueLiteral | FalseLiteral => true
case _ => false
}
val (isNull, isNullEvalCode) = if (!isNullLiteral) {
val v = addMutableState(JAVA_BOOLEAN, s"${variablePrefix}subExprIsNull")
(JavaCode.isNullGlobal(v), s"$v = ${eval.isNull};")
} else {
(eval.isNull, "")
}

val code = code"""
|${eval.code}
|$isNullEvalCode
|$value = ${eval.value};
"""

// Collects other subexpressions from the children.
val childrenSubExprs = mutable.ArrayBuffer.empty[SubExprEliminationState]
expr.foreach { e =>
subExprEliminationExprs.get(ExpressionEquals(e)) match {
localSubExprEliminationExprs.get(ExpressionEquals(e)) match {
case Some(state) => childrenSubExprs += state
case _ =>
}
}
val state = SubExprEliminationState(eval, childrenSubExprs.toSeq)
localSubExprEliminationExprsForNonSplit.put(ExpressionEquals(expr), state)
val state = SubExprEliminationState(
ExprCode(code, isNull, JavaCode.global(value, expr.dataType)),
childrenSubExprs.toSeq)
localSubExprEliminationExprs.put(ExpressionEquals(expr), state)
allStates += state
Seq(eval)
}
@@ -1188,38 +1227,18 @@ class CodegenContext extends Logging {
val needSplit = nonSplitCode.map(_.eval.code.length).sum > SQLConf.get.methodSplitThreshold
val (subExprsMap, exprCodes) = if (needSplit) {
if (inputVarsForAllFuncs.map(calculateParamLengthFromExprValues).forall(isValidParamLength)) {
val localSubExprEliminationExprs =
mutable.HashMap.empty[ExpressionEquals, SubExprEliminationState]

commonExprs.zipWithIndex.foreach { case (expr, i) =>
val eval = withSubExprEliminationExprs(localSubExprEliminationExprs.toMap) {
Seq(expr.genCode(this))
}.head

val value = addMutableState(javaType(expr.dataType), "subExprValue")

val isNullLiteral = eval.isNull match {
case TrueLiteral | FalseLiteral => true
case _ => false
}
val (isNull, isNullEvalCode) = if (!isNullLiteral) {
val v = addMutableState(JAVA_BOOLEAN, "subExprIsNull")
(JavaCode.isNullGlobal(v), s"$v = ${eval.isNull};")
} else {
(eval.isNull, "")
}

// Generate the code for this expression tree and wrap it in a function.
val fnName = freshName("subExpr")
val inputVars = inputVarsForAllFuncs(i)
val argList =
inputVars.map(v => s"${CodeGenerator.typeName(v.javaType)} ${v.variableName}")
val subExprState = localSubExprEliminationExprs.remove(ExpressionEquals(expr)).get
val fn =
s"""
|private void $fnName(${argList.mkString(", ")}) {
| ${eval.code}
| $isNullEvalCode
| $value = ${eval.value};
| ${subExprState.eval.code}
|}
""".stripMargin

@@ -1235,7 +1254,7 @@ class CodegenContext extends Logging {
val inputVariables = inputVars.map(_.variableName).mkString(", ")
val code = code"${addNewFunction(fnName, fn)}($inputVariables);"
val state = SubExprEliminationState(
ExprCode(code, isNull, JavaCode.global(value, expr.dataType)),
subExprState.eval.copy(code = code),
childrenSubExprs.toSeq)
localSubExprEliminationExprs.put(ExpressionEquals(expr), state)
}
@@ -1248,80 +1267,36 @@ class CodegenContext extends Logging {
throw SparkException.internalError(errMsg)
} else {
logInfo(errMsg)
(localSubExprEliminationExprsForNonSplit, Seq.empty)
(localSubExprEliminationExprs, Seq.empty)
}
}
} else {
(localSubExprEliminationExprsForNonSplit, Seq.empty)
(localSubExprEliminationExprs, Seq.empty)
}
SubExprCodes(subExprsMap.toMap, exprCodes.flatten)
}

/**
* Checks and sets up the state and codegen for subexpression elimination. This finds the
* common subexpressions, generates the functions that evaluate those expressions and populates
* the mapping of common subexpressions to the generated functions.
*/
private def subexpressionElimination(expressions: Seq[Expression]): Unit = {
// Add each expression tree and compute the common subexpressions.
expressions.foreach(equivalentExpressions.addExprTree(_))

// Get all the expressions that appear at least twice and set up the state for subexpression
// elimination.
val commonExprs = equivalentExpressions.getCommonSubexpressions
commonExprs.foreach { expr =>
val fnName = freshName("subExpr")
val isNull = addMutableState(JAVA_BOOLEAN, "subExprIsNull")
val value = addMutableState(javaType(expr.dataType), "subExprValue")

// Generate the code for this expression tree and wrap it in a function.
val eval = expr.genCode(this)
val fn =
s"""
|private void $fnName(InternalRow $INPUT_ROW) {
| ${eval.code}
| $isNull = ${eval.isNull};
| $value = ${eval.value};
|}
""".stripMargin

// Add a state and a mapping of the common subexpressions that are associate with this
// state. Adding this expression to subExprEliminationExprMap means it will call `fn`
// when it is code generated. This decision should be a cost based one.
//
// The cost of doing subexpression elimination is:
// 1. Extra function call, although this is probably *good* as the JIT can decide to
// inline or not.
// The benefit doing subexpression elimination is:
// 1. Running the expression logic. Even for a simple expression, it is likely more than 3
// above.
// 2. Less code.
// Currently, we will do this for all non-leaf only expression trees (i.e. expr trees with
// at least two nodes) as the cost of doing it is expected to be low.

val subExprCode = s"${addNewFunction(fnName, fn)}($INPUT_ROW);"
subexprFunctions += subExprCode
val state = SubExprEliminationState(
ExprCode(code"$subExprCode",
JavaCode.isNullGlobal(isNull),
JavaCode.global(value, expr.dataType)))
subExprEliminationExprs += ExpressionEquals(expr) -> state
}
}

/**
* Generates code for expressions. If doSubexpressionElimination is true, subexpression
* elimination will be performed. Subexpression elimination assumes that the code for each
* expression will be combined in the `expressions` order.
*/
def generateExpressions(
expressions: Seq[Expression],
doSubexpressionElimination: Boolean = false): Seq[ExprCode] = {
doSubexpressionElimination: Boolean = false): (Seq[ExprCode], String) = {
// We need to make sure that we do not reuse stateful expressions. This is needed for codegen
// as well because some expressions may implement `CodegenFallback`.
val cleanedExpressions = expressions.map(_.freshCopyIfContainsStatefulExpression())
if (doSubexpressionElimination) subexpressionElimination(cleanedExpressions)
cleanedExpressions.map(e => e.genCode(this))
if (doSubexpressionElimination) {
val subExprs = subexpressionElimination(cleanedExpressions)
val generatedExprs = withSubExprEliminationExprs(subExprs.states) {
cleanedExpressions.map(e => e.genCode(this))
}
val subExprCode = evaluateSubExprEliminationState(subExprs.states.values)
(generatedExprs, subExprCode)
} else {
(cleanedExpressions.map(e => e.genCode(this)), "")
}
}

/**
Original file line number Diff line number Diff line change
@@ -61,7 +61,8 @@ object GenerateMutableProjection extends CodeGenerator[Seq[Expression], MutableP
case (NoOp, _) => false
case _ => true
}
val exprVals = ctx.generateExpressions(validExpr.map(_._1), useSubexprElimination)
val (exprVals, evalSubexpr) =
ctx.generateExpressions(validExpr.map(_._1), useSubexprElimination)

// 4-tuples: (code for projection, isNull variable name, value variable name, column index)
val projectionCodes: Seq[(String, String)] = validExpr.zip(exprVals).map {
@@ -91,9 +92,6 @@ object GenerateMutableProjection extends CodeGenerator[Seq[Expression], MutableP
(code, update)
}

// Evaluate all the subexpressions.
val evalSubexpr = ctx.subexprFunctionsCode

val allProjections = ctx.splitExpressionsWithCurrentInputs(projectionCodes.map(_._1))
val allUpdates = ctx.splitExpressionsWithCurrentInputs(projectionCodes.map(_._2))

Original file line number Diff line number Diff line change
@@ -38,8 +38,8 @@ object GeneratePredicate extends CodeGenerator[Expression, BasePredicate] {
val ctx = newCodeGenContext()

// Do sub-expression elimination for predicates.
val eval = ctx.generateExpressions(Seq(predicate), useSubexprElimination).head
val evalSubexpr = ctx.subexprFunctionsCode
val (evalExprs, evalSubexpr) = ctx.generateExpressions(Seq(predicate), useSubexprElimination)
val eval = evalExprs.head

val codeBody = s"""
public SpecificPredicate generate(Object[] references) {
Loading