From 6ccb18b53c1deab69eb03cc140b38b3d333a2f30 Mon Sep 17 00:00:00 2001
From: Ben Hollis <ben.hollis@databricks.com>
Date: Wed, 7 Oct 2026 17:41:29 +0000
Subject: [PATCH 1/8] Keep UpdateFields compact through execution

Lower UpdateFields to a compact evaluator that carries unchanged field ordinals as metadata and generates code only for changed expressions. This prevents Catalyst expression growth with wide structs while preserving field extraction and evaluation semantics.

Add width-sensitive, short-circuit, and omitted-field regression coverage.
---
 .../expressions/complexTypeCreator.scala      | 269 ++++++++++++++++--
 .../sql/catalyst/optimizer/ComplexTypes.scala |   7 +-
 .../sql/catalyst/optimizer/Optimizer.scala    |   4 +-
 .../sql/catalyst/optimizer/UpdateFields.scala |   4 +-
 .../ReplaceUpdateFieldsExpressionSuite.scala  |  92 ++++++
 .../spark/sql/ColumnExpressionSuite.scala     | 138 ++++++++-
 6 files changed, 486 insertions(+), 28 deletions(-)
 create mode 100644 sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceUpdateFieldsExpressionSuite.scala

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala
index 2f5b05af043d3..4e1bcad3b9a80 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala
@@ -28,7 +28,7 @@ import org.apache.spark.sql.catalyst.expressions.Cast._
 import org.apache.spark.sql.catalyst.expressions.codegen._
 import org.apache.spark.sql.catalyst.expressions.codegen.Block._
 import org.apache.spark.sql.catalyst.parser.CatalystSqlParser
-import org.apache.spark.sql.catalyst.trees.{LeafLike, UnaryLike}
+import org.apache.spark.sql.catalyst.trees.{LeafLike, TreeNodeTag, UnaryLike}
 import org.apache.spark.sql.catalyst.trees.TreePattern._
 import org.apache.spark.sql.catalyst.util._
 import org.apache.spark.sql.errors.QueryCompilationErrors
@@ -764,12 +764,161 @@ case class DropField(name: String) extends StructFieldsOperation with LeafLike[E
     values.filterNot { case (field, _) => resolver(field.name, name) }
 }
 
+private case class UpdateFieldsExpression(
+    structExpr: Expression,
+    updatedValues: Seq[Expression],
+    // A negative ordinal marks an updated output, in updatedValues order.
+    outputOrdinals: Seq[Int],
+    outputFields: Seq[StructField]) extends Expression {
+
+  override def children: Seq[Expression] = structExpr +: updatedValues
+
+  override def nullable: Boolean = structExpr.nullable
+
+  override lazy val dataType: StructType = {
+    val sourceFields = structExpr.dataType.asInstanceOf[StructType]
+    var updatedIndex = 0
+    StructType(outputOrdinals.zipWithIndex.map { case (sourceOrdinal, index) =>
+      if (sourceOrdinal >= 0) {
+        sourceFields(sourceOrdinal)
+      } else {
+        val updatedValue = updatedValues(updatedIndex)
+        updatedIndex += 1
+        outputFields(index).copy(
+          dataType = updatedValue.dataType,
+          nullable = updatedValue.nullable)
+      }
+    })
+  }
+
+  override def eval(input: InternalRow): Any = {
+    val needsSource = nullable || outputOrdinals.exists(_ >= 0)
+    val sourceRow = if (needsSource) {
+      structExpr.eval(input).asInstanceOf[InternalRow]
+    } else {
+      null
+    }
+    if (sourceRow == null && nullable) {
+      null
+    } else {
+      val values = new Array[Any](outputOrdinals.length)
+      var updatedIndex = 0
+      var i = 0
+      while (i < outputOrdinals.length) {
+        val sourceOrdinal = outputOrdinals(i)
+        if (sourceOrdinal >= 0) {
+          values(i) = if (sourceRow == null) {
+            null
+          } else {
+            sourceRow.get(sourceOrdinal, dataType(i).dataType)
+          }
+        } else {
+          values(i) = updatedValues(updatedIndex).eval(input)
+          updatedIndex += 1
+        }
+        i += 1
+      }
+      new GenericInternalRow(values)
+    }
+  }
+
+  override protected def withNewChildrenInternal(
+      newChildren: IndexedSeq[Expression]): UpdateFieldsExpression =
+    copy(structExpr = newChildren.head, updatedValues = newChildren.tail)
+
+  override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+    val rowClass = classOf[GenericInternalRow].getName
+    val values = ctx.freshName("values")
+    val needsSource = nullable || outputOrdinals.exists(_ >= 0)
+    val sourceCode = if (needsSource) Some(structExpr.genCode(ctx)) else None
+    val sourceRow = ctx.freshName("sourceRow")
+
+    val copyValuesCode = if (outputOrdinals.exists(_ >= 0)) {
+      val index = ctx.freshName("i")
+      val sourceOrdinal = ctx.freshName("sourceOrdinal")
+      val sourceOrdinals = ctx.addReferenceObj("sourceOrdinals", outputOrdinals.toArray, "int[]")
+      val fieldTypes = ctx.addReferenceObj(
+        "fieldTypes",
+        dataType.fields.map(_.dataType).toArray,
+        s"${classOf[DataType].getName}[]")
+      s"""
+         |for (int $index = 0; $index < ${outputOrdinals.length}; $index++) {
+         |  int $sourceOrdinal = $sourceOrdinals[$index];
+         |  if ($sourceOrdinal >= 0) {
+         |    $values[$index] = $sourceRow == null ? null :
+         |      $sourceRow.get($sourceOrdinal, $fieldTypes[$index]);
+         |  }
+         |}
+       """.stripMargin
+    } else {
+      ""
+    }
+
+    val updatedOutputOrdinals = outputOrdinals.indices.filter(outputOrdinals(_) < 0)
+    assert(updatedOutputOrdinals.length == updatedValues.length)
+    val updatedCodes = updatedValues.zip(updatedOutputOrdinals).map { case (expr, ordinal) =>
+      val eval = expr.genCode(ctx)
+      s"""
+         |${eval.code}
+         |if (${eval.isNull}) {
+         |  $values[$ordinal] = null;
+         |} else {
+         |  $values[$ordinal] = ${eval.value};
+         |}
+       """.stripMargin
+    }
+    val updatedValuesCode = if (updatedCodes.nonEmpty) {
+      ctx.splitExpressionsWithCurrentInputs(
+        expressions = updatedCodes,
+        funcName = "updateFields",
+        extraArguments = "Object[]" -> values :: Nil)
+    } else {
+      ""
+    }
+
+    val outputCode = s"""
+       |Object[] $values = new Object[${outputOrdinals.length}];
+       |${copyValuesCode}
+       |${updatedValuesCode}
+       |${ev.value} = new $rowClass($values);
+       |$values = null;
+       """.stripMargin
+
+    sourceCode match {
+      case Some(source) if nullable =>
+        ev.copy(code = code"""
+          |${source.code}
+          |boolean ${ev.isNull} = ${source.isNull};
+          |InternalRow ${ev.value} = null;
+          |if (!${ev.isNull}) {
+          |  InternalRow $sourceRow = ${source.value};
+          |  $outputCode
+          |}
+        """.stripMargin)
+      case Some(source) =>
+        ev.copy(code = code"""
+          |${source.code}
+          |InternalRow $sourceRow = ${source.value};
+          |${CodeGenerator.javaType(dataType)} ${ev.value};
+          |$outputCode
+        """.stripMargin, isNull = FalseLiteral)
+      case None =>
+        ev.copy(code = code"""
+          |${CodeGenerator.javaType(dataType)} ${ev.value};
+          |$outputCode
+        """.stripMargin, isNull = FalseLiteral)
+    }
+  }
+}
+
 /**
  * Updates fields in a struct.
  */
 case class UpdateFields(structExpr: Expression, fieldOps: Seq[StructFieldsOperation])
   extends Unevaluable {
 
+  private lazy val generatedStructReadOwner = UpdateFields.generatedStructReadOwner(this)
+
   final override val nodePatterns: Seq[TreePattern] = Seq(UPDATE_FIELDS)
 
   override def checkInputDataTypes(): TypeCheckResult = {
@@ -783,7 +932,7 @@ case class UpdateFields(structExpr: Expression, fieldOps: Seq[StructFieldsOperat
           "inputSql" -> toSQLExpr(structExpr),
           "inputType" -> toSQLType(structExpr.dataType))
       )
-    } else if (newExprs.isEmpty) {
+    } else if (newFieldSources.isEmpty) {
       DataTypeMismatch(
         errorSubClass = "CANNOT_DROP_ALL_FIELDS",
         messageParameters = Map.empty
@@ -806,35 +955,107 @@ case class UpdateFields(structExpr: Expression, fieldOps: Seq[StructFieldsOperat
 
   override def prettyName: String = "update_fields"
 
-  private lazy val newFieldExprs: Seq[(StructField, Expression)] = {
-    def getFieldExpr(i: Int): Expression = structExpr match {
-      case c: CreateNamedStruct => c.valExprs(i)
-      case _ => GetStructField(structExpr, i)
-    }
-    val fieldsWithIndex = structExpr.dataType.asInstanceOf[StructType].fields.zipWithIndex
-    val existingFieldExprs: Seq[(StructField, Expression)] =
-      fieldsWithIndex.map { case (field, i) => (field, getFieldExpr(i)) }.toImmutableArraySeq
-    fieldOps.foldLeft(existingFieldExprs)((exprs, op) => op(exprs))
-  }
+  private lazy val newFieldSources: Seq[(StructField, UpdateFields.FieldSource)] =
+    UpdateFields.fieldSources(structExpr.dataType.asInstanceOf[StructType], fieldOps)
 
-  private lazy val newFields: Seq[StructField] = newFieldExprs.map(_._1)
+  private lazy val newFields: Seq[StructField] = newFieldSources.map(_._1)
 
-  lazy val newExprs: Seq[Expression] = newFieldExprs.map(_._2)
+  def newExpr(ordinal: Int): Expression = newFieldSources(ordinal)._2 match {
+    case UpdateFields.Original(sourceOrdinal) => getFieldExpr(sourceOrdinal)
+    case UpdateFields.Updated(expr) => expr
+  }
 
-  lazy val evalExpr: Expression = {
-    val createNamedStructExpr = CreateNamedStruct(newFieldExprs.flatMap {
-      case (field, expr) => Seq(Literal(field.name), expr)
-    })
+  private def getFieldExpr(ordinal: Int): Expression = structExpr match {
+    case c: CreateNamedStruct => c.valExprs(ordinal)
+    case _ => UpdateFields.markGeneratedStructRead(
+      GetStructField(structExpr, ordinal), generatedStructReadOwner)
+  }
 
-    if (structExpr.nullable) {
-      If(IsNull(structExpr), Literal(null, dataType), createNamedStructExpr)
-    } else {
-      createNamedStructExpr
+  lazy val evalExpr: Expression = With(structExpr) { case Seq(structRef) =>
+    // Nested generated reads must use the bound struct, while caller-supplied values stay separate.
+    def replaceStruct(expr: Expression): Expression = expr.transformDown {
+      case field @ GetStructField(child, _, _)
+          if UpdateFields.isGeneratedStructRead(field, generatedStructReadOwner) &&
+            child.eq(structExpr) =>
+        field.copy(child = structRef)
     }
+
+    val updatedValues = ArrayBuffer.empty[Expression]
+    val outputOrdinals = newFieldSources.map { case (_, source) =>
+      source match {
+        case UpdateFields.Original(ordinal) => structExpr match {
+          case namedStruct: CreateNamedStruct =>
+            updatedValues += namedStruct.valExprs(ordinal)
+            -1
+          case _ => ordinal
+        }
+        case UpdateFields.Updated(expr) =>
+          updatedValues += replaceStruct(expr)
+          -1
+      }
+    }
+    UpdateFieldsExpression(structRef, updatedValues.toSeq, outputOrdinals, newFields)
   }
 }
 
 object UpdateFields {
+  private sealed trait FieldSource
+  private case class Original(ordinal: Int) extends FieldSource
+  private case class Updated(expr: Expression) extends FieldSource
+
+  private def fieldSources(
+      structType: StructType,
+      fieldOps: Seq[StructFieldsOperation]): Seq[(StructField, FieldSource)] = {
+    val fields: Seq[(StructField, FieldSource)] = structType.fields.zipWithIndex.map {
+      case (field, ordinal) =>
+      field -> Original(ordinal)
+    }.toImmutableArraySeq
+    fieldOps.foldLeft(fields) {
+      case (current, op: WithField) =>
+        val updated = StructField(op.name, op.valExpr.dataType, op.valExpr.nullable) ->
+          Updated(op.valExpr)
+        val result = ArrayBuffer.empty[(StructField, FieldSource)]
+        var hasMatch = false
+        current.foreach { fieldSource =>
+          if (op.resolver(fieldSource._1.name, op.name)) {
+            hasMatch = true
+            result += updated
+          } else {
+            result += fieldSource
+          }
+        }
+        if (!hasMatch) result += updated
+        result.toSeq
+      case (current, op: DropField) =>
+        current.filterNot { case (field, _) => op.resolver(field.name, op.name) }
+      case (_, op) =>
+        throw SparkException.internalError(
+          s"Unsupported StructFieldsOperation: ${op.getClass.getName}")
+    }
+  }
+
+  private val GENERATED_STRUCT_READ = TreeNodeTag[ExprId]("generatedStructRead")
+  private val GENERATED_STRUCT_READ_OWNER = TreeNodeTag[ExprId]("generatedStructReadOwner")
+
+  private def generatedStructReadOwner(expr: Expression): ExprId = {
+    expr.getTagValue(GENERATED_STRUCT_READ_OWNER).getOrElse {
+      val owner = NamedExpression.newExprId
+      expr.setTagValue(GENERATED_STRUCT_READ_OWNER, owner)
+      owner
+    }
+  }
+
+  private def markGeneratedStructRead[T <: Expression](expr: T, owner: ExprId): T = {
+    expr.setTagValue(GENERATED_STRUCT_READ, owner)
+    expr
+  }
+
+  private def isGeneratedStructRead(expr: Expression, owner: ExprId): Boolean =
+    expr.getTagValue(GENERATED_STRUCT_READ).contains(owner)
+
+  private[catalyst] def isGeneratedStructRead(expr: Expression): Boolean =
+    expr.containsTag(GENERATED_STRUCT_READ)
+
   private def nameParts(fieldName: String): Seq[String] = {
     require(fieldName != null, "fieldName cannot be null")
 
@@ -878,7 +1099,9 @@ object UpdateFields {
         structExpr = newStruct,
         namePartsRemaining = namePartsRemaining.tail,
         valueFunc = valueFunc)
-      UpdateFields(structExpr, WithField(fieldName, newValue) :: Nil)
+      val updated = UpdateFields(structExpr, WithField(fieldName, newValue) :: Nil)
+      markGeneratedStructRead(newStruct, generatedStructReadOwner(updated))
+      updated
     }
   }
 }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ComplexTypes.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ComplexTypes.scala
index 34c7aa442d4ee..70a2205ca4346 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ComplexTypes.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ComplexTypes.scala
@@ -32,9 +32,12 @@ object SimplifyExtractValueOps extends Rule[LogicalPlan] {
       // Remove redundant field extraction.
       case GetStructField(createNamedStruct: CreateNamedStruct, ordinal, _) =>
         createNamedStruct.valExprs(ordinal)
-      case GetStructField(u: UpdateFields, ordinal, _)if !u.structExpr.isInstanceOf[UpdateFields] =>
+      // UpdateFields handles its nested-path reads after binding the input once.
+      case field @ GetStructField(u: UpdateFields, ordinal, _)
+          if !u.structExpr.isInstanceOf[UpdateFields] &&
+            !UpdateFields.isGeneratedStructRead(field) =>
         val structExpr = u.structExpr
-        u.newExprs(ordinal) match {
+        u.newExpr(ordinal) match {
           // if the struct itself is null, then any value extracted from it (expr) will be null
           // so we don't need to wrap expr in If(IsNull(struct), Literal(null, expr.dataType), expr)
           case expr: GetStructField if expr.child.semanticEquals(structExpr) => expr
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
index 7ba4493c0961e..97cbe03b00d2e 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
@@ -301,7 +301,9 @@ abstract class Optimizer(catalogManager: CatalogManager)
       RemoveNoopOperators),
     // This batch must be executed after the `RewriteSubquery` batch, which creates joins.
     Batch("NormalizeFloatingNumbers", Once, NormalizeFloatingNumbers),
-    Batch("ReplaceUpdateFieldsExpression", Once, ReplaceUpdateFieldsExpression)))
+    // Repeat because an update can contain another update. Keep the resulting With expressions so
+    // their definitions remain inside conditional branches until evaluation.
+    Batch("ReplaceUpdateFieldsExpression", fixedPoint, ReplaceUpdateFieldsExpression)))
 
     // remove any batches with no rules. this may happen when subclasses do not add optional rules.
     batches.filter(_.rules.nonEmpty)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/UpdateFields.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/UpdateFields.scala
index 9ca2f2a2e0ffa..1083040c396a0 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/UpdateFields.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/UpdateFields.scala
@@ -68,7 +68,9 @@ object OptimizeUpdateFields extends Rule[LogicalPlan] {
       val newWithFields = newNames.map(n => WithField(nameMap(n), newValues(n)))
       UpdateFields(structExpr, newWithFields.toSeq)
 
-    case UpdateFields(UpdateFields(struct, fieldOps1), fieldOps2) =>
+    // Flattening would separate a nested field read from the struct it reads.
+    case UpdateFields(UpdateFields(struct, fieldOps1), fieldOps2)
+        if !(fieldOps1 ++ fieldOps2).exists(_.exists(UpdateFields.isGeneratedStructRead)) =>
       UpdateFields(struct, fieldOps1 ++ fieldOps2)
   }
 
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceUpdateFieldsExpressionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceUpdateFieldsExpressionSuite.scala
new file mode 100644
index 0000000000000..50280a90ab9db
--- /dev/null
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceUpdateFieldsExpressionSuite.scala
@@ -0,0 +1,92 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.optimizer
+
+import org.apache.spark.sql.catalyst.dsl.expressions._
+import org.apache.spark.sql.catalyst.dsl.plans._
+import org.apache.spark.sql.catalyst.expressions.{Alias, AttributeReference, Expression, Literal,
+  UpdateFields, With, WithField}
+import org.apache.spark.sql.catalyst.plans.PlanTest
+import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan}
+import org.apache.spark.sql.catalyst.rules.RuleExecutor
+import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
+
+class ReplaceUpdateFieldsExpressionSuite extends PlanTest {
+  private object Optimize extends RuleExecutor[LogicalPlan] {
+    override protected def batches: Seq[Batch] =
+      Batch("ReplaceUpdateFieldsExpression", FixedPoint(10),
+        ReplaceUpdateFieldsExpression) :: Nil
+  }
+
+  private def nestedUpdates(struct: Expression, depth: Int): Expression = {
+    (0 until depth).foldLeft(struct) { (current, index) =>
+      UpdateFields(current, s"nested.field$index", Literal(index))
+    }
+  }
+
+  private def expressionNodeCount(nullable: Boolean, depth: Int): Int = {
+    val nested = if (nullable) {
+      $"nested".struct($"value".int)
+    } else {
+      $"nested".struct($"value".int).notNull
+    }
+    val struct = if (nullable) {
+      $"struct".struct(nested, $"a".int, $"b".int, $"c".int)
+    } else {
+      $"struct".struct(nested, $"a".int, $"b".int, $"c".int).notNull
+    }
+    val plan = LocalRelation(struct)
+      .select(Alias(nestedUpdates(struct, depth), "result")())
+    val optimized = Optimize.execute(plan)
+    assert(optimized.exists(_.expressions.exists(_.exists(_.isInstanceOf[With]))))
+    optimized.collect {
+      case node => node.expressions.map(_.collect { case expression => expression }.size).sum
+    }.sum
+  }
+
+  test("UpdateFields replacement grows linearly with nested update depth") {
+    Seq(true, false).foreach { nullable =>
+      val depthThree = expressionNodeCount(nullable, 3)
+      val depthSix = expressionNodeCount(nullable, 6)
+      assert(depthSix < depthThree * 3,
+        s"nullable=$nullable: depth 3 has $depthThree nodes, depth 6 has $depthSix")
+    }
+  }
+
+  test("UpdateFields replacement does not grow with struct width when reused") {
+    def nodeCount(width: Int, updateCount: Int): Int = {
+      val structType = StructType((0 until width).map { index =>
+        StructField(s"field$index", IntegerType)
+      })
+      val struct = AttributeReference("struct", structType)()
+      val updated = UpdateFields(struct, (0 until updateCount).map { index =>
+        WithField(s"field$index", Literal(index))
+      })
+      val optimized = Optimize.execute(LocalRelation(struct).select(
+        Alias(updated, "first")(), Alias(updated, "second")()))
+      optimized.expressions.map(_.collect { case expression => expression }.size).sum
+    }
+
+    Seq(1, 5).foreach { updateCount =>
+      val widthTen = nodeCount(10, updateCount)
+      val widthHundred = nodeCount(100, updateCount)
+      assert(widthHundred < widthTen * 2,
+        s"$updateCount updates: width 10 has $widthTen nodes, width 100 has $widthHundred")
+    }
+  }
+}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala
index c79fd0878ce25..9102f67ad2768 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala
@@ -28,7 +28,8 @@ import org.scalatest.matchers.should.Matchers._
 
 import org.apache.spark.{SparkException, SparkRuntimeException}
 import org.apache.spark.sql.UpdateFieldsBenchmark._
-import org.apache.spark.sql.catalyst.expressions.{InSet, Literal, NamedExpression, With}
+import org.apache.spark.sql.catalyst.expressions.{GetStructField, InSet, Literal, NamedExpression,
+  UpdateFields, With, WithField}
 import org.apache.spark.sql.catalyst.util.DateTimeTestUtils.{outstandingTimezonesIds, outstandingZoneIds}
 import org.apache.spark.sql.catalyst.util.DateTimeUtils
 import org.apache.spark.sql.catalyst.util.TimestampNanosTestUtils.foreachNanosPrecision
@@ -445,6 +446,141 @@ class ColumnExpressionSuite extends SharedSparkSession {
       SQLConf.CODEGEN_FACTORY_MODE.key -> "CODEGEN_ONLY")(f)
   }
 
+  test("withField should evaluate a nondeterministic struct expression once") {
+    onEachEvalPath {
+      val counter = new java.util.concurrent.atomic.AtomicLong()
+      val nextStruct = udf(() => {
+        val value = counter.getAndIncrement()
+        (value, value)
+      }).asNondeterministic()
+      val structExpr = nextStruct()
+      val df = spark.range(0, 10, 1, 1)
+      df.select(structExpr.withField("copy", structExpr)).collect().foreach { row =>
+        val result = row.getStruct(0)
+        val copy = result.getStruct(2)
+        assert(result.getLong(0) == result.getLong(1))
+        assert(copy.getLong(0) == copy.getLong(1))
+        assert(result.getLong(0) != copy.getLong(0))
+      }
+    }
+  }
+
+  test("withField should evaluate a reused nondeterministic value separately") {
+    onEachEvalPath {
+      val counter = new java.util.concurrent.atomic.AtomicLong()
+      val nextStruct = udf(() => {
+        val value = counter.getAndIncrement()
+        (value, value)
+      }).asNondeterministic()
+      val structExpr = nextStruct().expr
+      val updated = Column(UpdateFields(structExpr,
+        Seq(WithField("copy", GetStructField(structExpr, 0)))))
+      val df = spark.range(0, 10, 1, 1)
+      df.select(updated).collect().foreach { row =>
+        val result = row.getStruct(0)
+        assert(result.getLong(0) == result.getLong(1))
+        assert(result.getLong(0) != result.getLong(2))
+      }
+    }
+  }
+
+  test("withField should evaluate a reused nested update separately") {
+    onEachEvalPath {
+      val counter = new java.util.concurrent.atomic.AtomicLong()
+      val nextStruct = udf(() => {
+        val value = counter.getAndIncrement()
+        (value, value)
+      }).asNondeterministic()
+      val structExpr = nextStruct().expr
+      val nestedUpdate = UpdateFields(structExpr, Seq(WithField("extra", Literal(1))))
+      val updated = Column(UpdateFields(structExpr,
+        Seq(WithField("copy", GetStructField(nestedUpdate, 0)))))
+      val df = spark.range(0, 10, 1, 1)
+      df.select(updated).collect().foreach { row =>
+        val result = row.getStruct(0)
+        assert(result.getLong(0) == result.getLong(1))
+        assert(result.getLong(0) != result.getLong(2))
+      }
+    }
+  }
+
+  test("withField should reuse an updated struct for a nested path") {
+    onEachEvalPath {
+      val counter = new java.util.concurrent.atomic.AtomicLong()
+      val nextStruct = udf(() => {
+        val value = counter.getAndIncrement()
+        ((value, value), value)
+      }).asNondeterministic()
+      val structExpr = nextStruct().expr
+      val baseUpdate = UpdateFields(structExpr, Seq(WithField("extra", Literal(1))))
+      val updated = Column(UpdateFields(baseUpdate, "_1.copy", Literal(1)))
+      val df = spark.range(0, 10, 1, 1)
+      df.select(updated).collect().foreach { row =>
+        val result = row.getStruct(0)
+        val nested = result.getStruct(0)
+        assert(nested.getLong(0) == nested.getLong(1))
+        assert(nested.getLong(0) == result.getLong(1))
+      }
+    }
+  }
+
+  test("withField should preserve a nested path through a later update") {
+    onEachEvalPath {
+      val counter = new java.util.concurrent.atomic.AtomicLong()
+      val nextStruct = udf(() => {
+        val value = counter.getAndIncrement()
+        ((value, value), value)
+      }).asNondeterministic()
+      val structExpr = nextStruct().expr
+      val nestedUpdate = UpdateFields(structExpr, "_1.copy", Literal(1))
+      val updated = Column(UpdateFields(nestedUpdate, Seq(WithField("extra", Literal(1)))))
+      val df = spark.range(0, 10, 1, 1)
+      df.select(updated).collect().foreach { row =>
+        val result = row.getStruct(0)
+        val nested = result.getStruct(0)
+        assert(nested.getLong(0) == nested.getLong(1))
+        assert(nested.getLong(0) == result.getLong(1))
+      }
+    }
+  }
+
+  test("withField should not evaluate a value expression for null structs") {
+    onEachEvalPath {
+      withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
+        val structType = StructType(Seq(StructField("a", IntegerType, nullable = false)))
+        val df = spark.range(0, 3, 1, 1).select(
+          $"id",
+          when($"id" < 2, lit(null).cast(structType))
+            .otherwise(struct(lit(1)).cast(structType)).as("s"))
+        checkAnswer(
+          df.select($"s".withField("b", lit(1.0) / ($"id" - 1))),
+          Seq(Row(null), Row(null), Row(Row(1, 1.0))))
+      }
+    }
+  }
+
+  test("withField should remain inside a short-circuited branch") {
+    onEachEvalPath {
+      val throwingStruct = udf((id: Long) => {
+        if (id >= 0) throw new IllegalStateException(s"evaluated for $id") else ((id, id), id)
+      })
+      val updated = throwingStruct($"id").withField("_1.extra", lit(1))
+      checkAnswer(
+        Seq((false, 0L), (false, 1L)).toDF("keep", "id")
+          .filter($"keep" && updated.isNotNull).select($"id"),
+        Seq.empty)
+    }
+  }
+
+  test("dropFields should not evaluate omitted CreateNamedStruct fields") {
+    onEachEvalPath {
+      withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
+        val source = struct((lit(1) / $"id").as("dropped"), lit(1).as("kept"))
+        checkAnswer(spark.range(1).select(source.dropFields("dropped")), Row(Row(1)))
+      }
+    }
+  }
+
   test("SPARK-58902: BETWEEN on a nondeterministic input inside a conditional branch") {
     onEachEvalPath {
       // `BETWEEN` reads its input twice. Inlining the common expression into a branch gave each

From 3790650bdf01aaad796a5c4fa49abfd46372b807 Mon Sep 17 00:00:00 2001
From: Ben Hollis <ben.hollis@databricks.com>
Date: Wed, 7 Oct 2026 23:35:53 +0000
Subject: [PATCH 2/8] Document compact UpdateFields evaluation

---
 .../expressions/complexTypeCreator.scala      | 40 ++++++++++++++++++-
 1 file changed, 38 insertions(+), 2 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala
index 4e1bcad3b9a80..2180de023c77d 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala
@@ -764,10 +764,24 @@ case class DropField(name: String) extends StructFieldsOperation with LeafLike[E
     values.filterNot { case (field, _) => resolver(field.name, name) }
 }
 
+/**
+ * Builds a struct from fields copied from `structExpr` and values computed by `updatedValues`.
+ *
+ * `outputOrdinals` describes each output position. A nonnegative entry copies that field from
+ * `structExpr`; a negative entry evaluates the next expression in `updatedValues`. Storing copied
+ * fields as ordinals avoids adding a `GetStructField` child for every copied field. Row construction
+ * still visits every output field.
+ *
+ * For example, to transform `(a, b, c)` into `(a, newB, c, newD)`, `outputOrdinals` is
+ * `Seq(0, -1, 2, -1)` and `updatedValues` is `Seq(newB, newD)`. The first and third output fields
+ * copy source fields 0 and 2. The negative entries use `newB` and `newD`, in that order.
+ *
+ * If a nullable `structExpr` evaluates to null, the result is null and `updatedValues` are not
+ * evaluated.
+ */
 private case class UpdateFieldsExpression(
     structExpr: Expression,
     updatedValues: Seq[Expression],
-    // A negative ordinal marks an updated output, in updatedValues order.
     outputOrdinals: Seq[Int],
     outputFields: Seq[StructField]) extends Expression {
 
@@ -960,11 +974,16 @@ case class UpdateFields(structExpr: Expression, fieldOps: Seq[StructFieldsOperat
 
   private lazy val newFields: Seq[StructField] = newFieldSources.map(_._1)
 
+  /** Returns the expression for one output field without constructing the complete struct. */
   def newExpr(ordinal: Int): Expression = newFieldSources(ordinal)._2 match {
     case UpdateFields.Original(sourceOrdinal) => getFieldExpr(sourceOrdinal)
     case UpdateFields.Updated(expr) => expr
   }
 
+  /**
+   * Returns an input field expression and marks generated reads so they can share `structExpr`
+   * without changing equivalent expressions supplied to `withField`.
+   */
   private def getFieldExpr(ordinal: Int): Expression = structExpr match {
     case c: CreateNamedStruct => c.valExprs(ordinal)
     case _ => UpdateFields.markGeneratedStructRead(
@@ -972,7 +991,8 @@ case class UpdateFields(structExpr: Expression, fieldOps: Seq[StructFieldsOperat
   }
 
   lazy val evalExpr: Expression = With(structExpr) { case Seq(structRef) =>
-    // Nested generated reads must use the bound struct, while caller-supplied values stay separate.
+    // Read fields used to rebuild the input struct from structRef so structExpr is evaluated once.
+    // Do not rewrite expressions passed to withField; they must be evaluated independently.
     def replaceStruct(expr: Expression): Expression = expr.transformDown {
       case field @ GetStructField(child, _, _)
           if UpdateFields.isGeneratedStructRead(field, generatedStructReadOwner) &&
@@ -999,10 +1019,20 @@ case class UpdateFields(structExpr: Expression, fieldOps: Seq[StructFieldsOperat
 }
 
 object UpdateFields {
+  /** Describes where an output field gets its value without creating an expression for a copy. */
   private sealed trait FieldSource
+
+  /** Copies the field at `ordinal` from the input struct. */
   private case class Original(ordinal: Int) extends FieldSource
+
+  /** Computes the field from an expression supplied to `withField`. */
   private case class Updated(expr: Expression) extends FieldSource
 
+  /**
+   * Applies the requested field operations while retaining whether each output field is copied or
+   * computed. This provenance lets evaluation copy unchanged fields by ordinal and lets extraction
+   * return one field without constructing the complete struct.
+   */
   private def fieldSources(
       structType: StructType,
       fieldOps: Seq[StructFieldsOperation]): Seq[(StructField, FieldSource)] = {
@@ -1034,9 +1064,12 @@ object UpdateFields {
     }
   }
 
+  // Generated reads identify field accesses introduced while expanding nested paths. The owner
+  // prevents one UpdateFields from rewriting reads that belong to another UpdateFields.
   private val GENERATED_STRUCT_READ = TreeNodeTag[ExprId]("generatedStructRead")
   private val GENERATED_STRUCT_READ_OWNER = TreeNodeTag[ExprId]("generatedStructReadOwner")
 
+  /** Returns the stable identity used to associate generated reads with `expr`. */
   private def generatedStructReadOwner(expr: Expression): ExprId = {
     expr.getTagValue(GENERATED_STRUCT_READ_OWNER).getOrElse {
       val owner = NamedExpression.newExprId
@@ -1045,14 +1078,17 @@ object UpdateFields {
     }
   }
 
+  /** Marks a field read as generated by the UpdateFields identified by `owner`. */
   private def markGeneratedStructRead[T <: Expression](expr: T, owner: ExprId): T = {
     expr.setTagValue(GENERATED_STRUCT_READ, owner)
     expr
   }
 
+  /** Tests whether a field read was generated by the UpdateFields identified by `owner`. */
   private def isGeneratedStructRead(expr: Expression, owner: ExprId): Boolean =
     expr.getTagValue(GENERATED_STRUCT_READ).contains(owner)
 
+  /** Tests whether a field read was generated while expanding any nested update path. */
   private[catalyst] def isGeneratedStructRead(expr: Expression): Boolean =
     expr.containsTag(GENERATED_STRUCT_READ)
 

From ca280d68a095fc94649db954f2b1699e65646a05 Mon Sep 17 00:00:00 2001
From: Ben Hollis <ben.hollis@databricks.com>
Date: Sat, 10 Oct 2026 22:37:13 +0000
Subject: [PATCH 3/8] [SPARK-60067][SQL] Preserve nested UpdateFields semantics

Keep nested updates relative to their parent field to avoid duplicated inputs and tree growth while preserving overwrite pruning and chain flattening. Use typed generated copies and indexed interpreted evaluation.

Validation: 265 tests passed across seven SQL and Catalyst suites using the cached SBT classpath. SBT loading is blocked by Jackson BOM resolution.
---
 .../expressions/complexTypeCreator.scala      | 205 ++++++++----------
 .../sql/catalyst/optimizer/ComplexTypes.scala |   5 +-
 .../sql/catalyst/optimizer/UpdateFields.scala |  71 +++---
 .../spark/sql/ColumnExpressionSuite.scala     |  97 +++++++++
 4 files changed, 233 insertions(+), 145 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala
index 2180de023c77d..bd292b5af1412 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala
@@ -21,14 +21,14 @@ import scala.collection.mutable.ArrayBuffer
 
 import org.apache.spark.SparkException
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.analysis.{Resolver, TypeCheckResult, TypeCoercion, UnresolvedAttribute, UnresolvedExtractValue}
+import org.apache.spark.sql.catalyst.analysis.{Resolver, TypeCheckResult, TypeCoercion, TypeCoercionValidation, UnresolvedAttribute, UnresolvedExtractValue}
 import org.apache.spark.sql.catalyst.analysis.FunctionRegistry.{FUNC_ALIAS, FunctionBuilder}
 import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.DataTypeMismatch
 import org.apache.spark.sql.catalyst.expressions.Cast._
 import org.apache.spark.sql.catalyst.expressions.codegen._
 import org.apache.spark.sql.catalyst.expressions.codegen.Block._
 import org.apache.spark.sql.catalyst.parser.CatalystSqlParser
-import org.apache.spark.sql.catalyst.trees.{LeafLike, TreeNodeTag, UnaryLike}
+import org.apache.spark.sql.catalyst.trees.{CurrentOrigin, LeafLike, UnaryLike}
 import org.apache.spark.sql.catalyst.trees.TreePattern._
 import org.apache.spark.sql.catalyst.util._
 import org.apache.spark.sql.errors.QueryCompilationErrors
@@ -715,12 +715,6 @@ trait StructFieldsOperation extends Expression with Unevaluable {
   override def nullable: Boolean = throw SparkException.internalError(
     "StructFieldsOperation.nullable should not be called.")
 
-  /**
-   * Returns an updated list of StructFields and Expressions that will ultimately be used
-   * as the fields argument for [[StructType]] and as the children argument for
-   * [[CreateNamedStruct]] respectively inside of [[UpdateFields]].
-   */
-  def apply(values: Seq[(StructField, Expression)]): Seq[(StructField, Expression)]
 }
 
 /**
@@ -732,22 +726,6 @@ trait StructFieldsOperation extends Expression with Unevaluable {
 case class WithField(name: String, valExpr: Expression)
   extends StructFieldsOperation with UnaryLike[Expression] {
 
-  override def apply(values: Seq[(StructField, Expression)]): Seq[(StructField, Expression)] = {
-    val newFieldExpr = (StructField(name, valExpr.dataType, valExpr.nullable), valExpr)
-    val result = ArrayBuffer.empty[(StructField, Expression)]
-    var hasMatch = false
-    for (existingFieldExpr @ (existingField, _) <- values) {
-      if (resolver(existingField.name, name)) {
-        hasMatch = true
-        result += newFieldExpr
-      } else {
-        result += existingFieldExpr
-      }
-    }
-    if (!hasMatch) result += newFieldExpr
-    result.toSeq
-  }
-
   override def child: Expression = valExpr
 
   override def prettyName: String = "WithField"
@@ -759,9 +737,19 @@ case class WithField(name: String, valExpr: Expression)
 /**
  * Drop a field by name.
  */
-case class DropField(name: String) extends StructFieldsOperation with LeafLike[Expression] {
-  override def apply(values: Seq[(StructField, Expression)]): Seq[(StructField, Expression)] =
-    values.filterNot { case (field, _) => resolver(field.name, name) }
+case class DropField(name: String) extends StructFieldsOperation with LeafLike[Expression]
+
+/**
+ * Updates the current value of a field without duplicating the parent struct expression.
+ */
+case class UpdateNestedField(name: String, fieldOps: Seq[StructFieldsOperation])
+  extends StructFieldsOperation {
+
+  override def children: Seq[Expression] = fieldOps
+
+  override protected def withNewChildrenInternal(
+      newChildren: IndexedSeq[Expression]): Expression =
+    copy(fieldOps = newChildren.map(_.asInstanceOf[StructFieldsOperation]))
 }
 
 /**
@@ -789,6 +777,10 @@ private case class UpdateFieldsExpression(
 
   override def nullable: Boolean = structExpr.nullable
 
+  private lazy val indexedOrdinals = outputOrdinals.toArray
+  private lazy val indexedUpdatedValues = updatedValues.toArray
+  private lazy val needsSource = nullable || indexedOrdinals.exists(_ >= 0)
+
   override lazy val dataType: StructType = {
     val sourceFields = structExpr.dataType.asInstanceOf[StructType]
     var updatedIndex = 0
@@ -806,7 +798,6 @@ private case class UpdateFieldsExpression(
   }
 
   override def eval(input: InternalRow): Any = {
-    val needsSource = nullable || outputOrdinals.exists(_ >= 0)
     val sourceRow = if (needsSource) {
       structExpr.eval(input).asInstanceOf[InternalRow]
     } else {
@@ -815,11 +806,11 @@ private case class UpdateFieldsExpression(
     if (sourceRow == null && nullable) {
       null
     } else {
-      val values = new Array[Any](outputOrdinals.length)
+      val values = new Array[Any](indexedOrdinals.length)
       var updatedIndex = 0
       var i = 0
-      while (i < outputOrdinals.length) {
-        val sourceOrdinal = outputOrdinals(i)
+      while (i < indexedOrdinals.length) {
+        val sourceOrdinal = indexedOrdinals(i)
         if (sourceOrdinal >= 0) {
           values(i) = if (sourceRow == null) {
             null
@@ -827,7 +818,7 @@ private case class UpdateFieldsExpression(
             sourceRow.get(sourceOrdinal, dataType(i).dataType)
           }
         } else {
-          values(i) = updatedValues(updatedIndex).eval(input)
+          values(i) = indexedUpdatedValues(updatedIndex).eval(input)
           updatedIndex += 1
         }
         i += 1
@@ -843,27 +834,30 @@ private case class UpdateFieldsExpression(
   override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
     val rowClass = classOf[GenericInternalRow].getName
     val values = ctx.freshName("values")
-    val needsSource = nullable || outputOrdinals.exists(_ >= 0)
     val sourceCode = if (needsSource) Some(structExpr.genCode(ctx)) else None
     val sourceRow = ctx.freshName("sourceRow")
 
     val copyValuesCode = if (outputOrdinals.exists(_ >= 0)) {
-      val index = ctx.freshName("i")
-      val sourceOrdinal = ctx.freshName("sourceOrdinal")
-      val sourceOrdinals = ctx.addReferenceObj("sourceOrdinals", outputOrdinals.toArray, "int[]")
-      val fieldTypes = ctx.addReferenceObj(
-        "fieldTypes",
-        dataType.fields.map(_.dataType).toArray,
-        s"${classOf[DataType].getName}[]")
-      s"""
-         |for (int $index = 0; $index < ${outputOrdinals.length}; $index++) {
-         |  int $sourceOrdinal = $sourceOrdinals[$index];
-         |  if ($sourceOrdinal >= 0) {
-         |    $values[$index] = $sourceRow == null ? null :
-         |      $sourceRow.get($sourceOrdinal, $fieldTypes[$index]);
-         |  }
-         |}
-       """.stripMargin
+      val copiedFields = outputOrdinals.zipWithIndex.filter(_._1 >= 0)
+        .groupBy { case (_, ordinal) => dataType(ordinal).dataType }
+      val copyCodes = copiedFields.toSeq.map { case (fieldType, fields) =>
+        val index = ctx.freshName("i")
+        val sourceOrdinal = ctx.freshName("sourceOrdinal")
+        val ordinals = ctx.addReferenceObj("copyOrdinals",
+          fields.flatMap { case (source, output) => Seq(output, source) }.toArray, "int[]")
+        val value = CodeGenerator.getValue(sourceRow, fieldType, sourceOrdinal)
+        s"""
+           |for (int $index = 0; $index < ${fields.length * 2}; $index += 2) {
+           |  int $sourceOrdinal = $ordinals[$index + 1];
+           |  $values[$ordinals[$index]] =
+           |    $sourceRow == null || $sourceRow.isNullAt($sourceOrdinal) ? null : $value;
+           |}
+         """.stripMargin
+      }
+      ctx.splitExpressionsWithCurrentInputs(
+        expressions = copyCodes,
+        funcName = "copyFields",
+        extraArguments = Seq("InternalRow" -> sourceRow, "Object[]" -> values))
     } else {
       ""
     }
@@ -931,8 +925,6 @@ private case class UpdateFieldsExpression(
 case class UpdateFields(structExpr: Expression, fieldOps: Seq[StructFieldsOperation])
   extends Unevaluable {
 
-  private lazy val generatedStructReadOwner = UpdateFields.generatedStructReadOwner(this)
-
   final override val nodePatterns: Seq[TreePattern] = Seq(UPDATE_FIELDS)
 
   override def checkInputDataTypes(): TypeCheckResult = {
@@ -970,7 +962,7 @@ case class UpdateFields(structExpr: Expression, fieldOps: Seq[StructFieldsOperat
   override def prettyName: String = "update_fields"
 
   private lazy val newFieldSources: Seq[(StructField, UpdateFields.FieldSource)] =
-    UpdateFields.fieldSources(structExpr.dataType.asInstanceOf[StructType], fieldOps)
+    UpdateFields.fieldSources(structExpr, fieldOps)
 
   private lazy val newFields: Seq[StructField] = newFieldSources.map(_._1)
 
@@ -980,28 +972,19 @@ case class UpdateFields(structExpr: Expression, fieldOps: Seq[StructFieldsOperat
     case UpdateFields.Updated(expr) => expr
   }
 
-  /**
-   * Returns an input field expression and marks generated reads so they can share `structExpr`
-   * without changing equivalent expressions supplied to `withField`.
-   */
   private def getFieldExpr(ordinal: Int): Expression = structExpr match {
     case c: CreateNamedStruct => c.valExprs(ordinal)
-    case _ => UpdateFields.markGeneratedStructRead(
-      GetStructField(structExpr, ordinal), generatedStructReadOwner)
+    case _ => GetStructField(structExpr, ordinal)
   }
 
   lazy val evalExpr: Expression = With(structExpr) { case Seq(structRef) =>
-    // Read fields used to rebuild the input struct from structRef so structExpr is evaluated once.
-    // Do not rewrite expressions passed to withField; they must be evaluated independently.
-    def replaceStruct(expr: Expression): Expression = expr.transformDown {
-      case field @ GetStructField(child, _, _)
-          if UpdateFields.isGeneratedStructRead(field, generatedStructReadOwner) &&
-            child.eq(structExpr) =>
-        field.copy(child = structRef)
+    val source = structExpr match {
+      case namedStruct: CreateNamedStruct => namedStruct
+      case _ => structRef
     }
-
+    val fieldSources = UpdateFields.fieldSources(source, fieldOps)
     val updatedValues = ArrayBuffer.empty[Expression]
-    val outputOrdinals = newFieldSources.map { case (_, source) =>
+    val outputOrdinals = fieldSources.map { case (_, source) =>
       source match {
         case UpdateFields.Original(ordinal) => structExpr match {
           case namedStruct: CreateNamedStruct =>
@@ -1010,7 +993,7 @@ case class UpdateFields(structExpr: Expression, fieldOps: Seq[StructFieldsOperat
           case _ => ordinal
         }
         case UpdateFields.Updated(expr) =>
-          updatedValues += replaceStruct(expr)
+          updatedValues += expr
           -1
       }
     }
@@ -1034,14 +1017,24 @@ object UpdateFields {
    * return one field without constructing the complete struct.
    */
   private def fieldSources(
-      structType: StructType,
+      structExpr: Expression,
       fieldOps: Seq[StructFieldsOperation]): Seq[(StructField, FieldSource)] = {
+    val structType = structExpr.dataType.asInstanceOf[StructType]
+    def expression(source: FieldSource): Expression = source match {
+      case Original(ordinal) => structExpr match {
+        case namedStruct: CreateNamedStruct => namedStruct.valExprs(ordinal)
+        case _ => GetStructField(structExpr, ordinal)
+      }
+      case Updated(expr) => expr
+    }
     val fields: Seq[(StructField, FieldSource)] = structType.fields.zipWithIndex.map {
       case (field, ordinal) =>
       field -> Original(ordinal)
     }.toImmutableArraySeq
-    fieldOps.foldLeft(fields) {
-      case (current, op: WithField) =>
+    def applyOperation(
+        current: Seq[(StructField, FieldSource)],
+        operation: StructFieldsOperation): Seq[(StructField, FieldSource)] = operation match {
+      case op: WithField =>
         val updated = StructField(op.name, op.valExpr.dataType, op.valExpr.nullable) ->
           Updated(op.valExpr)
         val result = ArrayBuffer.empty[(StructField, FieldSource)]
@@ -1056,42 +1049,36 @@ object UpdateFields {
         }
         if (!hasMatch) result += updated
         result.toSeq
-      case (current, op: DropField) =>
+      case op: DropField =>
         current.filterNot { case (field, _) => op.resolver(field.name, op.name) }
-      case (_, op) =>
+      case op =>
         throw SparkException.internalError(
           s"Unsupported StructFieldsOperation: ${op.getClass.getName}")
     }
-  }
-
-  // Generated reads identify field accesses introduced while expanding nested paths. The owner
-  // prevents one UpdateFields from rewriting reads that belong to another UpdateFields.
-  private val GENERATED_STRUCT_READ = TreeNodeTag[ExprId]("generatedStructRead")
-  private val GENERATED_STRUCT_READ_OWNER = TreeNodeTag[ExprId]("generatedStructReadOwner")
-
-  /** Returns the stable identity used to associate generated reads with `expr`. */
-  private def generatedStructReadOwner(expr: Expression): ExprId = {
-    expr.getTagValue(GENERATED_STRUCT_READ_OWNER).getOrElse {
-      val owner = NamedExpression.newExprId
-      expr.setTagValue(GENERATED_STRUCT_READ_OWNER, owner)
-      owner
+    fieldOps.foldLeft(fields) { (current, operation) =>
+      val expanded = operation match {
+        case nested: UpdateNestedField =>
+          val currentType = StructType(current.map(_._1))
+          val field = ExtractValue(
+            Literal.create(null, currentType), Literal(nested.name), nested.resolver)
+            .asInstanceOf[GetStructField]
+          val updated = CurrentOrigin.withOrigin(nested.origin) {
+            expression(current(field.ordinal)._2) match {
+              case UpdateFields(input, previousOps) =>
+                UpdateFields(input, previousOps ++ nested.fieldOps)
+              case input => UpdateFields(input, nested.fieldOps)
+            }
+          }
+          if (updated.checkInputDataTypes().isFailure) {
+            TypeCoercionValidation.failOnTypeCheckResult(updated)
+          }
+          WithField(nested.name, updated)
+        case other => other
+      }
+      applyOperation(current, expanded)
     }
   }
 
-  /** Marks a field read as generated by the UpdateFields identified by `owner`. */
-  private def markGeneratedStructRead[T <: Expression](expr: T, owner: ExprId): T = {
-    expr.setTagValue(GENERATED_STRUCT_READ, owner)
-    expr
-  }
-
-  /** Tests whether a field read was generated by the UpdateFields identified by `owner`. */
-  private def isGeneratedStructRead(expr: Expression, owner: ExprId): Boolean =
-    expr.getTagValue(GENERATED_STRUCT_READ).contains(owner)
-
-  /** Tests whether a field read was generated while expanding any nested update path. */
-  private[catalyst] def isGeneratedStructRead(expr: Expression): Boolean =
-    expr.containsTag(GENERATED_STRUCT_READ)
-
   private def nameParts(fieldName: String): Seq[String] = {
     require(fieldName != null, "fieldName cannot be null")
 
@@ -1124,20 +1111,12 @@ object UpdateFields {
     if (namePartsRemaining.length == 1) {
       UpdateFields(structExpr, valueFunc(fieldName) :: Nil)
     } else {
-      val newStruct = if (structExpr.resolved) {
-        val resolver = SQLConf.get.resolver
-        ExtractValue(structExpr, Literal(fieldName), resolver)
-      } else {
-        UnresolvedExtractValue(structExpr, Literal(fieldName))
+      def nestedOperation(parts: Seq[String]): StructFieldsOperation = {
+        if (parts.length == 1) valueFunc(parts.head)
+        else UpdateNestedField(parts.head, Seq(nestedOperation(parts.tail)))
       }
-
-      val newValue = updateFieldsHelper(
-        structExpr = newStruct,
-        namePartsRemaining = namePartsRemaining.tail,
-        valueFunc = valueFunc)
-      val updated = UpdateFields(structExpr, WithField(fieldName, newValue) :: Nil)
-      markGeneratedStructRead(newStruct, generatedStructReadOwner(updated))
-      updated
+      UpdateFields(structExpr,
+        Seq(UpdateNestedField(fieldName, Seq(nestedOperation(namePartsRemaining.tail)))))
     }
   }
 }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ComplexTypes.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ComplexTypes.scala
index 70a2205ca4346..1a81584bc986e 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ComplexTypes.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ComplexTypes.scala
@@ -32,10 +32,7 @@ object SimplifyExtractValueOps extends Rule[LogicalPlan] {
       // Remove redundant field extraction.
       case GetStructField(createNamedStruct: CreateNamedStruct, ordinal, _) =>
         createNamedStruct.valExprs(ordinal)
-      // UpdateFields handles its nested-path reads after binding the input once.
-      case field @ GetStructField(u: UpdateFields, ordinal, _)
-          if !u.structExpr.isInstanceOf[UpdateFields] &&
-            !UpdateFields.isGeneratedStructRead(field) =>
+      case GetStructField(u: UpdateFields, ordinal, _) =>
         val structExpr = u.structExpr
         u.newExpr(ordinal) match {
           // if the struct itself is null, then any value extracted from it (expr) will be null
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/UpdateFields.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/UpdateFields.scala
index 1083040c396a0..f7ece835c7581 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/UpdateFields.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/UpdateFields.scala
@@ -21,7 +21,7 @@ import java.util.Locale
 
 import scala.collection.mutable
 
-import org.apache.spark.sql.catalyst.expressions.{Expression, UpdateFields, WithField}
+import org.apache.spark.sql.catalyst.expressions.{Expression, StructFieldsOperation, UpdateFields, UpdateNestedField, WithField}
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.trees.TreePattern.UPDATE_FIELDS
@@ -39,39 +39,54 @@ object OptimizeUpdateFields extends Rule[LogicalPlan] {
     }
   }
 
+  private def optimizeFieldOps(
+      fieldOps: Seq[StructFieldsOperation]): Seq[StructFieldsOperation] = {
+    val operations = fieldOps.map {
+      case nested: UpdateNestedField =>
+        nested.copy(fieldOps = optimizeFieldOps(nested.fieldOps))
+      case other => other
+    }
+    if (!operations.forall {
+        case _: WithField | _: UpdateNestedField => true
+        case _ => false
+      }) {
+      operations
+    } else {
+      val combined = mutable.LinkedHashMap.empty[String, StructFieldsOperation]
+      operations.foreach { operation =>
+        val name = operation match {
+          case WithField(name, _) => name
+          case UpdateNestedField(name, _) => name
+        }
+        val key = if (conf.caseSensitiveAnalysis) name else name.toLowerCase(Locale.ROOT)
+        val updated = operation match {
+          case nested: UpdateNestedField => combined.get(key) match {
+            case Some(previous: UpdateNestedField) =>
+              nested.copy(fieldOps = optimizeFieldOps(previous.fieldOps ++ nested.fieldOps))
+            case Some(WithField(_, value)) =>
+              WithField(name, UpdateFields(value, nested.fieldOps))
+            case _ => nested
+          }
+          case other => other
+        }
+        combined.update(key, updated)
+      }
+      combined.values.toSeq
+    }
+  }
+
   val optimizeUpdateFields: PartialFunction[Expression, Expression] = {
     case UpdateFields(structExpr, fieldOps)
       if fieldOps.forall(_.isInstanceOf[WithField]) &&
         canOptimize(fieldOps.map(_.asInstanceOf[WithField].name)) =>
-      val caseSensitive = conf.caseSensitiveAnalysis
+      UpdateFields(structExpr, optimizeFieldOps(fieldOps))
 
-      val withFields = fieldOps.map(_.asInstanceOf[WithField])
-      val names = withFields.map(_.name)
-      val values = withFields.map(_.valExpr)
-
-      val newNames = mutable.ArrayBuffer.empty[String]
-      val newValues = mutable.HashMap.empty[String, Expression]
-      // Used to remember the casing of the last instance
-      val nameMap = mutable.HashMap.empty[String, String]
-
-      names.zip(values).foreach { case (name, value) =>
-        val normalizedName = if (caseSensitive) name else name.toLowerCase(Locale.ROOT)
-        if (nameMap.contains(normalizedName)) {
-          newValues += normalizedName -> value
-        } else {
-          newNames += normalizedName
-          newValues += normalizedName -> value
-        }
-        nameMap += normalizedName -> name
-      }
-
-      val newWithFields = newNames.map(n => WithField(nameMap(n), newValues(n)))
-      UpdateFields(structExpr, newWithFields.toSeq)
-
-    // Flattening would separate a nested field read from the struct it reads.
-    case UpdateFields(UpdateFields(struct, fieldOps1), fieldOps2)
-        if !(fieldOps1 ++ fieldOps2).exists(_.exists(UpdateFields.isGeneratedStructRead)) =>
+    case UpdateFields(UpdateFields(struct, fieldOps1), fieldOps2) =>
       UpdateFields(struct, fieldOps1 ++ fieldOps2)
+
+    case UpdateFields(struct, fieldOps)
+        if fieldOps.exists(_.isInstanceOf[UpdateNestedField]) =>
+      UpdateFields(struct, optimizeFieldOps(fieldOps))
   }
 
   def apply(plan: LogicalPlan): LogicalPlan = plan.resolveExpressionsWithPruning(
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala
index 9102f67ad2768..932f005f3c68f 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala
@@ -504,6 +504,103 @@ class ColumnExpressionSuite extends SharedSparkSession {
     }
   }
 
+  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 =>
+      val updated = (0 until count).foldLeft(col("s")) { (column, index) =>
+        column.withField(s"_1.x$index", lit(index))
+      }
+      val query = df.select(updated).queryExecution
+      Seq(query.analyzed, query.optimizedPlan).foreach { plan =>
+        val nodes = plan.flatMap(_.expressions).map(_.collect { case _ => 1 }.size).sum
+        assert(nodes < 30 * count, s"$count updates produced $nodes expression nodes")
+      }
+      checkAnswer(df.select(updated.dropFields((0 until count).map(i => s"_1.x$i"): _*)),
+        Row(Row(Row(1, 2), 3)))
+    }
+  }
+
+  test("nested field updates should prune dropped and overwritten values") {
+    onEachEvalPath {
+      withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
+        val df = Seq((((1, 2), 3), 0)).toDF("s", "z")
+        val updated = col("s").withField("_1.d", lit(1) / col("z"))
+        checkAnswer(df.select(updated.dropFields("_1")), Row(Row(3)))
+        checkAnswer(df.select(updated.withField("_1.d", lit(4))),
+          Row(Row(Row(1, 2, 4), 3)))
+        checkAnswer(df.select(updated.getField("_2")), Row(3))
+        val projected = df.select(
+          col("s").withField("x", lit(1) / col("z")).as("t"))
+        checkAnswer(projected.select(col("t").withField("y", lit(2)).getField("_2")), Row(3))
+      }
+    }
+  }
+
+  test("nested field updates should share an unresolved nondeterministic input") {
+    onEachEvalPath {
+      val counter = new java.util.concurrent.atomic.AtomicLong()
+      val nextStruct = udf((id: Long) => {
+        val value = counter.getAndIncrement() + id
+        ((value, value), value)
+      }).asNondeterministic()
+      val updated = nextStruct(col("id")).withField("_1.copy", lit(1))
+        .withField("_1.extra", lit(2)).withField("extra", lit(3))
+      spark.range(10).select(updated).collect().foreach { row =>
+        val result = row.getStruct(0)
+        val nested = result.getStruct(0)
+        assert(nested.getLong(0) == nested.getLong(1))
+        assert(nested.getLong(0) == result.getLong(1))
+      }
+    }
+  }
+
+  test("nested field updates should discard paths overwritten before resolution") {
+    val df = Seq((1, 2)).toDF("a", "b").select(struct(col("a"), col("b")).as("s"))
+    checkAnswer(df.select(col("s").withField("a.x", lit(3)).withField("a", lit(4))),
+      Row(Row(4, 2)))
+    checkAnswer(df.select(col("s").withField("missing.x", lit(3))
+      .withField("missing", lit(4))), Row(Row(1, 2, 4)))
+  }
+
+  test("withField should copy TIME fields from vectorized Parquet structs") {
+    withSQLConf(
+      SQLConf.TIME_TYPE_ENABLED.key -> "true",
+      SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "true",
+      SQLConf.PARQUET_VECTORIZED_READER_NESTED_COLUMN_ENABLED.key -> "true",
+      SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true") {
+      withTempPath { path =>
+        spark.sql("SELECT named_struct('t', TIME '12:34:56', 'a', 1) AS s")
+          .write.parquet(path.getCanonicalPath)
+        val df = spark.read.parquet(path.getCanonicalPath)
+        checkAnswer(df.select(col("s").withField("b", lit(2))),
+          spark.sql("SELECT named_struct('t', TIME '12:34:56', 'a', 1, 'b', 2)"))
+      }
+    }
+  }
+
+  test("withField should copy geospatial fields from vectorized Parquet structs") {
+    val point = "010100000000000000000031400000000000001C40"
+      .grouped(2).map(Integer.parseInt(_, 16).toByte).toArray
+    Seq(
+      GeometryType(4326) -> Geometry.fromWKB(point, 4326),
+      GeographyType(4326) -> Geography.fromWKB(point, 4326)).foreach { case (dataType, value) =>
+      withSQLConf(
+        SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "true",
+        SQLConf.PARQUET_VECTORIZED_READER_NESTED_COLUMN_ENABLED.key -> "true",
+        SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true") {
+        withTempPath { path =>
+          val schema = StructType(Seq(StructField("s", StructType(Seq(
+            StructField("g", dataType), StructField("a", IntegerType))))))
+          spark.createDataFrame(
+            sparkContext.parallelize(Seq(Row(Row(value, 1)))), schema)
+            .write.parquet(path.getCanonicalPath)
+          val df = spark.read.parquet(path.getCanonicalPath)
+          checkAnswer(df.select(col("s").withField("b", lit(2))), Row(Row(value, 1, 2)))
+        }
+      }
+    }
+  }
+
   test("withField should reuse an updated struct for a nested path") {
     onEachEvalPath {
       val counter = new java.util.concurrent.atomic.AtomicLong()

From ce59df630dadca20f896bc60b08546ff3671065a Mon Sep 17 00:00:00 2001
From: Ben Hollis <ben.hollis@databricks.com>
Date: Sun, 11 Oct 2026 01:49:01 +0000
Subject: [PATCH 4/8] [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.
---
 .../expressions/EquivalentExpressions.scala   | 32 +++++++---
 .../spark/sql/catalyst/expressions/With.scala |  2 +-
 .../expressions/complexTypeCreator.scala      | 13 +++-
 .../SubexpressionEliminationSuite.scala       | 59 ++++++++++++++++-
 .../spark/sql/ColumnExpressionSuite.scala     | 64 +++++++++++++++++++
 5 files changed, 157 insertions(+), 13 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala
index 97844b82349d0..f8822c72727bc 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala
+++ b/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.
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/With.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/With.scala
index 85388c271885d..1fd9f441b1536 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/With.scala
+++ b/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)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala
index bd292b5af1412..551facab098ff 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala
+++ b/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
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SubexpressionEliminationSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SubexpressionEliminationSuite.scala
index c8adeb82c9f53..eb0dddc14a086 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SubexpressionEliminationSuite.scala
+++ b/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)))
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala
index 932f005f3c68f..48003d551ef36 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala
+++ b/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 =>

From edc65d8d50fb24ad2476b3f372c1b06000fc0c8e Mon Sep 17 00:00:00 2001
From: Ben Hollis <ben.hollis@databricks.com>
Date: Sun, 11 Oct 2026 03:30:04 +0000
Subject: [PATCH 5/8] Avoid unnecessary scopes when lowering flat struct
 updates

Introduce With only for nested updates to avoid planning overhead for flat operations. Add a regression test for replacements, drops, and nullability.

Validation: 364 tests passed across 11 suites; planning and projection benchmarks rerun.
---
 .../expressions/complexTypeCreator.scala      | 41 +++++++++++--------
 .../ReplaceUpdateFieldsExpressionSuite.scala  | 20 ++++++++-
 2 files changed, 43 insertions(+), 18 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala
index 551facab098ff..a91c937d00bea 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala
@@ -988,27 +988,34 @@ case class UpdateFields(structExpr: Expression, fieldOps: Seq[StructFieldsOperat
     case _ => GetStructField(structExpr, ordinal)
   }
 
-  lazy val evalExpr: Expression = With(structExpr) { case Seq(structRef) =>
-    val source = structExpr match {
-      case namedStruct: CreateNamedStruct => namedStruct
-      case _ => structRef
-    }
-    val fieldSources = UpdateFields.fieldSources(source, fieldOps)
-    val updatedValues = ArrayBuffer.empty[Expression]
-    val outputOrdinals = fieldSources.map { case (_, source) =>
-      source match {
-        case UpdateFields.Original(ordinal) => structExpr match {
-          case namedStruct: CreateNamedStruct =>
-            updatedValues += namedStruct.valExprs(ordinal)
+  lazy val evalExpr: Expression = {
+    def buildExpression(structRef: Expression): Expression = {
+      val source = structExpr match {
+        case namedStruct: CreateNamedStruct => namedStruct
+        case _ => structRef
+      }
+      val fieldSources = UpdateFields.fieldSources(source, fieldOps)
+      val updatedValues = ArrayBuffer.empty[Expression]
+      val outputOrdinals = fieldSources.map { case (_, source) =>
+        source match {
+          case UpdateFields.Original(ordinal) => structExpr match {
+            case namedStruct: CreateNamedStruct =>
+              updatedValues += namedStruct.valExprs(ordinal)
+              -1
+            case _ => ordinal
+          }
+          case UpdateFields.Updated(expr) =>
+            updatedValues += expr
             -1
-          case _ => ordinal
         }
-        case UpdateFields.Updated(expr) =>
-          updatedValues += expr
-          -1
       }
+      UpdateFieldsExpression(structRef, updatedValues.toSeq, outputOrdinals, newFields)
+    }
+    if (fieldOps.exists(_.isInstanceOf[UpdateNestedField])) {
+      With(structExpr) { case Seq(structRef) => buildExpression(structRef) }
+    } else {
+      buildExpression(structExpr)
     }
-    UpdateFieldsExpression(structRef, updatedValues.toSeq, outputOrdinals, newFields)
   }
 }
 
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceUpdateFieldsExpressionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceUpdateFieldsExpressionSuite.scala
index 50280a90ab9db..044243fb9dc31 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceUpdateFieldsExpressionSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceUpdateFieldsExpressionSuite.scala
@@ -19,7 +19,7 @@ package org.apache.spark.sql.catalyst.optimizer
 
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.dsl.plans._
-import org.apache.spark.sql.catalyst.expressions.{Alias, AttributeReference, Expression, Literal,
+import org.apache.spark.sql.catalyst.expressions.{Alias, AttributeReference, DropField, Expression, Literal,
   UpdateFields, With, WithField}
 import org.apache.spark.sql.catalyst.plans.PlanTest
 import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan}
@@ -68,6 +68,24 @@ class ReplaceUpdateFieldsExpressionSuite extends PlanTest {
     }
   }
 
+  test("UpdateFields replacement does not introduce scopes for non-nested operations") {
+    val structType = StructType(Seq(
+      StructField("a", IntegerType), StructField("b", IntegerType)))
+    Seq(true, false).foreach { nullable =>
+      val struct = AttributeReference("struct", structType, nullable)()
+      Seq(
+        Seq(WithField("a", Literal(1))),
+        Seq(DropField("b")),
+        Seq(WithField("a", Literal(1)), DropField("b"))).foreach { operations =>
+        val optimized = Optimize.execute(LocalRelation(struct).select(
+          Alias(UpdateFields(struct, operations), "result")()))
+        assert(!optimized.exists(_.expressions.exists(_.exists(_.isInstanceOf[With]))))
+        assert(!optimized.exists(_.expressions.exists(_.exists(_.isInstanceOf[UpdateFields]))))
+        assert(optimized.output.head.nullable == nullable)
+      }
+    }
+  }
+
   test("UpdateFields replacement does not grow with struct width when reused") {
     def nodeCount(width: Int, updateCount: Int): Int = {
       val structType = StructType((0 until width).map { index =>

From 9980682bc1b5c3ec0aebe2d7077963a98592f7c2 Mon Sep 17 00:00:00 2001
From: Ben Hollis <ben.hollis@databricks.com>
Date: Sun, 11 Oct 2026 04:00:17 +0000
Subject: [PATCH 6/8] [SPARK-60067][SQL] Reduce generated field-copy overhead

Use arithmetic indexing for contiguous copied fields and omit redundant null checks. Retain ordinal-array copying for sparse fields.

Add codegen regressions for field nullability and contiguous copies with shifted output positions. Validation: 366 tests passed across 11 suites; local projection benchmarks rerun.
---
 .../expressions/complexTypeCreator.scala      | 31 +++++++---
 .../ReplaceUpdateFieldsExpressionSuite.scala  | 57 ++++++++++++++++++-
 2 files changed, 78 insertions(+), 10 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala
index a91c937d00bea..4f6ce4d40cee1 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala
@@ -850,18 +850,33 @@ private case class UpdateFieldsExpression(
 
     val copyValuesCode = if (outputOrdinals.exists(_ >= 0)) {
       val copiedFields = outputOrdinals.zipWithIndex.filter(_._1 >= 0)
-        .groupBy { case (_, ordinal) => dataType(ordinal).dataType }
-      val copyCodes = copiedFields.toSeq.map { case (fieldType, fields) =>
+        .groupBy { case (_, ordinal) =>
+          (dataType(ordinal).dataType, dataType(ordinal).nullable)
+        }
+      val copyCodes = copiedFields.toSeq.map { case ((fieldType, fieldNullable), fields) =>
         val index = ctx.freshName("i")
         val sourceOrdinal = ctx.freshName("sourceOrdinal")
-        val ordinals = ctx.addReferenceObj("copyOrdinals",
-          fields.flatMap { case (source, output) => Seq(output, source) }.toArray, "int[]")
+        val (firstSource, firstOutput) = fields.head
+        val contiguous = fields.zipWithIndex.forall { case ((source, output), offset) =>
+          source == firstSource + offset && output == firstOutput + offset
+        }
+        val (loopEnd, loopStep, sourceIndex, outputIndex) = if (contiguous) {
+          (fields.length, 1, s"$index + $firstSource", s"$index + $firstOutput")
+        } else {
+          val ordinals = ctx.addReferenceObj("copyOrdinals",
+            fields.flatMap { case (source, output) => Seq(output, source) }.toArray, "int[]")
+          (fields.length * 2, 2, s"$ordinals[$index + 1]", s"$ordinals[$index]")
+        }
         val value = CodeGenerator.getValue(sourceRow, fieldType, sourceOrdinal)
+        val copiedValue = if (fieldNullable) {
+          s"$sourceRow.isNullAt($sourceOrdinal) ? null : $value"
+        } else {
+          value
+        }
         s"""
-           |for (int $index = 0; $index < ${fields.length * 2}; $index += 2) {
-           |  int $sourceOrdinal = $ordinals[$index + 1];
-           |  $values[$ordinals[$index]] =
-           |    $sourceRow == null || $sourceRow.isNullAt($sourceOrdinal) ? null : $value;
+           |for (int $index = 0; $index < $loopEnd; $index += $loopStep) {
+           |  int $sourceOrdinal = $sourceIndex;
+           |  $values[$outputIndex] = $copiedValue;
            |}
          """.stripMargin
       }
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceUpdateFieldsExpressionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceUpdateFieldsExpressionSuite.scala
index 044243fb9dc31..7a940c68b8931 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceUpdateFieldsExpressionSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceUpdateFieldsExpressionSuite.scala
@@ -17,10 +17,12 @@
 
 package org.apache.spark.sql.catalyst.optimizer
 
+import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.dsl.plans._
-import org.apache.spark.sql.catalyst.expressions.{Alias, AttributeReference, DropField, Expression, Literal,
-  UpdateFields, With, WithField}
+import org.apache.spark.sql.catalyst.expressions.{Alias, AttributeReference, BoundReference,
+  DropField, Expression, Literal, UnsafeProjection, UpdateFields, With, WithField}
+import org.apache.spark.sql.catalyst.expressions.codegen.CodegenContext
 import org.apache.spark.sql.catalyst.plans.PlanTest
 import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan}
 import org.apache.spark.sql.catalyst.rules.RuleExecutor
@@ -86,6 +88,57 @@ class ReplaceUpdateFieldsExpressionSuite extends PlanTest {
     }
   }
 
+  test("UpdateFields codegen checks only nullable copied fields") {
+    Seq(true, false).foreach { structNullable =>
+      Seq(true, false).foreach { fieldNullable =>
+        val structType = StructType(Seq(
+          StructField("a", IntegerType, nullable = false),
+          StructField("b", IntegerType, nullable = fieldNullable)))
+        val struct = BoundReference(0, structType, structNullable)
+        val expression = UpdateFields(struct, Seq(WithField("a", Literal(7)))).evalExpr
+        val code = expression.genCode(new CodegenContext).code.toString
+        val checksCopiedField = code.linesIterator.exists { line =>
+          line.contains("sourceRow") && line.contains(".isNullAt(")
+        }
+        assert(checksCopiedField == fieldNullable)
+        val projection = UnsafeProjection.create(Seq(expression))
+        val result = projection(InternalRow(InternalRow(1, 2))).getStruct(0, 2)
+        assert(result.getInt(0) == 7)
+        assert(result.getInt(1) == 2)
+        if (fieldNullable) {
+          assert(projection(InternalRow(InternalRow(1, null))).getStruct(0, 2).isNullAt(1))
+        }
+        if (structNullable) {
+          assert(projection(InternalRow(null)).isNullAt(0))
+        }
+      }
+    }
+  }
+
+  test("UpdateFields codegen does not load ordinal arrays for contiguous copies") {
+    Seq(5, 100, 500).foreach { width =>
+      val structType = StructType((0 until width).map { index =>
+        StructField(s"field$index", IntegerType, nullable = false)
+      })
+      val struct = BoundReference(0, structType, nullable = false)
+      Seq(false, true).foreach { dropField =>
+        val operations = Seq(WithField("field0", Literal(7))) ++
+          (if (dropField) Seq(DropField("field1")) else Nil)
+        val expression = UpdateFields(struct, operations).evalExpr
+        val context = new CodegenContext
+        expression.genCode(context)
+        assert(!context.references.exists(_.isInstanceOf[Array[Int]]))
+        val outputWidth = if (dropField) width - 1 else width
+        val result = UnsafeProjection.create(Seq(expression))(
+          InternalRow(InternalRow.fromSeq(0 until width))).getStruct(0, outputWidth)
+        assert(result.getInt(0) == 7)
+        (1 until outputWidth).foreach { index =>
+          assert(result.getInt(index) == index + (if (dropField) 1 else 0))
+        }
+      }
+    }
+  }
+
   test("UpdateFields replacement does not grow with struct width when reused") {
     def nodeCount(width: Int, updateCount: Int): Int = {
       val structType = StructType((0 until width).map { index =>

From 049f52ef8c22270fd855e897c272cd45b8e05e74 Mon Sep 17 00:00:00 2001
From: Ben Hollis <ben.hollis@databricks.com>
Date: Mon, 12 Oct 2026 16:50:44 +0000
Subject: [PATCH 7/8] [SPARK-60067][SQL] Preserve pruning and names in compact
 field updates

Prune discarded named-struct values and preserve update-chain simplification and legacy generated names. Remove quadratic ordinal indexing, restore single-pass lowering, and correct CSE test comments.

Validation: 372 tests passed across 11 suites.
---
 .../expressions/complexTypeCreator.scala      | 68 ++++++++++-----
 .../sql/catalyst/optimizer/ComplexTypes.scala |  3 +-
 .../sql/catalyst/optimizer/Optimizer.scala    |  4 +-
 .../spark/sql/catalyst/util/package.scala     |  3 +
 .../SubexpressionEliminationSuite.scala       | 11 ++-
 .../ReplaceUpdateFieldsExpressionSuite.scala  | 87 ++++++++++++++++++-
 .../spark/sql/ColumnExpressionSuite.scala     | 58 +++++++++++++
 7 files changed, 199 insertions(+), 35 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala
index 4f6ce4d40cee1..c2346245454e4 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala
@@ -747,6 +747,8 @@ case class UpdateNestedField(name: String, fieldOps: Seq[StructFieldsOperation])
 
   override def children: Seq[Expression] = fieldOps
 
+  override def prettyName: String = "WithField"
+
   override protected def withNewChildrenInternal(
       newChildren: IndexedSeq[Expression]): Expression =
     copy(fieldOps = newChildren.map(_.asInstanceOf[StructFieldsOperation]))
@@ -794,18 +796,18 @@ private case class UpdateFieldsExpression(
 
   override lazy val dataType: StructType = {
     val sourceFields = structExpr.dataType.asInstanceOf[StructType]
-    var updatedIndex = 0
-    StructType(outputOrdinals.zipWithIndex.map { case (sourceOrdinal, index) =>
+    val updated = indexedUpdatedValues.iterator
+    StructType(indexedOrdinals.iterator.zip(outputFields.iterator).map {
+      case (sourceOrdinal, field) =>
       if (sourceOrdinal >= 0) {
         sourceFields(sourceOrdinal)
       } else {
-        val updatedValue = updatedValues(updatedIndex)
-        updatedIndex += 1
-        outputFields(index).copy(
+        val updatedValue = updated.next()
+        field.copy(
           dataType = updatedValue.dataType,
           nullable = updatedValue.nullable)
       }
-    })
+    }.toArray)
   }
 
   override def eval(input: InternalRow): Any = {
@@ -888,7 +890,7 @@ private case class UpdateFieldsExpression(
       ""
     }
 
-    val updatedOutputOrdinals = outputOrdinals.indices.filter(outputOrdinals(_) < 0)
+    val updatedOutputOrdinals = indexedOrdinals.indices.filter(indexedOrdinals(_) < 0)
     assert(updatedOutputOrdinals.length == updatedValues.length)
     val updatedCodes = updatedValues.zip(updatedOutputOrdinals).map { case (expr, ordinal) =>
       val eval = expr.genCode(ctx)
@@ -987,6 +989,27 @@ case class UpdateFields(structExpr: Expression, fieldOps: Seq[StructFieldsOperat
 
   override def prettyName: String = "update_fields"
 
+  private[sql] def sqlString(
+      render: Expression => String,
+      quoteField: String => String): String = {
+    def updateSql(inputSql: String, operations: Seq[StructFieldsOperation]): String = {
+      val operationsSql = operations.map {
+        case nested: UpdateNestedField =>
+          val fieldSql = if (structExpr.resolved) {
+            s"$inputSql.${quoteField(nested.name)}"
+          } else {
+            s"$inputSql[${render(Literal(nested.name))}]"
+          }
+          s"WithField(${updateSql(fieldSql, nested.fieldOps)})"
+        case operation => render(operation)
+      }
+      s"$prettyName(${(inputSql +: operationsSql).mkString(", ")})"
+    }
+    updateSql(render(structExpr), fieldOps)
+  }
+
+  override def sql: String = sqlString(_.sql, quoteIdentifier)
+
   private lazy val newFieldSources: Seq[(StructField, UpdateFields.FieldSource)] =
     UpdateFields.fieldSources(structExpr, fieldOps)
 
@@ -1005,20 +1028,11 @@ case class UpdateFields(structExpr: Expression, fieldOps: Seq[StructFieldsOperat
 
   lazy val evalExpr: Expression = {
     def buildExpression(structRef: Expression): Expression = {
-      val source = structExpr match {
-        case namedStruct: CreateNamedStruct => namedStruct
-        case _ => structRef
-      }
-      val fieldSources = UpdateFields.fieldSources(source, fieldOps)
+      val fieldSources = UpdateFields.fieldSources(structRef, fieldOps)
       val updatedValues = ArrayBuffer.empty[Expression]
       val outputOrdinals = fieldSources.map { case (_, source) =>
         source match {
-          case UpdateFields.Original(ordinal) => structExpr match {
-            case namedStruct: CreateNamedStruct =>
-              updatedValues += namedStruct.valExprs(ordinal)
-              -1
-            case _ => ordinal
-          }
+          case UpdateFields.Original(ordinal) => ordinal
           case UpdateFields.Updated(expr) =>
             updatedValues += expr
             -1
@@ -1026,10 +1040,20 @@ case class UpdateFields(structExpr: Expression, fieldOps: Seq[StructFieldsOperat
       }
       UpdateFieldsExpression(structRef, updatedValues.toSeq, outputOrdinals, newFields)
     }
-    if (fieldOps.exists(_.isInstanceOf[UpdateNestedField])) {
-      With(structExpr) { case Seq(structRef) => buildExpression(structRef) }
-    } else {
-      buildExpression(structExpr)
+    structExpr match {
+      case namedStruct: CreateNamedStruct =>
+        val values = namedStruct.valExprs.toArray
+        CreateNamedStruct(newFieldSources.flatMap { case (field, source) =>
+          val value = source match {
+            case UpdateFields.Original(ordinal) => values(ordinal)
+            case UpdateFields.Updated(expr) => expr
+          }
+          Seq(Literal(field.name), value)
+        })
+      case _ if fieldOps.exists(_.isInstanceOf[UpdateNestedField]) =>
+        With(structExpr) { case Seq(structRef) => buildExpression(structRef) }
+      case _ =>
+        buildExpression(structExpr)
     }
   }
 }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ComplexTypes.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ComplexTypes.scala
index 1a81584bc986e..ca577333150bc 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ComplexTypes.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ComplexTypes.scala
@@ -32,7 +32,8 @@ object SimplifyExtractValueOps extends Rule[LogicalPlan] {
       // Remove redundant field extraction.
       case GetStructField(createNamedStruct: CreateNamedStruct, ordinal, _) =>
         createNamedStruct.valExprs(ordinal)
-      case GetStructField(u: UpdateFields, ordinal, _) =>
+      case GetStructField(u: UpdateFields, ordinal, _)
+          if !u.structExpr.isInstanceOf[UpdateFields] =>
         val structExpr = u.structExpr
         u.newExpr(ordinal) match {
           // if the struct itself is null, then any value extracted from it (expr) will be null
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
index 97cbe03b00d2e..7ba4493c0961e 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
@@ -301,9 +301,7 @@ abstract class Optimizer(catalogManager: CatalogManager)
       RemoveNoopOperators),
     // This batch must be executed after the `RewriteSubquery` batch, which creates joins.
     Batch("NormalizeFloatingNumbers", Once, NormalizeFloatingNumbers),
-    // Repeat because an update can contain another update. Keep the resulting With expressions so
-    // their definitions remain inside conditional branches until evaluation.
-    Batch("ReplaceUpdateFieldsExpression", fixedPoint, ReplaceUpdateFieldsExpression)))
+    Batch("ReplaceUpdateFieldsExpression", Once, ReplaceUpdateFieldsExpression)))
 
     // remove any batches with no rules. this may happen when subclasses do not add optional rules.
     batches.filter(_.rules.nonEmpty)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/package.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/package.scala
index 75a1d6a046081..0d9648cedf746 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/package.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/package.scala
@@ -90,6 +90,9 @@ package object util extends Logging {
   private def usePrettyExpression(
       e: Expression,
       shouldTrimTempResolvedColumn: Boolean = false): Expression = e transform {
+    case u: UpdateFields if u.fieldOps.exists(_.isInstanceOf[UpdateNestedField]) =>
+      PrettyAttribute(
+        u.sqlString(toPrettySQL(_, shouldTrimTempResolvedColumn), identity), u.dataType)
     case a: Attribute => new PrettyAttribute(a)
     case Literal(s: UTF8String, StringType) => PrettyAttribute(s.toString, StringType)
     case Literal(v, t: NumericType) if v != null => PrettyAttribute(v.toString, t)
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SubexpressionEliminationSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SubexpressionEliminationSuite.scala
index eb0dddc14a086..2b20894c3fa43 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SubexpressionEliminationSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SubexpressionEliminationSuite.scala
@@ -625,16 +625,15 @@ class SubexpressionEliminationSuite extends SparkFunSuite with ExpressionEvalHel
         assert(states.nonEmpty, s"nothing was recorded for $expr (peel = $peel)")
         // The `With` itself stays a candidate wherever it is reached: deduplicating it as a whole
         // is safe, since it carries its own definitions and brings their slots into scope wherever
-        // it is generated. The exception is an `And` root with the peel on, where the peel lands on
-        // the `With` and the guard drops it along with its subtree -- no opportunity is lost there,
-        // because the peel never hands back the node it lands on, only that node's children.
+        // it is generated. With an `And` root and the peel on, the `With` itself is skipped; its
+        // child is still traversed, but candidates with unbound references are rejected.
         if (!(peel && expr.isInstanceOf[And])) {
           assert(states.exists(_.expr.isInstanceOf[With]),
             s"the With itself stopped being a candidate for $expr (peel = $peel)")
         }
         states.foreach { state =>
-          // A candidate may hold a reference only if it also holds the `With` that binds it -- the
-          // whole `With`, or something above it, is a legal candidate; anything below it is not.
+          // A candidate may hold a reference only if it also holds the `With` that binds it.
+          // Reference-free candidates below a `With` are legal too.
           // The ids are collected over the whole candidate rather than per reference, so a
           // reference sitting beside its binder rather than under it would pass; `With.apply`
           // cannot build that shape.
@@ -642,7 +641,7 @@ class SubexpressionEliminationSuite extends SparkFunSuite with ExpressionEvalHel
           val boundIds =
             state.expr.collect { case inScope: With => inScope.defs.map(_.id) }.flatten.toSet
           assert(refIds.subsetOf(boundIds),
-            s"a candidate below the With was added for $expr (peel = $peel): ${state.expr}")
+            s"an unbound reference was added for $expr (peel = $peel): ${state.expr}")
           assert(!state.expr.isInstanceOf[CommonExpressionDef],
             s"an unevaluable definition was added for $expr (peel = $peel): ${state.expr}")
         }
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceUpdateFieldsExpressionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceUpdateFieldsExpressionSuite.scala
index 7a940c68b8931..d667885185f93 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceUpdateFieldsExpressionSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceUpdateFieldsExpressionSuite.scala
@@ -18,20 +18,23 @@
 package org.apache.spark.sql.catalyst.optimizer
 
 import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.analysis.{UnresolvedAttribute, UnresolvedExtractValue}
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.dsl.plans._
 import org.apache.spark.sql.catalyst.expressions.{Alias, AttributeReference, BoundReference,
-  DropField, Expression, Literal, UnsafeProjection, UpdateFields, With, WithField}
+  CreateNamedStruct, DropField, Expression, GetStructField, Literal, UnsafeProjection, UpdateFields,
+  With, WithField}
 import org.apache.spark.sql.catalyst.expressions.codegen.CodegenContext
 import org.apache.spark.sql.catalyst.plans.PlanTest
 import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan}
 import org.apache.spark.sql.catalyst.rules.RuleExecutor
-import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
+import org.apache.spark.sql.catalyst.util.{quoteIdentifier, toPrettySQL}
+import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType}
 
 class ReplaceUpdateFieldsExpressionSuite extends PlanTest {
   private object Optimize extends RuleExecutor[LogicalPlan] {
     override protected def batches: Seq[Batch] =
-      Batch("ReplaceUpdateFieldsExpression", FixedPoint(10),
+      Batch("ReplaceUpdateFieldsExpression", Once,
         ReplaceUpdateFieldsExpression) :: Nil
   }
 
@@ -56,6 +59,7 @@ class ReplaceUpdateFieldsExpressionSuite extends PlanTest {
       .select(Alias(nestedUpdates(struct, depth), "result")())
     val optimized = Optimize.execute(plan)
     assert(optimized.exists(_.expressions.exists(_.exists(_.isInstanceOf[With]))))
+    assert(!optimized.exists(_.expressions.exists(_.exists(_.isInstanceOf[UpdateFields]))))
     optimized.collect {
       case node => node.expressions.map(_.collect { case expression => expression }.size).sum
     }.sum
@@ -70,6 +74,42 @@ class ReplaceUpdateFieldsExpressionSuite extends PlanTest {
     }
   }
 
+  test("UpdateFields replacement prunes named struct inputs in one pass") {
+    val source = CreateNamedStruct(Seq(
+      Literal("nested"), CreateNamedStruct(Seq(Literal("a"), Literal(1))),
+      Literal("discarded"), Literal(100)))
+    val updated = UpdateFields(UpdateFields(source, "nested.b", Literal(2)), "discarded")
+    val plan = LocalRelation().select(Alias(updated, "result")())
+    val replaced = ReplaceUpdateFieldsExpression(OptimizeUpdateFields(plan))
+    val expressions = replaced.flatMap(_.expressions)
+    assert(!expressions.exists(_.exists(_.isInstanceOf[UpdateFields])))
+    assert(!expressions.exists(_.exists(_.isInstanceOf[With])))
+    assert(!expressions.exists(_.exists(_ == Literal(100))))
+    comparePlans(ReplaceUpdateFieldsExpression(replaced), replaced)
+  }
+
+  test("nested UpdateFields SQL matches expanded operations") {
+    Seq("nested", "a.b", "a`b").foreach { name =>
+      val nestedType = StructType(Seq(StructField("a", IntegerType)))
+      val structType = StructType(Seq(StructField(name, nestedType)))
+      val resolved = AttributeReference("s", structType)()
+      Seq(resolved, UnresolvedAttribute("s")).foreach { struct =>
+        val nested = if (struct.resolved) {
+          GetStructField(struct, 0, Some(name))
+        } else {
+          UnresolvedExtractValue(struct, Literal(name))
+        }
+        val expanded = UpdateFields(struct, Seq(WithField(name,
+          UpdateFields(nested, Seq(WithField("a", Literal(2)))))))
+        val compact = UpdateFields(struct, s"${quoteIdentifier(name)}.a", Literal(2))
+        assert(compact.sql == expanded.sql)
+        if (struct.resolved) {
+          assert(toPrettySQL(compact) == toPrettySQL(expanded))
+        }
+      }
+    }
+  }
+
   test("UpdateFields replacement does not introduce scopes for non-nested operations") {
     val structType = StructType(Seq(
       StructField("a", IntegerType), StructField("b", IntegerType)))
@@ -160,4 +200,45 @@ class ReplaceUpdateFieldsExpressionSuite extends PlanTest {
         s"$updateCount updates: width 10 has $widthTen nodes, width 100 has $widthHundred")
     }
   }
+
+  test("wide UpdateFields replacements preserve rewritten field types and values") {
+    for {
+      width <- Seq(100, 500)
+      nullable <- Seq(false, true)
+    } {
+      val structType = StructType((0 until width).map { index =>
+        StructField(s"field$index", IntegerType, nullable = false)
+      })
+      val struct = BoundReference(0, structType, nullable)
+      val operations = (0 until width by 3).map { index =>
+        WithField(s"field$index", Literal(index))
+      }.toList
+      val expression = UpdateFields(struct, operations).evalExpr
+      val rewritten = expression.withNewChildren(
+        expression.children.head +: expression.children.tail.map(_ => Literal("updated")))
+      val expectedType = StructType((0 until width).map { index =>
+        StructField(s"field$index",
+          if (index % 3 == 0) StringType else IntegerType, nullable = false)
+      })
+      assert(rewritten.dataType == expectedType)
+      assert(rewritten.nullable == nullable)
+      val input = InternalRow(InternalRow.fromSeq(0 until width))
+      val interpreted = rewritten.eval(input).asInstanceOf[InternalRow]
+      val projection = UnsafeProjection.create(Seq(rewritten))
+      val generated = projection(input).getStruct(0, width)
+      Seq(interpreted, generated).foreach { result =>
+        (0 until width).foreach { index =>
+          if (index % 3 == 0) {
+            assert(result.getUTF8String(index).toString == "updated")
+          } else {
+            assert(result.getInt(index) == index)
+          }
+        }
+      }
+      if (nullable) {
+        assert(rewritten.eval(InternalRow(null)) == null)
+        assert(projection(InternalRow(null)).isNullAt(0))
+      }
+    }
+  }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala
index 48003d551ef36..2a101c5629718 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala
@@ -600,6 +600,64 @@ class ColumnExpressionSuite extends SharedSparkSession {
     }
   }
 
+  test("struct field updates should prune discarded aggregate expressions") {
+    onEachEvalPath {
+      withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
+        val df = Seq((1L, 0L)).toDF("x", "z")
+        val source = struct(sum($"x").as("a"), sum($"x" / $"z").as("b"))
+        checkAnswer(df.agg(source.dropFields("b")), Row(Row(1L)))
+        checkAnswer(df.agg(source.withField("b", lit(2))), Row(Row(1L, 2)))
+        checkAnswer(df.agg(source.as("s")).select(col("s").dropFields("b")),
+          Row(Row(1L)))
+        val nested = struct(source.as("inner"), sum($"x").as("c"))
+        checkAnswer(df.agg(nested.withField("inner.d", lit(2)).dropFields("inner.b")),
+          Row(Row(Row(1L, 2), 1L)))
+      }
+    }
+  }
+
+  test("extracting a re-updated nested struct should prune unrelated updates") {
+    onEachEvalPath {
+      withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
+        val nestedType = StructType(Seq(StructField("x", IntegerType)))
+        val structType = StructType(Seq(StructField("a", nestedType)))
+        Seq(false, true).foreach { nullable =>
+          val schema = StructType(Seq(
+            StructField("s", structType, nullable),
+            StructField("z", IntegerType, nullable = false)))
+          val rows = Seq(Row(Row(Row(1)), 0), Row(Row(null), 0)) ++
+            (if (nullable) Seq(Row(null, 0)) else Nil)
+          val df = spark.createDataFrame(sparkContext.parallelize(rows), schema)
+          val input = if (nullable) {
+            col("s").getField("a").withField("x", lit(1) / col("z"))
+          } else {
+            col("s").withField("a.x", lit(1) / col("z")).getField("a")
+          }
+          val updated = input.withField("y", lit(2)).getField("y")
+          checkAnswer(df.select(updated), Seq(Row(2), Row(null)) ++
+            (if (nullable) Seq(Row(null)) else Nil))
+          val projected = df.select(input.as("t"))
+          checkAnswer(projected.select(col("t").withField("y", lit(2)).getField("y")),
+            Seq(Row(2), Row(null)) ++ (if (nullable) Seq(Row(null)) else Nil))
+        }
+      }
+    }
+  }
+
+  test("nested field updates should preserve generated column names") {
+    val df = Seq((((1, 2), (3, 4)), 0)).toDF("s", "z")
+    val added = df.select(
+      col("s").withField("_1.x", lit(1)),
+      col("s").withField("_2.x", lit(1)))
+    assert(added.columns.toSeq == Seq(
+      "update_fields(s, WithField(update_fields(s._1, WithField(1))))",
+      "update_fields(s, WithField(update_fields(s._2, WithField(1))))"))
+    val dropped = df.select(col("s").dropFields("_1._2"), col("s").dropFields("_2._2"))
+    assert(dropped.columns.toSeq == Seq(
+      "update_fields(s, WithField(update_fields(s._1, dropfield())))",
+      "update_fields(s, WithField(update_fields(s._2, dropfield())))"))
+  }
+
   test("nested field updates should share an unresolved nondeterministic input") {
     onEachEvalPath {
       val counter = new java.util.concurrent.atomic.AtomicLong()

From 013eaf997e150448bcdd75116cbf1d91bc22bca2 Mon Sep 17 00:00:00 2001
From: Ben Hollis <ben.hollis@databricks.com>
Date: Mon, 12 Oct 2026 17:33:42 +0000
Subject: [PATCH 8/8] [SPARK-60067][SQL] Simplify field-source matching

Match field sources directly when building output ordinals, removing a redundant nested match without changing behavior.
---
 .../catalyst/expressions/complexTypeCreator.scala    | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala
index c2346245454e4..02c3446fd13fb 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala
@@ -1030,13 +1030,11 @@ case class UpdateFields(structExpr: Expression, fieldOps: Seq[StructFieldsOperat
     def buildExpression(structRef: Expression): Expression = {
       val fieldSources = UpdateFields.fieldSources(structRef, fieldOps)
       val updatedValues = ArrayBuffer.empty[Expression]
-      val outputOrdinals = fieldSources.map { case (_, source) =>
-        source match {
-          case UpdateFields.Original(ordinal) => ordinal
-          case UpdateFields.Updated(expr) =>
-            updatedValues += expr
-            -1
-        }
+      val outputOrdinals = fieldSources.map {
+        case (_, UpdateFields.Original(ordinal)) => ordinal
+        case (_, UpdateFields.Updated(expr)) =>
+          updatedValues += expr
+          -1
       }
       UpdateFieldsExpression(structRef, updatedValues.toSeq, outputOrdinals, newFields)
     }