[SPARK-30027][SQL] Support codegen for aggregate filters in HashAggregateExec #27019
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
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
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 |
|---|---|---|
|
|
@@ -109,59 +109,39 @@ case class ProjectExec(projectList: Seq[NamedExpression], child: SparkPlan) | |
| } | ||
| } | ||
|
|
||
|
||
| trait GeneratePredicateHelper extends PredicateHelper { | ||
| self: CodegenSupport => | ||
|
|
||
| protected def generatePredicateCode( | ||
| ctx: CodegenContext, | ||
| condition: Expression, | ||
| inputAttrs: Seq[Attribute], | ||
| inputExprCode: Seq[ExprCode]): String = { | ||
| val (notNullPreds, otherPreds) = splitConjunctivePredicates(condition).partition { | ||
| case IsNotNull(a) => isNullIntolerant(a) && a.references.subsetOf(AttributeSet(inputAttrs)) | ||
| case _ => false | ||
| } | ||
| val nonNullAttrExprIds = notNullPreds.flatMap(_.references).distinct.map(_.exprId) | ||
| val outputAttrs = outputWithNullability(inputAttrs, nonNullAttrExprIds) | ||
| generatePredicateCode( | ||
| ctx, inputAttrs, inputExprCode, outputAttrs, notNullPreds, otherPreds, | ||
| nonNullAttrExprIds) | ||
| } | ||
|
|
||
| protected def generatePredicateCode( | ||
| ctx: CodegenContext, | ||
| inputAttrs: Seq[Attribute], | ||
| inputExprCode: Seq[ExprCode], | ||
| outputAttrs: Seq[Attribute], | ||
| notNullPreds: Seq[Expression], | ||
| otherPreds: Seq[Expression], | ||
| nonNullAttrExprIds: Seq[ExprId]): String = { | ||
| /** | ||
| * Generates code for `c`, using `in` for input attributes and `attrs` for nullability. | ||
| */ | ||
| def genPredicate(c: Expression, in: Seq[ExprCode], attrs: Seq[Attribute]): String = { | ||
| val bound = BindReferences.bindReference(c, attrs) | ||
| val evaluated = evaluateRequiredVariables(inputAttrs, in, c.references) | ||
|
|
||
| // Generate the code for the predicate. | ||
| val ev = ExpressionCanonicalizer.execute(bound).genCode(ctx) | ||
|
|
@@ -195,10 +175,10 @@ case class FilterExec(condition: Expression, child: SparkPlan) | |
| if (idx != -1 && !generatedIsNotNullChecks(idx)) { | ||
| generatedIsNotNullChecks(idx) = true | ||
| // Use the child's output. The nullability is what the child produced. | ||
| genPredicate(notNullPreds(idx), inputExprCode, inputAttrs) | ||
| } else if (nonNullAttrExprIds.contains(r.exprId) && !extraIsNotNullAttrs.contains(r)) { | ||
| extraIsNotNullAttrs += r | ||
| genPredicate(IsNotNull(r), inputExprCode, inputAttrs) | ||
| } else { | ||
| "" | ||
| } | ||
|
|
@@ -208,18 +188,61 @@ case class FilterExec(condition: Expression, child: SparkPlan) | |
| // enforced them with the IsNotNull checks above. | ||
| s""" | ||
| |$nullChecks | ||
| |${genPredicate(c, inputExprCode, outputAttrs)} | ||
| """.stripMargin.trim | ||
| }.mkString("\n") | ||
|
|
||
| val nullChecks = notNullPreds.zipWithIndex.map { case (c, idx) => | ||
| if (!generatedIsNotNullChecks(idx)) { | ||
| genPredicate(c, inputExprCode, inputAttrs) | ||
| } else { | ||
| "" | ||
| } | ||
| }.mkString("\n") | ||
|
|
||
| s""" | ||
| |$generated | ||
| |$nullChecks | ||
| """.stripMargin | ||
| } | ||
| } | ||
|
|
||
| /** Physical plan for Filter. */ | ||
| case class FilterExec(condition: Expression, child: SparkPlan) | ||
| extends UnaryExecNode with CodegenSupport with GeneratePredicateHelper { | ||
|
|
||
| // Split out all the IsNotNulls from condition. | ||
| private val (notNullPreds, otherPreds) = splitConjunctivePredicates(condition).partition { | ||
| case IsNotNull(a) => isNullIntolerant(a) && a.references.subsetOf(child.outputSet) | ||
| case _ => false | ||
| } | ||
|
|
||
| // The columns that will filtered out by `IsNotNull` could be considered as not nullable. | ||
| private val notNullAttributes = notNullPreds.flatMap(_.references).distinct.map(_.exprId) | ||
|
|
||
| // Mark this as empty. We'll evaluate the input during doConsume(). We don't want to evaluate | ||
| // all the variables at the beginning to take advantage of short circuiting. | ||
| override def usedInputs: AttributeSet = AttributeSet.empty | ||
|
|
||
| override def output: Seq[Attribute] = outputWithNullability(child.output, notNullAttributes) | ||
|
|
||
| override lazy val metrics = Map( | ||
| "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows")) | ||
|
|
||
| override def inputRDDs(): Seq[RDD[InternalRow]] = { | ||
| child.asInstanceOf[CodegenSupport].inputRDDs() | ||
| } | ||
|
|
||
| protected override def doProduce(ctx: CodegenContext): String = { | ||
| child.asInstanceOf[CodegenSupport].produce(ctx, this) | ||
| } | ||
|
|
||
| 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) | ||
|
|
||
| // Reset the isNull to false for the not-null columns, then the followed operators could | ||
| // generate better code (remove dead branches). | ||
| val resultVars = input.zipWithIndex.map { case (ev, i) => | ||
|
|
@@ -232,8 +255,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 { | ||
| | $predicateCode | ||
| | $numOutput.add(1); | ||
| | ${consume(ctx, resultVars)} | ||
| |} while(false); | ||
|
|
||
Oops, something went wrong.
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.
I don't understand why "while(false)" can take an effect here. Could you explain why it is needed here?
Uh oh!
There was an error while loading. Please reload this page.
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.
NOTE in above code comment would be enough to explain why, right? It still executes only once, but be able to exit the specific code block via
continueinstead of exiting the function/method in the middle of the code.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.
Got it. Thanks for the explanation.
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 for the reply, @HeartSaVioR. Since I forgot to add the generated code, I've added it in the PR description.