(spark) branch master updated: [SPARK-56032][SQL] Support subexpression elimination in `FilterExec` whole-stage codegen

[do...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

dongjoon pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/master by this push:
     new 21d69c4f9098 [SPARK-56032][SQL] Support subexpression elimination in `FilterExec` whole-stage codegen
21d69c4f9098 is described below

commit 21d69c4f9098c2081b1bb22eb6be35112ec4809c
Author: yangjie01 <ya...@baidu.com>
AuthorDate: Fri Mar 20 05:31:11 2026 -0700

    [SPARK-56032][SQL] Support subexpression elimination in `FilterExec` whole-stage codegen
    
    ### What changes were proposed in this pull request?
    This PR adds subexpression elimination support to `FilterExec.doConsume` in the whole-stage codegen path.
    
    The implementation follows the established CSE pattern used by `ProjectExec`, `HashAggregateExec`, and `AggregateCodegenSupport`:
    
    1. Collect `otherPreds` (excluding `notNullPreds`, which are simple `IsNotNull` checks with no CSE value) and run `ctx.subexpressionEliminationForWholeStageCodegen` to discover common subexpressions.
    2. Wrap `generatePredicateCode` inside `ctx.withSubExprEliminationExprs` so that `genCode` calls within can look up pre-computed subexpressions.
    3. Emit the CSE pre-computation code via `ctx.evaluateSubExprEliminationState` before the predicate checks.
    
    Key design decisions:
    - **Only `otherPreds` participate in CSE** — `notNullPreds` are guaranteed to be `IsNotNull` checks (by FilterExec's constructor partitioning logic) with no CSE value; including them would interfere with equivalence analysis.
    - **Pre-evaluate input variables before CSE analysis** — `FilterExec` sets `usedInputs = AttributeSet.empty` to defer input evaluation for short-circuit optimization. However, `subexpressionEliminationForWholeStageCodegen`'s internal `getLocalInputVariableValues` has a side effect: it clears `ctx.currentVars[i].code` for input variables referenced by common subexpressions. If `notNullPreds` reference the same input variables, `generatePredicateCode`'s `evaluateRequiredVariables` would [...]
    - **CSE evaluation is placed before predicate short-circuit checks** — This is an intentional tradeoff: common subexpressions are evaluated unconditionally even if an earlier `notNull` check would have short-circuited. For expensive shared expressions (e.g., `from_json` appearing in 500 predicates), the benefit of evaluating once vs N times far outweighs the cost of losing short-circuit on the CSE portion. When there are no common subexpressions, `subExprsCode` is empty and this path  [...]
    
    ### Why are the changes needed?
    Support subexpression elimination in FilterExec whole-stage codegen, and with this fix, the `from_json` codegen (PR #48466) can be safely re-enabled in a follow-up PR.
    
    ### Does this PR introduce _any_ user-facing change?
    No.
    
    ### How was this patch tested?
    Added two new test cases in `WholeStageCodegenSuite`
    
    ### Was this patch authored or co-authored using generative AI tooling?
    Generated-by: Claude Sonnet 4.6
    
    Closes #54862 from LuciferYang/SPARK-56032.
    
    Lead-authored-by: yangjie01 <ya...@baidu.com>
    Co-authored-by: YangJie <ya...@baidu.com>
    Signed-off-by: Dongjoon Hyun <do...@apache.org>
---
 .../sql/execution/basicPhysicalOperators.scala     | 31 ++++++++++++++--
 .../sql/execution/WholeStageCodegenSuite.scala     | 42 +++++++++++++++++++++-
 2 files changed, 70 insertions(+), 3 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala
index f2cbfd2154a4..49da71e2b0af 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala
@@ -253,8 +253,33 @@ 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)
+    // Apply CSE to otherPreds only (notNullPreds are simple IsNotNull checks with no CSE value).
+    val (inputVarsCode, subExprsCode, predicateCode) =
+      if (conf.subexpressionEliminationEnabled && otherPreds.nonEmpty) {
+        val boundOtherPreds = otherPreds.map(
+          BindReferences.bindReference(_, output))
+        // Pre-evaluate input variables before CSE analysis: CSE clears
+        // ctx.currentVars[i].code as a side effect, which causes "is not an rvalue"
+        // errors if notNullPreds reference the same inputs.
+        val otherPredInputAttrs = AttributeSet(otherPreds.flatMap(_.references))
+        val inputVarsEvalCode = evaluateRequiredVariables(
+          child.output, input, otherPredInputAttrs)
+
+        val subExprs = ctx.subexpressionEliminationForWholeStageCodegen(boundOtherPreds)
+        val predCode: String = {
+          var code = ""
+          ctx.withSubExprEliminationExprs(subExprs.states) {
+            code = generatePredicateCode(
+              ctx, child.output, input, output, notNullPreds, otherPreds, notNullAttributes)
+            Seq.empty
+          }
+          code
+        }
+        (inputVarsEvalCode, ctx.evaluateSubExprEliminationState(subExprs.states.values), predCode)
+      } else {
+        ("", "", 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).
@@ -268,6 +293,8 @@ case class FilterExec(condition: Expression, child: SparkPlan)
     // Note: wrap in "do { } while (false);", so the generated checks can jump out with "continue;"
     s"""
        |do {
+       |  $inputVarsCode
+       |  $subExprsCode
        |  $predicateCode
        |  $numOutput.add(1);
        |  ${consume(ctx, resultVars)}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala
index 6004f6f76394..08e0a92e516a 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala
@@ -25,11 +25,12 @@ import org.apache.spark.sql.catalyst.expressions.codegen.{ByteCodeStats, CodeAnd
 import org.apache.spark.sql.execution.adaptive.DisableAdaptiveExecutionSuite
 import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, SortAggregateExec}
 import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec
+import org.apache.spark.sql.execution.debug.codegenString
 import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, ShuffledHashJoinExec, SortMergeJoinExec}
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.SharedSparkSession
-import org.apache.spark.sql.types.{IntegerType, StringType, StructType}
+import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType}
 
 // Disable AQE because the WholeStageCodegenExec is added when running QueryStageExec
 class WholeStageCodegenSuite extends QueryTest with SharedSparkSession
@@ -944,4 +945,43 @@ class WholeStageCodegenSuite extends QueryTest with SharedSparkSession
       }
     }
   }
+
+  test("SPARK-56032: subexpression elimination in FilterExec codegen") {
+    val schema = StructType(Seq(
+      StructField("a", IntegerType, nullable = true),
+      StructField("b", IntegerType, nullable = true)))
+    val data = spark.sparkContext.parallelize(Seq(
+      Row(1, 5), Row(null, 3), Row(4, null), Row(5, 6), Row(7, 8), Row(2, 3)))
+
+    val expected = Seq(Row(1, 5), Row(5, 6), Row(2, 3))
+
+    Seq("1", Int.MaxValue.toString).foreach { splitThreshold =>
+      def testFilterCSE(cseEnabled: Boolean): String = {
+        withSQLConf(
+          SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> cseEnabled.toString,
+          SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true",
+          SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+          SQLConf.CODEGEN_METHOD_SPLIT_THRESHOLD.key -> splitThreshold) {
+          val df = spark.createDataFrame(data, schema)
+          val filtered = df.where(
+            "a IS NOT NULL AND (a + b) > 3 AND (a + b) < 15 AND (a + b) != 10")
+          val plan = filtered.queryExecution.executedPlan
+          assert(plan.exists(_.isInstanceOf[WholeStageCodegenExec]),
+            "Filter should be in whole-stage codegen")
+          checkAnswer(filtered, expected)
+          codegenString(plan)
+        }
+      }
+
+      val cseCode = testFilterCSE(cseEnabled = true)
+      val noCseCode = testFilterCSE(cseEnabled = false)
+
+      val addExactPattern = "addExact".r
+      val cseAddCount = addExactPattern.findAllIn(cseCode).length
+      val noCseAddCount = addExactPattern.findAllIn(noCseCode).length
+      assert(cseAddCount < noCseAddCount,
+        s"CSE should reduce repeated evaluation (splitThreshold=$splitThreshold): " +
+          s"addExact appears $cseAddCount times with CSE vs $noCseAddCount times without")
+    }
+  }
 }


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org