[SPARK-60067][SQL] Preserve CSE through scoped field updates

Let subexpression elimination inspect generated With bodies without moving unbound references, lazy definitions, or guarded values outside their scopes. Describe UpdateFields evaluation guards so shared values remain eligible without evaluating null-struct updates eagerly.

Validated 363 tests across 11 SQL, Catalyst, codegen, and optimizer suites using recompiled sources and the cached SBT classpath.

Author: bhollis-dbx

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala

@@ -23,7 +23,7 @@ import scala.collection.mutable
 
 import org.apache.spark.SparkException
 import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
-import org.apache.spark.sql.catalyst.trees.TreePattern.{LAMBDA_VARIABLE, PLAN_EXPRESSION}
+import org.apache.spark.sql.catalyst.trees.TreePattern.{COMMON_EXPR_REF, LAMBDA_VARIABLE, PLAN_EXPRESSION}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.util.Utils
 
@@ -62,7 +62,7 @@ class EquivalentExpressions(
       expr: Expression,
       map: mutable.HashMap[ExpressionEquals, ExpressionStats],
       useCount: Int = 1): Boolean = {
-    if (expr.deterministic) {
+    if (expr.deterministic && supportedExpression(expr)) {
       val wrapper = ExpressionEquals(expr)
       map.get(wrapper) match {
         case Some(stats) =>
@@ -159,12 +159,10 @@ class EquivalentExpressions(
   // that the peel lands on rather than about the `And`/`Or` chain above it.
   private def childrenToRecurse(expr: Expression): Seq[Expression] = expr match {
     case _: CodegenFallback => Nil
-    // A `CommonExpressionRef` cannot be evaluated ahead of the `With` that binds it, for the same
-    // reason a `LambdaVariable` cannot be evaluated ahead of its loop. Do not descend, so that no
-    // subtree holding a reference -- nor a `CommonExpressionDef`, which is unevaluable -- becomes
-    // a candidate. The `With` itself may still be deduplicated as a whole, which is safe: it
-    // carries its own definitions and brings their slots into scope wherever it is generated.
-    case _: With => Nil
+    // Definitions are lazy. Only the generated body exposes always-evaluated subexpressions;
+    // supportedExpression excludes candidates whose references would escape their binding scope.
+    case withExpr: With =>
+      if (withExpr.refUnderCodegenFallback) Nil else Seq(withExpr.child)
     // Peeling here is redundant for safety: `updateExprTree` peels before it descends, so this
     // method sees the peeled node either way. It stays because dropping it would pass
     // `updateExprTree` the input itself rather than the operand its peel lands on.
@@ -186,11 +184,29 @@ class EquivalentExpressions(
     // `LambdaVariable` is usually used as a loop variable, which can't be evaluated ahead of the
     // loop. So we can't evaluate sub-expressions containing `LambdaVariable` at the beginning.
     !(e.containsPattern(LAMBDA_VARIABLE) ||
+      e.isInstanceOf[CommonExpressionDef] ||
+      hasUnboundCommonExpressionRef(e) ||
       // `PlanExpression` wraps query plan. To compare query plans of `PlanExpression` on executor,
       // can cause error like NPE.
       (e.containsPattern(PLAN_EXPRESSION) && Utils.isInRunningSparkTask))
   }
 
+  private def hasUnboundCommonExpressionRef(
+      expression: Expression,
+      boundIds: Set[CommonExpressionId] = Set.empty): Boolean = {
+    if (!expression.containsPattern(COMMON_EXPR_REF)) {
+      false
+    } else {
+      expression match {
+        case ref: CommonExpressionRef => !boundIds.contains(ref.id)
+        case withExpr: With =>
+          val innerIds = boundIds ++ withExpr.defs.map(_.id)
+          withExpr.children.exists(hasUnboundCommonExpressionRef(_, innerIds))
+        case other => other.children.exists(hasUnboundCommonExpressionRef(_, boundIds))
+      }
+    }
+  }
+
   /**
    * Adds the expression to this data structure recursively. Stops if a matching expression
    * is found. That is, if `expr` has already been added, its children are not added.

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/With.scala

@@ -165,7 +165,7 @@ case class With(child: Expression, defs: Seq[CommonExpressionDef])
    * falls back (`With` is not a `CodegenFallback`, so it has to be named rather than matched). Such
    * a reference needs its cell bound and cleared; the generated code clears codegen flags instead.
    */
-  @transient private lazy val refUnderCodegenFallback: Boolean = {
+  @transient private[expressions] lazy val refUnderCodegenFallback: Boolean = {
     val ids = defs.map(_.id).toSet
     def holdsMyRef(e: Expression): Boolean = e.exists {
       case r: CommonExpressionRef => ids.contains(r.id)

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala

@@ -771,7 +771,7 @@ private case class UpdateFieldsExpression(
     structExpr: Expression,
     updatedValues: Seq[Expression],
     outputOrdinals: Seq[Int],
-    outputFields: Seq[StructField]) extends Expression {
+    outputFields: Seq[StructField]) extends ConditionalExpression {
 
   override def children: Seq[Expression] = structExpr +: updatedValues
 
@@ -781,6 +781,17 @@ private case class UpdateFieldsExpression(
   private lazy val indexedUpdatedValues = updatedValues.toArray
   private lazy val needsSource = nullable || indexedOrdinals.exists(_ >= 0)
 
+  override def alwaysEvaluatedInputs: Seq[Expression] =
+    (if (needsSource) Seq(structExpr) else Nil) ++ (if (nullable) Nil else updatedValues)
+
+  override def branchGroups: Seq[Seq[Expression]] = Nil
+
+  override def withNewAlwaysEvaluatedInputs(
+      inputs: Seq[Expression]): ConditionalExpression =
+    copy(
+      structExpr = if (needsSource) inputs.head else structExpr,
+      updatedValues = if (nullable) updatedValues else if (needsSource) inputs.tail else inputs)
+
   override lazy val dataType: StructType = {
     val sourceFields = structExpr.dataType.asInstanceOf[StructType]
     var updatedIndex = 0

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SubexpressionEliminationSuite.scala

@@ -546,12 +546,65 @@ class SubexpressionEliminationSuite extends SparkFunSuite with ExpressionEvalHel
     assert(equivalence1.getCommonSubexpressions.size == 1)
   }
 
-  test("SPARK-58902: no candidate below a With, with or without the short-circuit peel") {
+  test("With exposes independent subexpressions but not unbound references") {
+    val input = AttributeReference("input", IntegerType)()
+    val shared = Add(input, Literal(1))
+    val expression = With(input) { case Seq(ref) =>
+      Add(ref, Multiply(shared, shared))
+    }
+    val equivalence = new EquivalentExpressions
+    equivalence.addExprTree(expression)
+    assert(equivalence.getCommonSubexpressions == Seq(shared))
+    assert(equivalence.getAllExprStates().forall { state =>
+      state.expr == expression || !state.expr.exists(_.isInstanceOf[CommonExpressionRef])
+    })
+  }
+
+  test("With candidates must bind references in their lexical scope") {
+    val input = AttributeReference("input", IntegerType)()
+    val shared = Add(input, Literal(1))
+    var inner: Expression = null
+    val outer = With(input) { case Seq(outerRef) =>
+      inner = With(input) { case Seq(innerRef) =>
+        Add(Add(outerRef, innerRef), shared)
+      }
+      Add(inner, shared)
+    }
+    val equivalence = new EquivalentExpressions
+    equivalence.addExprTree(outer)
+    assert(equivalence.getExprState(outer).nonEmpty)
+    assert(equivalence.getExprState(inner).isEmpty)
+    assert(equivalence.getCommonSubexpressions == Seq(shared))
+    assert(!equivalence.addExpr(inner))
+    val closed = With(input) { case Seq(ref) => Add(ref, Literal(1)) }
+    val sibling = new CommonExpressionRef(closed.defs.head)
+    assert(!equivalence.addExpr(Add(closed, sibling)))
+  }
+
+  test("With does not expose lazy definitions or interpreted children") {
+    val input = AttributeReference("input", IntegerType)()
+    val shared = Add(input, Literal(1))
+    val lazyDefinition = With(Multiply(shared, shared)) { case Seq(ref) =>
+      If(Literal(false), ref, input)
+    }
+    val fallback = With(input) { case Seq(ref) =>
+      Add(CodegenFallbackExpression(ref), Multiply(shared, shared))
+    }
+    val guarded = With(input) { case Seq(ref) =>
+      If(GreaterThan(ref, Literal(0)), Multiply(shared, shared), Literal(0))
+    }
+    Seq(lazyDefinition, fallback, guarded).foreach { expression =>
+      val equivalence = new EquivalentExpressions
+      equivalence.addExprTree(expression)
+      assert(equivalence.getCommonSubexpressions.isEmpty)
+    }
+  }
+
+  test("SPARK-58902: candidates below a With cannot hold unbound references") {
     // A `CommonExpressionRef` can only be evaluated inside the `With` that binds it: the codegen
     // slots exist only while that `With` is being generated, and `getCommonExpr` throws otherwise.
     // Subexpression elimination generates its candidates outside every `With` scope, so a candidate
-    // holding a reference would fail the query. `childrenToRecurse` therefore stops at a `With`,
-    // including after `skipForShortcut` has peeled an `And`/`Or` down to one.
+    // holding an unbound reference would fail the query.
     val a = AttributeReference("a", IntegerType)()
     val memoized = With(Add(a, a)) { case Seq(ref) =>
       And(GreaterThan(ref, Literal(0)), LessThan(ref, Literal(10)))

sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala

@@ -504,6 +504,70 @@ class ColumnExpressionSuite extends SharedSparkSession {
     }
   }
 
+  test("field updates should preserve shared value subexpression elimination") {
+    withSQLConf(SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "true",
+      SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true") {
+      for {
+        nested <- Seq(false, true)
+        external <- Seq(false, true)
+        opaque <- Seq(false, true)
+      } {
+        val calls = spark.sparkContext.longAccumulator
+        val sourceCalls = spark.sparkContext.longAccumulator
+        val value = udf((id: Long) => {
+          calls.add(1)
+          id + 1
+        }).apply($"id")
+        val source = if (opaque) {
+          udf((id: Long) => {
+            sourceCalls.add(1)
+            ((id, id), id)
+          }).asNonNullable().asNondeterministic()($"id")
+        } else {
+          struct(struct($"id", $"id").as("_1"), $"id".as("_2"))
+        }
+        val prefix = if (nested) "_1." else ""
+        val updated = source.withField(s"${prefix}first", value)
+          .withField(s"${prefix}second", value)
+        val columns = if (external) Seq(updated, value) else Seq(updated)
+        val rows = spark.range(0, 10, 1, 1).select(columns: _*).collect()
+        rows.zipWithIndex.foreach { case (row, index) =>
+          val result = if (nested) row.getStruct(0).getStruct(0) else row.getStruct(0)
+          assert(result.getLong(2) == index + 1L)
+          assert(result.getLong(3) == index + 1L)
+          if (external) assert(row.getLong(1) == index + 1L)
+        }
+        val expectedCalls = if (nested && opaque) 10 * (if (external) 3 else 2) else 10
+        assert(calls.value == expectedCalls,
+          s"nested=$nested external=$external opaque=$opaque calls=${calls.value}")
+        if (opaque) assert(sourceCalls.value == 10)
+      }
+    }
+  }
+
+  test("field updates should not hoist guarded shared values") {
+    onEachEvalPath {
+      withSQLConf(SQLConf.ANSI_ENABLED.key -> "true",
+        SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "true") {
+        val schema = StructType(Seq(StructField("a", IntegerType, nullable = false)))
+        val source = when($"id" === 0, lit(null).cast(schema))
+          .otherwise(struct(lit(1).as("a")))
+        val shared = lit(1L) / $"id"
+        checkAnswer(
+          spark.range(0, 2, 1, 1).select(
+            source.withField("b", shared).withField("c", shared)),
+          Seq(Row(null), Row(Row(1, 1.0, 1.0))))
+        val throwingSource = udf((id: Long) => {
+          if (id >= 0) throw new IllegalStateException(s"evaluated for $id") else (id, id)
+        }).asNonNullable()
+        checkAnswer(
+          spark.range(1).select(
+            throwingSource($"id").withField("_1", lit(1L)).withField("_2", lit(2L))),
+          Row(Row(1L, 2L)))
+      }
+    }
+  }
+
   test("nested field updates on unresolved columns should remain linear") {
     val df = Seq((((1, 2), 3), 0)).toDF("s", "z")
     Seq(4, 8, 12).foreach { count =>