Skip to content

Commit

Permalink
[Only Test][Don't Review] FilterExec Codegen SubexpressionElimination
Browse files Browse the repository at this point in the history
  • Loading branch information
panbingkun committed Jan 20, 2025
1 parent 4443078 commit e167903
Showing 1 changed file with 22 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -248,9 +248,27 @@ case class FilterExec(condition: Expression, child: SparkPlan)

override def doConsume(ctx: CodegenContext, input: Seq[ExprCode], row: ExprCode): String = {
val numOutput = metricTerm(ctx, "numOutputRows")

val predicateCode = generatePredicateCode(
ctx, child.output, input, output, notNullPreds, otherPreds, notNullAttributes)
val (predicateCode, localValInputs) = if (conf.subexpressionEliminationEnabled) {
val bound = BindReferences.bindReference(condition, child.output)
val exprs = Seq(bound)
val subExprs = ctx.subexpressionEliminationForWholeStageCodegen(exprs)
val genVars = ctx.withSubExprEliminationExprs(subExprs.states) {
exprs.map(_.genCode(ctx))
}
val ev = genVars.head
val nullCheck = if (bound.nullable) { s"${ev.isNull} || " } else { "" }
val predicateCode =
s"""
|${ctx.evaluateSubExprEliminationState(subExprs.states.values)}
|${genVars.map(_.code.code).mkString("\n")}
|if ($nullCheck!${ev.value}) continue;
""".stripMargin
(predicateCode, subExprs.exprCodesNeedEvaluate)
} else {
val predicateCode = generatePredicateCode(
ctx, child.output, input, output, notNullPreds, otherPreds, notNullAttributes)
(predicateCode, Seq.empty)
}

// Reset the isNull to false for the not-null columns, then the followed operators could
// generate better code (remove dead branches).
Expand All @@ -264,6 +282,7 @@ case class FilterExec(condition: Expression, child: SparkPlan)
// Note: wrap in "do { } while (false);", so the generated checks can jump out with "continue;"
s"""
|do {
| ${evaluateVariables(localValInputs)}
| $predicateCode
| $numOutput.add(1);
| ${consume(ctx, resultVars)}
Expand Down

0 comments on commit e167903

Please sign in to comment.