[SPARK-60067[SQL] Keep UpdateFields compact through execution

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,175 @@ 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],
+    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 +946,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 +969,129 @@ 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)
+  /** 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
+  }
 
-  lazy val evalExpr: Expression = {
-    val createNamedStructExpr = CreateNamedStruct(newFieldExprs.flatMap {
-      case (field, expr) => Seq(Literal(field.name), 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)
+  }
 
-    if (structExpr.nullable) {
-      If(IsNull(structExpr), Literal(null, dataType), createNamedStructExpr)
-    } else {
-      createNamedStructExpr
+  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 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 {
+  /** 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)] = {
+    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}")
+    }
+  }
+
+  // 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
+    }
+  }
+
+  /** 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")
 
@@ -878,7 +1135,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
     }
   }
 }

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

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)

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)
   }