HyukjinKwon commented on code in PR #59283:
URL: https://github.com/apache/spark/pull/59283#discussion_r4254325346


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

Review Comment:
   **Non-blocking (P2):** `u.dataType` runs `UpdateFields.fieldSources`, which 
starts with an unchecked `structExpr.dataType.asInstanceOf[StructType]`. 
`toSQLExpr` goes through this case when it builds the 
`DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE` error for an `UpdateFields` whose 
input is not a struct. That happens from `CheckAnalysis`, and also from 
`fieldSources` itself via `failOnTypeCheckResult`. Either way, the cast throws 
before the `AnalysisException` is constructed.
   
   For example, `spark.range(1).select($"id".withField("a.b", lit(1)))` now 
fails with a raw `java.lang.ClassCastException`. So do `$"a".withField("b.x.y", 
lit(2))` and `$"a".dropFields("b.x.y")` when `a.b` is an INT. The error has no 
error class, message parameters, or query context, and Spark Connect clients 
get the same unclassified error. In the merge target, these calls failed 
analysis with an `AnalysisException` that names the field and its type. Valid 
nested updates are not affected. The existing "called on a non-StructType 
column" and "intermediate field is not a struct" tests use one- and two-part 
paths, which never leave an `UpdateNestedField` in the failing node, so they 
still pass.
   
   Taking the data type only from a well-typed update keeps the rendered text 
and avoids the cast. `PrettyAttribute` defaults to `NullType`, and 
`checkInputDataTypes()` reuses the same cached field sources that `dataType` 
uses:
   
   ```scala
       case u: UpdateFields if 
u.fieldOps.exists(_.isInstanceOf[UpdateNestedField]) =>
         val sql = u.sqlString(toPrettySQL(_, shouldTrimTempResolvedColumn), 
identity)
         if (u.checkInputDataTypes().isSuccess) {
           PrettyAttribute(sql, u.dataType)
         } else {
           PrettyAttribute(sql)
         }
   ```
   
   **Verification:**
   - **Regression:** withField and dropFields with a two-part path on a 
non-struct column fail with an AnalysisException whose condition is 
DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE.
   - **Regression:** withField and dropFields with a path that continues at 
least two parts past a non-struct intermediate field fail with an 
AnalysisException, not ClassCastException.
   - **Compatibility:** Generated column names for valid single nested 
withField and dropFields calls still match the existing expected strings.



##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala:
##########
@@ -806,35 +987,131 @@ 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, 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)
-    })
+  private def getFieldExpr(ordinal: Int): Expression = structExpr match {
+    case c: CreateNamedStruct => c.valExprs(ordinal)
+    case _ => GetStructField(structExpr, ordinal)
+  }
 
-    if (structExpr.nullable) {
-      If(IsNull(structExpr), Literal(null, dataType), createNamedStructExpr)
+  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
+        }
+      }
+      UpdateFieldsExpression(structRef, updatedValues.toSeq, outputOrdinals, 
newFields)

Review Comment:
   Thanks, verified. A struct(...) input now lowers to a CreateNamedStruct with 
only the surviving values, and no copy of the original struct or With 
definition is kept. Dropped and overwritten aggregates, Python UDFs, and 
subqueries are no longer in the lowered tree. The new tests cover drops, 
overwrites, nested paths, and projection inlining, and check that the discarded 
value and With are gone.
   
   <!-- SPARK_DEV_REVIEW_REPLY 
{"feedback_id":"inline:4241668312","thread_id":"inline:4241668312","verdict_sha256":"376a843e0801b33f093c16ee386ce4fd6a12672aa472f4cfb03d6ca9bc61f79e"}
 -->



##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ComplexTypes.scala:
##########
@@ -32,9 +32,9 @@ 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] =>
+      case GetStructField(u: UpdateFields, ordinal, _) =>

Review Comment:
   Thanks, verified. The guard is back, so the chain merges before the 
extraction is simplified. The new ANSI tests cover nested extraction, 
projection inlining, and nullable inputs. Agreed that the nullable top-level 
variant was already in the merge target.
   
   <!-- SPARK_DEV_REVIEW_REPLY 
{"feedback_id":"inline:4241668315","thread_id":"inline:4241668315","verdict_sha256":"376a843e0801b33f093c16ee386ce4fd6a12672aa472f4cfb03d6ca9bc61f79e"}
 -->



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to