You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by we...@apache.org on 2026/08/06 08:24:02 UTC
(spark) branch master updated: [SPARK-58501][SQL] Eliminate common subexpressions in ExpandExec
This is an automated email from the ASF dual-hosted git repository.
cloud-fan 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 97e8de6a443e [SPARK-58501][SQL] Eliminate common subexpressions in ExpandExec
97e8de6a443e is described below
commit 97e8de6a443e83585863dfac750c3dcfb482c8f6
Author: James Xu <xu...@gmail.com>
AuthorDate: Thu Aug 6 16:23:46 2026 +0800
[SPARK-58501][SQL] Eliminate common subexpressions in ExpandExec
### What changes were proposed in this pull request?
Wire `ExpandExec` into the whole-stage codegen subexpression elimination framework:
- All branch expressions are bound once up front and analyzed together. A subexpression shared across branches (or repeated within a branch) is evaluated once per input row, before the branch loop, and each branch references the cached value. This is semantics-preserving because every branch of an Expand consumes the same input row.
- Both the branch-invariant column generation and the per-branch switch/case generation resolve repeated subtrees to the cached values.
- The switch/case function splitting (SPARK-35329) now passes the eliminated subexpression variables into the split functions as parameters, so large Expands keep compiling correctly.
- Gated by the existing `spark.sql.subexpressionElimination.enabled` conf (no new configuration).
- Adds an `ExpandBenchmark` case.
### Why are the changes needed?
Conditional-aggregate rollup queries (e.g. "N-day active users" dashboards) stack many conditional aggregates whose conditions share one expensive subexpression:
```sql
SELECT
COUNT(DISTINCT IF(datediff(date '2026-01-01',
from_unixtime(unix_timestamp(ts, 'yyyy-MM-dd HH:mm:ss'))) <= 1, uid, NULL)) AS uv_1d,
COUNT(DISTINCT IF(datediff(date '2026-01-01',
from_unixtime(unix_timestamp(ts, 'yyyy-MM-dd HH:mm:ss'))) <= 7, uid, NULL)) AS uv_7d,
SUM(IF(datediff(date '2026-01-01',
from_unixtime(unix_timestamp(ts, 'yyyy-MM-dd HH:mm:ss'))) <= 1, 1, 0)) AS pv_1d
-- ... more conditional aggregates over longer windows
FROM traffic
```
`RewriteDistinctAggregates` gives each distinct group its own Expand branch with the condition expression verbatim, so the shared subexpression is compiled into every branch body and re-evaluated once per branch per input row. In the benchmark query above it is evaluated 18 times per input row (9 distinct-aggregate conditions + 9 regular-aggregate conditions across 10 branches); with this PR it is evaluated once.
### Does this PR introduce _any_ user-facing change?
No. Query results are unchanged; only the number of evaluations per input row changes.
### How was this patch tested?
- Extended the `WholeStageCodegenSuite` test "Expand should eliminate common subexpressions across branches" (SQL-based, via conditional `COUNT DISTINCT` aggregates):
- asserts the shared subexpression is evaluated once per input row with elimination enabled, and once per branch with it disabled;
- forces both switch/case code paths deterministically via `spark.sql.codegen.methodSplitThreshold` (a tiny threshold forces the SPARK-35329 function splitting, a huge threshold keeps the bodies inline), asserting on the generated `switchCaseCode` functions and verifying correctness on both paths.
- Ran the full `WholeStageCodegenSuite` (57 tests passed).
- Benchmark: new `ExpandBenchmark` case modeling the rollup above, 5M rows, 3 iterations, Apple M4 Pro / JDK 17 (local run): shared subexpression evaluations per input row 18 -> 1; average runtime 82485 ms -> 51457 ms (15701 -> 9756 ns/row), 1.6X faster.
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #57703 from xumingming/expand-cse.
Authored-by: James Xu <xu...@gmail.com>
Signed-off-by: Wenchen Fan <we...@databricks.com>
---
.../apache/spark/sql/execution/ExpandExec.scala | 88 +++++++++++++++-------
.../sql/execution/WholeStageCodegenSuite.scala | 63 +++++++++++++++-
.../sql/execution/benchmark/ExpandBenchmark.scala | 69 +++++++++++++++++
3 files changed, 191 insertions(+), 29 deletions(-)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/ExpandExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/ExpandExec.scala
index ba1238564348..c40e79e61d78 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/ExpandExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/ExpandExec.scala
@@ -137,6 +137,11 @@ case class ExpandExec(
*
* We use a for loop here so we only includes one copy of the consume code and avoid code
* size explosion.
+ *
+ * In addition, when subexpression elimination is enabled, common subexpressions shared by
+ * the branch expressions (e.g. an expensive condition repeated in many branches) are
+ * evaluated only once per input row, before the loop, since all the branches consume the
+ * same input row.
*/
// Tracks whether a column has the same output for all rows.
@@ -147,42 +152,62 @@ case class ExpandExec(
projections.map(p => p(colIndex)).toSet.size == 1
}.toArray
+ // Bind all the branch expressions once up front, so that identical expressions appearing in
+ // different branches bind to identical trees and can be deduplicated by subexpression
+ // elimination below.
+ val boundProjections: Seq[Seq[Expression]] = projections.map { exprs =>
+ BindReferences.bindReferences(exprs, child.output)
+ }
+
+ // Set up subexpression elimination over all the branch expressions. This deduplicates
+ // repeated subexpressions both within a branch and across branches. The code evaluating the
+ // common subexpressions is emitted once before the branch loop (see the end of this method).
+ val subExprs: SubExprCodes = if (conf.subexpressionEliminationEnabled) {
+ ctx.subexpressionEliminationForWholeStageCodegen(boundProjections.flatten)
+ } else {
+ SubExprCodes(Map.empty, Seq.empty)
+ }
+
// Part 1: declare variables for each column
// If a column has the same value for all output rows, then we also generate its computation
// right after declaration. Otherwise its value is computed in the part 2.
- lazy val attributeSeq: AttributeSeq = child.output
- val outputColumns = output.indices.map { col =>
- val firstExpr = projections.head(col)
- if (sameOutput(col)) {
- // This column is the same across all output rows. Just generate code for it here.
- BindReferences.bindReference(firstExpr, attributeSeq).genCode(ctx)
- } else {
- val isNull = ctx.addMutableState(
- CodeGenerator.JAVA_BOOLEAN,
- "resultIsNull",
- v => s"$v = true;")
- val value = ctx.addMutableState(
- CodeGenerator.javaType(firstExpr.dataType),
- "resultValue",
- v => s"$v = ${CodeGenerator.defaultValue(firstExpr.dataType)};")
+ val outputColumns = ctx.withSubExprEliminationExprs(subExprs.states) {
+ output.indices.map { col =>
+ val firstExpr = boundProjections.head(col)
+ if (sameOutput(col)) {
+ // This column is the same across all output rows. Just generate code for it here.
+ firstExpr.genCode(ctx)
+ } else {
+ val isNull = ctx.addMutableState(
+ CodeGenerator.JAVA_BOOLEAN,
+ "resultIsNull",
+ v => s"$v = true;")
+ val value = ctx.addMutableState(
+ CodeGenerator.javaType(firstExpr.dataType),
+ "resultValue",
+ v => s"$v = ${CodeGenerator.defaultValue(firstExpr.dataType)};")
- ExprCode(
- JavaCode.isNullVariable(isNull),
- JavaCode.variable(value, firstExpr.dataType))
+ ExprCode(
+ JavaCode.isNullVariable(isNull),
+ JavaCode.variable(value, firstExpr.dataType))
+ }
}
}
// Part 2: switch/case statements
- val switchCaseExprs = projections.zipWithIndex.map { case (exprs, row) =>
- val (exprCodesWithIndices, inputVarSets) = exprs.indices.flatMap { col =>
- if (!sameOutput(col)) {
- val boundExpr = BindReferences.bindReference(exprs(col), attributeSeq)
- val exprCode = boundExpr.genCode(ctx)
- val inputVars = CodeGenerator.getLocalInputVariableValues(ctx, boundExpr)._1
- Some(((col, exprCode), inputVars))
- } else {
- None
- }
+ val switchCaseExprs = projections.indices.map { row =>
+ val colsToGenerate = projections(row).indices.filter(col => !sameOutput(col))
+ val exprCodes = ctx.withSubExprEliminationExprs(subExprs.states) {
+ colsToGenerate.map(col => boundProjections(row)(col).genCode(ctx))
+ }
+ val (exprCodesWithIndices, inputVarSets) = colsToGenerate.zip(exprCodes).map {
+ case (col, exprCode) =>
+ // Pass `subExprs.states` so that the input variables of the split switch/case
+ // functions below include the variables holding the common subexpression values,
+ // which are evaluated outside of the split functions.
+ val inputVars = CodeGenerator.getLocalInputVariableValues(
+ ctx, boundProjections(row)(col), subExprs.states)._1
+ ((col, exprCode), inputVars)
}.unzip
val inputVars = inputVarSets.foldLeft(Set.empty[VariableValue])(_ ++ _)
@@ -239,7 +264,14 @@ case class ExpandExec(
val i = ctx.freshName("i")
// these column have to declared before the loop.
val evaluate = evaluateVariables(outputColumns)
+ // The input variables used by the common subexpressions have to be evaluated first, then
+ // the common subexpressions themselves, both before the loop since every branch consumes
+ // the same input row.
+ val evaluateSubExprInputs = evaluateVariables(subExprs.exprCodesNeedEvaluate)
+ val evaluateSubExprs = ctx.evaluateSubExprEliminationState(subExprs.states.values)
s"""
+ |$evaluateSubExprInputs
+ |$evaluateSubExprs
|$evaluate
|for (int $i = 0; $i < ${projections.length}; $i ++) {
| switch ($i) {
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 3f0c0ab97c40..85e961d71ff9 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
@@ -28,7 +28,7 @@ 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.debug.{codegenString, codegenStringSeq}
import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, ShuffledHashJoinExec, SortMergeJoinExec}
import org.apache.spark.sql.functions._
import org.apache.spark.sql.internal.SQLConf
@@ -1556,4 +1556,65 @@ class WholeStageCodegenSuite extends SharedSparkSession
s"subexpressionElimination.filterExec.enabled should reduce repeated evaluation: " +
s"addExact appears $enabledCount times when enabled vs $disabledCount times when disabled")
}
+
+ test("Expand should eliminate common subexpressions across branches") {
+ // The conditions of the two conditional COUNT DISTINCT aggregates share `sinh(v)`, which
+ // `RewriteDistinctAggregates` places in different branches of the Expand. Since all the
+ // branches of an Expand consume the same input row, with subexpression elimination
+ // `sinh(v)` should be evaluated once per input row (before the branch loop) instead of
+ // once per branch.
+ def runQuery(): String = {
+ val df = spark.range(10)
+ .selectExpr("id as k", "cast(id as double) as v")
+ .selectExpr(
+ "count(DISTINCT IF(sinh(v) > 5.0, k, NULL))",
+ "count(DISTINCT IF(sinh(v) > 50.0, v, NULL))")
+ val plan = df.queryExecution.executedPlan
+ assert(plan.exists(_.isInstanceOf[ExpandExec]), "Expand is expected")
+ checkAnswer(df, Row(7L, 5L))
+ codegenStringSeq(plan).map(_._2).mkString
+ }
+
+ val sinhPattern = "java\\.lang\\.Math\\.sinh".r
+ val enabledCode = withSQLConf(SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "true") {
+ runQuery()
+ }
+ val disabledCode = withSQLConf(SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "false") {
+ runQuery()
+ }
+ assert(sinhPattern.findAllIn(enabledCode).length == 1,
+ "sinh(v) should be evaluated only once per input row with subexpression elimination")
+ assert(sinhPattern.findAllIn(disabledCode).length == 2,
+ "sinh(v) should be evaluated once per branch without subexpression elimination")
+
+ // Whether the switch/case bodies are split into separate functions (SPARK-35329)
+ // depends on the amount of code generated for the branches, so force both code paths
+ // explicitly instead of relying on the default methodSplitThreshold. With a tiny
+ // threshold the bodies are split and the eliminated subexpression variables have to
+ // be passed to the split functions as parameters; with a huge threshold the bodies
+ // stay inline and reference the variables directly. ExpandExec names each split
+ // function via `ctx.freshName("switchCaseCode")` (e.g. `switchCaseCode_0`) and emits
+ // both its definition and call site into the generated source, so the substring
+ // "switchCaseCode" appears in the generated code if and only if a split happened.
+ val splitCode = withSQLConf(
+ SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "true",
+ SQLConf.CODEGEN_METHOD_SPLIT_THRESHOLD.key -> "1") {
+ runQuery()
+ }
+ assert(splitCode.contains("switchCaseCode"),
+ "switch/case bodies should be split into separate functions with a tiny " +
+ "methodSplitThreshold")
+ assert(sinhPattern.findAllIn(splitCode).length == 1,
+ "sinh(v) should be evaluated only once per input row with function splitting")
+
+ val inlineCode = withSQLConf(
+ SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "true",
+ SQLConf.CODEGEN_METHOD_SPLIT_THRESHOLD.key -> "1000000") {
+ runQuery()
+ }
+ assert(!inlineCode.contains("switchCaseCode"),
+ "switch/case bodies should stay inline with a large methodSplitThreshold")
+ assert(sinhPattern.findAllIn(inlineCode).length == 1,
+ "sinh(v) should be evaluated only once per input row without function splitting")
+ }
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/ExpandBenchmark.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/ExpandBenchmark.scala
index 27e67e6b3634..586dfec19620 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/ExpandBenchmark.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/ExpandBenchmark.scala
@@ -33,6 +33,11 @@ import org.apache.spark.sql.internal.SQLConf
* OptimizeExpand rule that pre-aggregates data before the Expand. Controlled by
* spark.sql.optimizer.optimizeExpandRatio (default -1 = disabled).
*
+ * It also measures subexpression elimination in Expand: conditional-aggregate
+ * queries whose conditions all share the same expensive subexpression, which
+ * ends up in every branch of the Expand. Controlled by
+ * spark.sql.subexpressionElimination.enabled (default true).
+ *
* To run this benchmark:
* {{{
* 1. build/sbt "sql/Test/runMain <this class>"
@@ -119,6 +124,62 @@ object ExpandBenchmark extends SqlBasedBenchmark {
benchmark.run()
}
+ /**
+ * Prepares a traffic-like fact table: one row per page view, with a user id,
+ * a date-string partition column `pt` (format `yyyyMMdd`), and a metric value.
+ */
+ private def prepareTrafficTable(name: String, N: Long, userMod: Int): Unit = {
+ spark.range(N)
+ .selectExpr(
+ s"cast(id % $userMod as bigint) as user_id",
+ "date_format(date_add(date '2023-06-01', cast(id % 730 as int)), 'yyyyMMdd') as pt",
+ "cast(id % 3600 as bigint) as value")
+ .createOrReplaceTempView(name)
+ }
+
+ /**
+ * Builds a traffic/BI "N-day active users" rollup query: for each of the
+ * given day windows, one conditional COUNT(DISTINCT) and one conditional SUM
+ * whose conditions all share the same expensive datetime subexpression
+ * (parse `pt`, format it back, then `datediff` against a fixed date).
+ * `RewriteDistinctAggregates` places those conditions in the branches of an
+ * Expand, so the shared subexpression is codegen'd once per branch without
+ * subexpression elimination.
+ */
+ private def conditionalAggsQuery(table: String, windows: Seq[Int]): String = {
+ val daysSince = "datediff(" +
+ "from_unixtime(unix_timestamp('20250601', 'yyyyMMdd')), " +
+ "from_unixtime(unix_timestamp(pt, 'yyyyMMdd')))"
+ val metrics = windows.flatMap { n =>
+ Seq(
+ s"count(distinct if($daysSince <= $n, pt, null)) as count_${n}d",
+ s"sum(if($daysSince <= $n, value, null)) as sum_${n}d")
+ }
+ s"""SELECT user_id,
+ | ${metrics.mkString(",\n ")}
+ |FROM $table
+ |GROUP BY user_id""".stripMargin
+ }
+
+ private def subExprEliminationBenchmark(
+ title: String, N: Long, table: String): Unit = {
+ val benchmark = new Benchmark(title, N, output = output)
+ val sql = conditionalAggsQuery(table, Seq(1, 7, 14, 30, 60, 90, 180, 365, 730))
+
+ benchmark.addCase("CSE disabled (re-evaluate per branch)", numIters = 3) { _ =>
+ withSQLConf(SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "false") {
+ spark.sql(sql).noop()
+ }
+ }
+ benchmark.addCase("CSE enabled (evaluate once per row)", numIters = 3) { _ =>
+ withSQLConf(SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "true") {
+ spark.sql(sql).noop()
+ }
+ }
+
+ benchmark.run()
+ }
+
override def runBenchmarkSuite(mainArgs: Array[String]): Unit = {
val N = 10L << 20 // ~10M rows
@@ -212,5 +273,13 @@ object ExpandBenchmark extends SqlBasedBenchmark {
benchDataChar.run()
}
+
+ runBenchmark("Expand: subexpression elimination across branches") {
+ val numRows = 5L << 20 // ~5M rows
+ prepareTrafficTable("expand_traffic", numRows, userMod = 1000000)
+ subExprEliminationBenchmark(
+ "9 conditional COUNT(DISTINCT) + 9 conditional SUM sharing one subexpression",
+ numRows, "expand_traffic")
+ }
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org