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

[bhollis-dbx]

What changes were proposed in this pull request?

Lower UpdateFields to a compact evaluator that stores unchanged field ordinals as metadata and evaluates changed fields separately. Support interpreted execution and code generation without expanding the Catalyst tree by struct width.

Why are the changes needed?

Reconstructing every field makes Catalyst plan size proportional to struct width, increasing optimizer, binding, and code-generation costs for wide and nested structs.

Does this PR introduce any user-facing change?

No.

How was this patch tested?

Added Catalyst coverage for expression growth by update depth and struct width. Added SQL coverage for nullable, nondeterministic, reused, nested, short-circuited, omitted, and throwable expressions across interpreted and code-generated execution. Ran the focused ReplaceUpdateFieldsExpressionSuite, OptimizeWithFieldsSuite, ComplexTypesSuite, and ColumnExpressionSuite tests.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Codex

[HyukjinKwon]

Review summary

The compact evaluator is a good direction for wide structs, but the nested-path sharing scheme causes regressions on the main Column API path. The eq-based sharing in evalExpr breaks for unresolved inputs, and the new guards turn off the flatten/simplify rewrites that used to keep UpdateFields chains linear. As a result, chained nested-path withField/dropFields calls on col()-style columns grow exponentially, both during analysis (the guard also disables the analyzer's Substitution-batch flatten) and at execution. The same guards, together with the tag getFieldExpr now puts on ordinary reads, mean dropped or overwritten values can be evaluated again, which can turn previously successful queries into ANSI errors.

Separately, the generated copy loop's generic get breaks on ColumnarRow structs with TIME/GEOMETRY/GEOGRAPHY fields. The interpreted evaluator indexes Scala Lists per row. There are also smaller cleanups: StructFieldsOperation.apply is now duplicated, and one evalExpr comment is misleading. The new tests only use pre-resolved expressions or a single optimizer rule, so they don't cover the paths where these problems show up. There is also an open question about per-row cost compared with the old lowering.

Cross-cutting observations

• The new UpdateFieldsExpression copy path has only been checked on UnsafeRow-backed inputs. Its generic getter breaks ColumnarRow sources under codegen, its List-backed state makes interpreted eval quadratic in width, and its per-row cost compared with the old typed CreateNamedStruct path is unmeasured. Switching to typed getters and indexed state also gives a fairer baseline for the UpdateFieldsBenchmark comparison that pass-a-3 asks for.
  Separation: Each issue has its own mechanism and fix. The benchmark question stays advisory and does not gate either defect fix.

Findings

9 total: 0 P0, 2 P1, 4 P2, 3 P3.

Blocking (P1)

• Nested-path sharing relies on object identity that analysis does not preserve — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala:999 — see inline.
• Guard also disables the analyzer's early flatten, making analysis exponential for nested-path chains — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/UpdateFields.scala:73 — see inline.

Non-blocking (P2)

• Blocking flatten makes dropped or overwritten nested values evaluate (and fail) — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/UpdateFields.scala:73
  Because flattening is now skipped for any chain that contains a generated nested read, a later dropFields/withField no longer removes the earlier value before evaluation. With ANSI on (the default), Seq((((1, 2), 3), 0)).toDF("s", "z").select($"s".withField("_1.d", lit(1) / $"z").dropFields("_1")) returned Row(Row(3)) before. Now UpdateFields(U1, [DropField(_1)]) lowers to With(U1) { ref => UpdateFieldsExpression(ref, [], [1]) }, which evaluates U1 (including 1 / z) and throws DIVIDE_BY_ZERO. The same happens for .withField("a.b", expensive).withField("a.b", other), where the overwritten value is still computed, and for getField on such a chain, which builds the whole chain. This also goes against the intent of the new dropFields should not evaluate omitted CreateNamedStruct fields test.

  See Shared repair plan 1 in the review body.

• Generic get in the codegen copy loop fails on ColumnarRow for TIME/GEOMETRY/GEOGRAPHY fields — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala:863 — see inline.

• getFieldExpr tags ordinary extraction reads, so later simplification and flattening are blocked — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala:989 — see inline.

• New tests bypass the analyzer path where the sharing and growth guarantees break — sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala:507 — see inline.

Nit (P3)

• fieldSources duplicates StructFieldsOperation.apply, leaving the trait contract unused and stale — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala:1061 — see inline.
• Interpreted eval indexes Scala Lists per row, making it quadratic in struct width — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala:822 — see inline.
• evalExpr comment says withField values are not rewritten, but replaceStruct rewrites them — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala:995 — see inline.

Shared repair plans

Shared repair plan 1

Covered findings:

• Nested-path sharing relies on object identity that analysis does not preserve — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala:999
• Blocking flatten makes dropped or overwritten nested values evaluate (and fail) — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/UpdateFields.scala:73
• Guard also disables the analyzer's early flatten, making analysis exponential for nested-path chains — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/UpdateFields.scala:73
• getFieldExpr tags ordinary extraction reads, so later simplification and flattening are blocked — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala:989
• New tests bypass the analyzer path where the sharing and growth guarantees break — sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala:507

Recommended change: Replace the duplicated-input-plus-tag scheme. Represent a nested-path update as an operation on the current value of the parent field, with a node-owned, input-relative reference instead of a copy of the struct expression. fieldSources/evalExpr bind that reference to structRef, so the input is evaluated once by construction. newExpr/SimplifyExtractValueOps bind it to GetStructField(structExpr, ordinal). OptimizeUpdateFields keeps flattening chains and composes a nested update with a preceding WithField of the same name. Remove GENERATED_STRUCT_READ/OWNER, the getFieldExpr tagging, and both guards. Add full-pipeline DataFrame tests (analyzer plus optimizer) for chained nested-path withField/dropFields on unresolved columns, including drop or overwrite after a nested update.

Why this works: Because the nested read refers to its own UpdateFields' input rather than to a second copy of it, analysis cannot break the sharing, flattening stays semantically valid (ops apply in order to the current struct), and the existing rewrites again keep chains linear and prune dropped or overwritten values.

Scope: Rework how UpdateFields encodes nested-path reads and restore the chain-collapsing rewrites, plus end-to-end tests.

Compatibility: Results, nullability, and error behavior of withField/dropFields for deterministic inputs stay identical to the merge target.

Risks: A new reference leaf must have a correct dataType/nullability before and after resolution, and must be handled by the single-pass resolver. Dedupe and flatten composition must preserve the order of operations on the same field name under case-insensitive resolution.

Constraints: Keep the compact UpdateFieldsExpression evaluator and its null and short-circuit semantics. Keep Column API behavior identical for deterministic inputs.

Success: A chain of N nested-path withField/dropFields calls on an unresolved column analyzes and optimizes to a tree of size linear in N. Each nested-path update in such a chain is evaluated at most once per row, and a nondeterministic struct input is evaluated once. Values dropped or overwritten later in a chain are not evaluated, including withField("_1.d", 1 / z).dropFields("_1") under ANSI. Ordinary withField(...).getField(...) reads simplify as in the merge target.

Verification:

• Regression: A DataFrame chain of many nested-path withField/dropFields calls on col()-style columns completes analysis and optimization with a bounded, linear plan node count.
• Regression: Dropping or overwriting a nested-path value whose expression would raise an ANSI error returns the merge-target result instead of throwing.
• Behavior: A nondeterministic, input-dependent struct expression used through the Column API with nested-path withField produces internally consistent fields under both interpreted and codegen evaluation.
• Compatibility: The existing UpdateFields optimizer and Column withField/dropFields suites keep their expected plans and results.

PR description suggestions

• The title is missing a closing bracket. It should be [SPARK-60067][SQL] Keep UpdateFields compact through execution.
• "Does this PR introduce any user-facing change? No." is not quite accurate. With this change, withField/dropFields evaluate a nondeterministic struct input once per row instead of once per copied field, which the new tests assert. Please mention that observable change.

Decision challenges

Has per-row execution cost been measured against the old lowering?

Two things I verified in the code: EquivalentExpressions.childrenToRecurse returns Nil for With, so nothing inside a lowered UpdateFields is a subexpression-elimination candidate anymore. Before, a non-nullable input lowered to a bare CreateNamedStruct whose children were candidates; for example, f($"id") shared between two withField values, or with another projected column, is now evaluated twice. Separately, copied fields go through the generic get(ordinal, DataType) (PhysicalDataType.apply plus an instanceof chain per field per row) instead of typed getters. What I couldn't confirm is whether this matters overall. Could you share UpdateFieldsBenchmark numbers before and after, for both narrow and wide structs? It may also be worth skipping the With when the input is a cheap leaf (attribute or bound reference) or when no nested read was rebound, and emitting typed getters for copied fields.

Generated by Omnigent on Databricks.

[HyukjinKwon]

Blocking (P1): This only shares the input when the generated read's child is the same object as structExpr. On the common Column API path the struct is unresolved when UpdateFields.apply builds the tree, for example col("s").withField("a.b", lit(1)).withField("a.c", lit(2)). ColumnResolutionHelper.innerResolve then resolves each occurrence of that shared instance separately (mapChildren, no memoization), so after analysis the two positions hold distinct but equal copies and eq fails. Because the new guards in OptimizeUpdateFields and SimplifyExtractValueOps also skip these chains, nothing collapses the copies either. At each level of the chain the previous update appears twice, so n chained nested-path calls run about 2^(n-1) lowered copies of the first update per row, and CSE can't merge them because it doesn't descend into With. The merge target reduced the same chain to a single UpdateFields(s, a -> UpdateFields(GetStructField(s, a), ...)). For nondeterministic inputs, the evaluate-once guarantee doesn't hold on this path either.

See Shared repair plan 1 in the review body.

[bhollis-dbx]

Replaced identity-based sharing with UpdateNestedField, which applies nested operations to the current field value without duplicating the parent input. Added unresolved-column and input-dependent nondeterministic regressions.

[HyukjinKwon]

Thanks, verified on the current head. UpdateNestedField applies nested operations to the current field value, so the parent input is no longer duplicated, and the With in evalExpr evaluates a non-literal input once. The new unresolved-column and nondeterministic-input tests cover the path I raised.

[HyukjinKwon]

Nit (P3): fieldSources re-implements the WithField/DropField name matching instead of going through StructFieldsOperation.apply, and nothing in production calls apply anymore. That leaves two copies of the same semantics that can drift apart, a trait scaladoc ("...used as the fields argument for StructType and as the children argument for CreateNamedStruct inside UpdateFields") that is no longer true, and an internal error for any other operation. Could apply be made generic over the value type (e.g. apply[T](values: Seq[(StructField, T)], updated: Expression => T)) and called from here? Or remove it and update the doc so there's a single source of truth.

Verification:

• Compatibility: Existing withField/dropFields name-matching behavior, including case sensitivity and duplicate names, is unchanged.

[bhollis-dbx]

Removed the unused StructFieldsOperation.apply implementations and updated the documentation. Name matching now lives in fieldSources.

[HyukjinKwon]

Thanks, verified. The apply implementations and their scaladoc are gone, and name matching now lives only in fieldSources.

[HyukjinKwon]

Blocking (P1): This guard also runs in the analyzer: OptimizeUpdateFields is in the Substitution batch (Analyzer.scala:641-645) specifically because "when manipulating deeply nested schema, UpdateFields expression tree could be very complex and make analysis impossible". Each nested-path step u_k = UpdateFields(u_{k-1}, a -> UpdateFields(ExtractValue(u_{k-1}, a), op_k)) references u_{k-1} twice. Before, flatten plus dedupe collapsed every step before resolution, so the tree stayed linear. Now (fieldOps1 ++ fieldOps2).exists(_.exists(isGeneratedStructRead)) is true for every nested-path step, so nothing collapses, and resolution and later traversals walk about 2^N occurrences for N chained calls. Examples are col("s").withField("a.x1", ..)...withField("a.xN", ..) or dropFields("a.x1", ..., "a.xN"), which Column.dropFields folds into such a chain. Programmatic updates of many nested fields can stall analysis before execution even starts.

See Shared repair plan 1 in the review body.

[bhollis-dbx]

Removed the guard and restored chain flattening. Nested operations compose before resolution. Added full-pipeline tests for chained nested withField and dropFields calls on unresolved columns.

[HyukjinKwon]

Thanks, verified. The flatten case is back without the guard, nested operations merge in optimizeFieldOps, and the new test bounds analyzed and optimized plan size for chained nested updates on unresolved columns.

[HyukjinKwon]

Non-blocking (P2): The previous lowering read each copied field through GetStructField codegen, i.e. CodeGenerator.getValue with typed getters (getLong, getBinaryView, getInterval, ...). The generic get(ordinal, DataType) here isn't supported by every InternalRow: ColumnarRow.get (ColumnarRow.java:191-235) has no branch for TimeType, GeometryType, GeographyType, or CalendarIntervalType and throws _LEGACY_ERROR_TEMP_3155. Structs from the vectorized Parquet reader arrive as ColumnarRow (vector.getStruct(rowId)), so with geospatial enabled by default, spark.read.parquet(p).select($"s".withField("b", lit(1))) on s STRUCT<g: GEOMETRY(4326), a: INT> now fails at runtime, where it used to succeed. The same goes for TIME fields when the time type is enabled. This only affects the code-generated path; interpreted GetStructField already used the generic get. Emitting CodeGenerator.getValue(sourceRow, fieldType, ordinal) per copied ordinal would keep the previous compatibility.

Verification:

• Regression: A vectorized Parquet struct with an unchanged GEOMETRY or TIME field passes through withField under default codegen without SparkUnsupportedOperationException.

[bhollis-dbx]

Generated copies now use CodeGenerator.getValue, grouped by field type. Added vectorized Parquet regressions for unchanged TIME, GEOMETRY, and GEOGRAPHY fields.

[HyukjinKwon]

Thanks, verified. Copied fields now go through CodeGenerator.getValue per type group, and the new vectorized Parquet tests cover TIME, GEOMETRY, and GEOGRAPHY.

[HyukjinKwon]

Non-blocking (P2): newExpr's only caller is SimplifyExtractValueOps, so this tag ends up on ordinary single-level extractions, not just on nested-path reads. $"t".withField("y", lit(2)).getField("a") becomes a tagged GetStructField(t, a), and the tag stays because copyTagsFrom only fills untagged nodes. If t is later inlined as s.withField("x", lit(1) / $"s.b") (e.g. by CollapseProject), the new !isGeneratedStructRead(field) guard refuses to simplify GetStructField(UpdateFields(s, x), a) to s.a. The full update, including x, is then evaluated, so with ANSI on and s.b = 0 this now fails with DIVIDE_BY_ZERO where it used to return rows. A simplified read like this inside an ops list also trips the flatten guard at UpdateFields.scala:73, even though no nested-path read would be separated. This conflicts with the comment that says generated reads are "field accesses introduced while expanding nested paths".

See Shared repair plan 1 in the review body.

[bhollis-dbx]

Removed the generated-read tags and simplification guards. Ordinary field extraction simplifies again, with a regression covering an unused throwing update.

[HyukjinKwon]

Thanks, the tag-based block is gone and ordinary extraction simplifies again. One related regression: the revision also removed the merge-target !u.structExpr.isInstanceOf[UpdateFields] guard in SimplifyExtractValueOps. An extraction from a two-level chain is now rewritten to If(IsNull(innerUpdate), ...) before OptimizeUpdateFields can merge the chain, so the inner update's values run. I left the details in a new comment on ComplexTypes.scala.

[HyukjinKwon]

Nit (P3): outputOrdinals and updatedValues are Scala 2.13 Lists here whenever a WithField is present (ArrayBuffer.toSeq returns a List, and so does its .map). This loop calls outputOrdinals.length, outputOrdinals(i), and updatedValues(updatedIndex) for every field of every row, so interpreted evaluation is O(width^2) per row. That's roughly 125k+ list hops per row for a 500-field struct. needsSource is also recomputed with an exists scan on every row. Codegen isn't affected because it copies into an int[], but interpreted mode is what you get when codegen falls back on very large expressions. Building indexed collections (toArray/IndexedSeq, or @transient lazy val arrays) and hoisting needsSource into a lazy val would keep this linear.

Verification:

• Inspection: The per-row eval loop indexes only array-backed or IndexedSeq state, and needsSource is computed once per expression instance.

[bhollis-dbx]

Converted per-row indexed state to arrays and cached needsSource. The evaluation loop no longer indexes Scala Lists or rescans source requirements.

[HyukjinKwon]

Thanks, the per-row loop now uses the arrays and is linear. The same List indexing is still in dataType (outputFields(index), updatedValues(updatedIndex)) and in doGenCode (outputOrdinals.indices.filter(outputOrdinals(_) < 0)), which makes planning quadratic in width. Reusing indexedOrdinals/indexedUpdatedValues or zipped iterators there would fix it.

[HyukjinKwon]

Nit (P3): This reads as the opposite of what the next lines do. replaceStruct is applied to every Updated(expr), and for a nested path that value is the generated UpdateFields(GetStructField(structExpr, a), ...) whose tagged read is rewritten to structRef. The actual rule is narrower: only reads generated by this update's own nested-path expansion are rebound, and other reads of structExpr inside withField values are left alone. Maybe something like: "Leave other reads in withField values unchanged; they must be evaluated independently."

Verification:

• Inspection: The comment states the owner-tagged nested-read rule and no longer claims that withField values are never rewritten.

[bhollis-dbx]

Removed the comment along with the identity/tag-based rewriting machinery. Nested updates are now represented explicitly.

[HyukjinKwon]

Thanks, verified. The comment went away with the rewriting machinery it described.

[HyukjinKwon]

Non-blocking (P2): This test (and preserve a nested path through a later update) builds UpdateFields directly over nextStruct().expr, a zero-argument UDF that is already resolved. Identity therefore survives analysis and the eq-based sharing works. Users reach this code through Column.withField on input-dependent expressions such as udf(...)($"id").withField("_1.copy", lit(1)) or col("s")..., where analysis rebuilds the struct. ReplaceUpdateFieldsExpressionSuite also runs only ReplaceUpdateFieldsExpression on pre-resolved AttributeReference inputs, skipping both the analyzer and the operator-optimization batch. So none of the new tests would fail for the exponential chains, the evaluated dropped values, or the double evaluation on the normal DataFrame/Connect path. Could you add full-pipeline DataFrame tests for chained nested-path withField/dropFields on unresolved columns, including a drop or overwrite after a nested update?

See Shared repair plan 1 in the review body.

[bhollis-dbx]

Added DataFrame tests through analysis and optimization for unresolved nested chains, later drops and overwrites, and input-dependent nondeterministic structs.

[HyukjinKwon]

Thanks, verified. The new DataFrame tests go through analysis and optimization for unresolved nested chains, later drops and overwrites, and an input-dependent nondeterministic struct, which is what I asked for.

[bhollis-dbx]

Changes since the initial commit

• Nested updates now operate on the current parent-field value without duplicating the input expression. Restored chain flattening, pruning of overwritten/dropped values, and field-extraction simplification.
• Generated copies use typed getters; interpreted evaluation uses arrays rather than repeated sequence indexing.
• Preserved common-subexpression elimination through With, while respecting lexical bindings, lazy definitions, codegen fallback, and null guards. Shared unconditional UDF values are evaluated once per row.
• Flat updates no longer introduce unnecessary With scopes. This removed the measured planning regression for deeply nested, grouped updates.
• Generated copies use arithmetic indexing when source and output fields are contiguous, avoiding ordinal-array lookups. Sparse copies retain the indexed path. Null checks are omitted for non-nullable fields, and the source row is not rechecked inside the already-guarded copy loop.

Validation

366 tests passed across 11 SQL, Catalyst, optimizer, and codegen suites. New regressions cover nested chains, overwritten/dropped values, evaluate-once behavior, typed copies, CSE scope safety, nullability, and contiguous copies with shifted output positions.

Changed sources were compiled against the cached SBT test classpath and run with ScalaTest. Normal SBT project loading remains blocked by Jackson BOM resolution; runtests does not support the OSS checkout.

Benchmarks

Latest local projection microbenchmark, expanded struct construction versus compact execution, in ns/row; lower is better:

Fields	Baseline codegen	Compact codegen	Baseline interpreted	Compact interpreted
5	24	29	131	43
100	308	243	970	284
500	3295	1361	5198	1356

This measures replacing the first field of a non-nullable integer struct, leaving contiguous copied fields. Codegen is faster at 100 and 500 fields; the five-field case still costs about 5 ns/row more than baseline. Interpreted evaluation is faster at every measured width. Unrolling small copies did not improve the five-field result, so that experiment was discarded.

The latest projection run uses longer warmup and more iterations than the earlier measurements; compare baseline and compact within this table rather than absolute timings across runs.

Planning results from the preceding run of UpdateFieldsBenchmark, upstream versus compact, in average ms:

Case	Baseline non-nullable	Compact non-nullable	Baseline nullable	Compact nullable
Repeated dotted-path updates, 3 depths	31	2	509	1
Grouped updates per level, 100 depths	561	549	603	548

These are local microbenchmarks, not end-to-end SQL throughput measurements. UpdateFieldsBenchmark processes zero rows and measures planning only. The projection results do not establish performance for sparse copies, mixed field types, or nullable structs.

[HyukjinKwon]

Review summary

This revision resolves the previous round's blocking issues. Nested-path updates are now UpdateNestedField operations applied to the current value of the parent field, so the input is bound once without relying on object identity, the analyzer and optimizer flatten rewrites are back, and the generated-read tags are gone. The new end-to-end DataFrame tests cover chained nested-path updates on col() inputs and pruning of dropped or overwritten nested values. The ColumnarRow copy failure (now typed getters, with TIME and geospatial Parquet tests), the duplicated StructFieldsOperation.apply, the misleading evalExpr comment, and per-row List indexing in interpreted eval are fixed as well. On the earlier per-row cost question, CSE now looks inside With bodies, and a new test checks that shared values are still deduplicated.

The remaining findings are new in this round, and none is blocking:

• When the input is struct(...), the lowered UpdateFieldsExpression keeps the original CreateNamedStruct as a child. Rules that walk the whole plan after lowering, such as PhysicalAggregation and ExtractPythonUDFs, still compute dropped aggregates and Python UDFs, so dropping a field that holds a failing aggregate can fail a query that used to succeed.
• With the SimplifyExtractValueOps guard removed, reading one field from a re-updated nested struct can wrap the unmerged inner update in IsNull, which evaluates that whole update. This is new for non-nullable struct inputs and can raise ANSI errors.
• Nested-path updates render as updatenestedfield(...), which changes generated column names and lets updates on different parent fields collide.
• UpdateFieldsExpression.dataType and doGenCode still index List-backed arguments by position, which is quadratic in struct width at planning and codegen time.
• The new comment on the ReplaceUpdateFieldsExpression batch says the rule must repeat, but one top-down pass already lowers nested updates.
• Two comments in the SPARK-58902 CSE test still describe the old rule that excluded everything below a With.

Cross-cutting observations

• Two independent paths still evaluate values that a later dropFields or withField discards: the CreateNamedStruct lowering keeps the original struct as an unevaluated child, and extracting from a re-updated nested struct wraps the unmerged update in IsNull. The new pruning tests cover projections over attribute and literal-struct inputs only, so neither path is exercised.
  Separation: The two paths are in different rules with disjoint fixes. Each repair closes its own finding, and neither needs coordination through a shared owner.

Findings

6 total: 0 P0, 0 P1, 2 P2, 4 P3.

Non-blocking (P2)

• A struct(...) input stays in the lowered tree with its dropped and overwritten values — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala:1027 — see inline.
• Removing the chain guard lets a field read evaluate an unmerged inner update — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ComplexTypes.scala:35 — see inline.

Nit (P3)

• Nested-path updates render as updatenestedfield(...) in generated names — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala:745 — see inline.
• ReplaceUpdateFieldsExpression does not need a fixed-point batch — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala:304 — see inline.
• SPARK-58902 test still explains its assertion with the removed With rule — sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SubexpressionEliminationSuite.scala:607 — see inline.
• dataType and doGenCode still index List-backed arguments by position — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala:795 — see inline.

Re-review status

Prior AI findings: 9 addressed, 0 still present; additional unresolved findings in this review: 6.

New attribution: 3 newly introduced, 3 late catch, 0 previously raised, 0 unattributed.

Remaining prior AI findings

No prior AI findings remain.

Existing discussions

• existing discussion — The author's update addresses the previous review. Identity and tag-based sharing is replaced by UpdateNestedField, chain flattening and extraction simplification are restored, typed getters and array-backed state are used, and benchmarks answer the per-row cost question. The pinned source supports these claims, except the claim that pruning of dropped or overwritten values was restored: for CreateNamedStruct inputs the lowering still keeps the dropped fields in the tree, where planner walkers evaluate them.
• existing discussion — The generated-read tags and the flatten and simplification guards are gone, and ordinary extraction simplifies again, as the new prune test shows. The revision also removed the merge-target guard that kept SimplifyExtractValueOps from rewriting an extraction from a two-level update chain. That reintroduces the same consequence through a different mechanism: an unused update value is evaluated when an extraction cannot be simplified away.
• existing discussion — The per-row eval loop now uses the indexedOrdinals and indexedUpdatedValues arrays and a cached needsSource, so interpreted evaluation is linear. The same List-backed positional indexing remains in UpdateFieldsExpression.dataType and doGenCode, which now dominates planning-time cost for wide structs.

PR description suggestions

• The title is missing the closing bracket after the JIRA ID. It should read [SPARK-60067][SQL] Keep UpdateFields compact through execution.
• "Does this PR introduce any user-facing change? No." understates the change. withField and dropFields on a nondeterministic struct input now evaluate it once per row instead of once per copied field, so the copied fields agree with each other, as the new tests assert. Please describe that change.

Generated by Omnigent on Databricks.

[HyukjinKwon]

Non-blocking (P2): For a CreateNamedStruct input, every output value already comes from namedStruct.valExprs (all outputOrdinals are -1), but the original struct is still passed here as structExpr, and it is also kept as the With definition on line 1030 for nested paths. needsSource is false, so it is never evaluated per row, but it is still a child, together with the values that dropFields or a later withField removed.

Rules that walk the whole expression tree after ReplaceUpdateFieldsExpression still see those values. For example, with ANSI on:

Seq((1L, 0L)).toDF("x", "z")
  .agg(struct(sum($"x").as("a"), sum($"x" / $"z").as("b")).dropFields("b"))

PhysicalAggregation collects AggregateExpressions from the whole result expression (planning/patterns.scala:301-307), so it still computes sum(x / z) and fails with DIVIDE_BY_ZERO. Before this PR the same query returned Row(Row(1)). Python UDFs in a dropped field are extracted and run the same way, because ExtractPythonUDFs runs after the optimizer batches, and a scalar subquery in a dropped field is still planned. Inputs reach this path through direct struct(...).dropFields/withField calls and also when CollapseProject inlines a projected struct(...) into a later update.

Row-level evaluation in a plain projection is not affected, which is why dropFields should not evaluate omitted CreateNamedStruct fields passes. For nested paths over struct(...) inputs, the retained With definition also keeps a full copy of the input struct at every nesting level, so the lowered tree grows quadratically with depth there.

Verification:

• Regression: An aggregate query that drops a struct(...) field computing a failing ANSI expression returns the surviving fields instead of throwing.
• Regression: A dropped struct(...) field holding a side-effecting or failing value is not evaluated when the update is consumed by an aggregate or a projection.
• Inspection: The lowered expression for a nested-path update over a struct(...) input contains no With and holds each surviving value once.
• Compatibility: Existing withField/dropFields results and nullability over struct(...) inputs are unchanged.

[bhollis-dbx]

Fixed in 049f52e. Named-struct inputs now lower to a CreateNamedStruct containing only surviving values; the original struct and unused With definitions are removed. Added failing-aggregate regressions covering dropped/overwritten fields, nested paths, and projection inlining, plus inspection assertions that discarded values and With nodes are absent.

[HyukjinKwon]

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.

[HyukjinKwon]

Non-blocking (P2): Before this PR, the guard if !u.structExpr.isInstanceOf[UpdateFields] left an extraction from an update chain alone until OptimizeUpdateFields had merged the chain. OptimizeUpdateFields runs before SimplifyExtractValueOps in the batch, so without the guard this arm can see a chain that is not merged yet and rewrite it to If(IsNull(innerUpdate), null, value) whenever the inner update is nullable. NullPropagation only folds IsNull over non-nullable children, so after lowering IsNull evaluates the whole inner update, including values the extracted field does not depend on.

For example, with a non-nullable struct s whose field a is a nullable struct, ANSI on, and z = 0:

col("s").withField("a.x", lit(1) / col("z")).getField("a")
  .withField("y", lit(2)).getField("y")

In one pass the inner extraction becomes UpdateFields(s.a, [x = 1 / z]) and the outer one becomes If(IsNull(UpdateFields(s.a, [x = 1 / z])), null, 2), which throws DIVIDE_BY_ZERO. With the guard, the next iteration merged the chain into UpdateFields(s.a, [x, y]) and simplified the read to If(IsNull(s.a), null, 2), which returns 2. The same shape appears when CollapseProject inlines an earlier withColumn that updated s.a.

Nullable top-level inputs already behaved this way before the PR; non-nullable ones are newly affected. Keeping the guard costs one more optimizer iteration, which the merge target already paid.

Verification:

• Regression: An extraction from a re-updated, extracted nested struct does not evaluate the discarded inner update values, for non-nullable and nullable inputs.
• Compatibility: Existing SimplifyExtractValueOps and UpdateFields optimizer results for single-level updates are unchanged.

[bhollis-dbx]

Restored the guard in 049f52e so chains merge before extraction simplification. Added ANSI-failure regressions for nested extraction and projection-inlined inputs, plus direct extracted-nested chains with nullable inputs. This fixes the newly introduced regression; it does not claim to fix the preexisting nullable-top-level variant.

[HyukjinKwon]

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.

[HyukjinKwon]

Nit (P3): UpdateNestedField overrides neither prettyName nor sql, so Expression.sql renders it as updatenestedfield(<child ops>): the lowercased class name, without the parent field name. Unaliased nested-path withField/dropFields columns get their name from this rendering. df.select(col("s").withField("a.b", lit(1))) used to produce a column named update_fields(s, WithField(update_fields(s.a, WithField(1)))) and now produces update_fields(s, updatenestedfield(WithField(1))); dropFields("a.b") changes the same way to update_fields(s, updatenestedfield(dropfield())).

This exposes an internal class name in the output schema, in classic and Connect, and updates on different parents now collide: withField("a.x", lit(1)) and withField("c.x", lit(1)) get the same name. Analyzed-plan EXPLAIN output and sqlExpr error parameters that render the outer update change the same way. Query results are not affected, and no test pins the generated names.

Verification:

• Compatibility: Generated column names of unaliased nested-path withField and dropFields match the merge-target names.
• Regression: Two nested updates on different parent fields produce distinct generated names.

[bhollis-dbx]

Fixed in 049f52e. Nested updates render like the legacy expanded operations, retaining parent-field names without reconstructing expanded trees. Tests compare SQL rendering and generated names, including dotted/backtick field names, and verify that different parents produce distinct names.

[HyukjinKwon]

Thanks, single nested updates now render with the legacy names, and the test checks that different parents get distinct names. Chained nested calls still get different names than before. For example, withField("a.x", ...).withField("a.y", ...) or dropFields("a.b", "a.c") used to embed the first update in the name, and now render as one merged update on s.a. I raised this as a question in this review: keep the old names or document the new ones.

[HyukjinKwon]

Nit (P3): ReplaceUpdateFieldsExpression is transformAllExpressionsWithPruning(...) { case u: UpdateFields => u.evalExpr } (UpdateFields.scala:99-103), a top-down transform that also recurses into the children of each replacement. Every UpdateFields that evalExpr creates (the nested updates built by fieldSources, or an input kept as a With definition) is a child of the returned UpdateFieldsExpression (structExpr +: updatedValues) or With (child +: defs), so the same pass lowers it. The second fixed-point iteration only confirms that nothing changed, so the comment describes a requirement the rule does not have.

Moving off Once also takes this batch out of the test-mode idempotence check that RuleExecutor runs on Once batches, and adds one more traversal of the plan.

Verification:

• Regression: Nested-path and chained withField/dropFields queries optimize with no remaining UpdateFields, and the test-mode Once idempotence check passes.

[bhollis-dbx]

Restored Once and removed the incorrect comment in 049f52e. Tests verify that nested updates fully lower in one pass and that applying the rule again changes nothing.

[HyukjinKwon]

Thanks, verified. The batch is Once again and the comment is gone. The new suite checks that nested updates lower fully in one pass and that a second application changes nothing.

[HyukjinKwon]

Nit (P3): This lead comment now describes the bound-reference rule, but two later comments in the same test still describe the removed one: around line 629 the guard is said to drop the With "along with its subtree", and around line 637 the comment says that below a With "anything below it is not" a legal candidate.

childrenToRecurse now returns Seq(withExpr.child) unless refUnderCodegenFallback is set, and supportedExpression only rejects candidates with unbound references. The new With exposes independent subexpressions but not unbound references test asserts that a reference-free subexpression below a With is a candidate, which contradicts both comments. Updating them to the bound-reference rule would keep this test's explanation accurate.

Verification:

• Inspection: The test comments describe the bound-reference rule enforced by supportedExpression and asserted by the new With tests.

[bhollis-dbx]

Updated in 049f52e. The comments now describe the bound-reference rule: independent expressions below With remain eligible; candidates with unbound references are rejected.

[HyukjinKwon]

Thanks, verified. Both comments now describe the bound-reference rule that the new With tests assert.

[HyukjinKwon]

Nit (P3): outputFields(index) and updatedValues(updatedIndex) are positional lookups, and the sequences evalExpr passes in are Lists on Scala 2.13 (updatedValues.toSeq from an ArrayBuffer, and the fieldSources result after any WithField). Each lookup is O(index), so building this type is O(F^2) for an F-field struct, and because it is a per-instance lazy val it is recomputed for every copy a tree transform makes. doGenCode has the same pattern in outputOrdinals.indices.filter(outputOrdinals(_) < 0) (line 891).

The per-row eval loop already uses indexedOrdinals and indexedUpdatedValues after the earlier review. Using those arrays here too, or walking the sequences with iterators, would keep planning and codegen linear in width for the wide structs this PR targets.

Verification:

• Inspection: dataType and doGenCode access outputs through arrays or a single iterator pass.
• Behavior: Wide-struct updates (hundreds of fields) keep identical output types and results under interpreted and codegen evaluation.

[bhollis-dbx]

Fixed in 049f52e. Type construction uses indexed ordinals and iterators; codegen uses indexed ordinals. Added 100/500-field checks for output types, nullability, null inputs, and interpreted/codegen results.

[HyukjinKwon]

Thanks, verified. dataType now walks the indexed ordinals and updated values with iterators, and doGenCode uses indexedOrdinals. The 100- and 500-field tests cover types, nullability, and interpreted and codegen results.

[bhollis-dbx]

Addressed the six new findings in 049f52e:

• Lower named-struct updates to a CreateNamedStruct containing only surviving values, removing discarded expressions and unused With definitions.
• Restore the extraction guard so update chains merge before simplification introduces null checks.
• Preserve legacy nested-update SQL rendering and generated column names without rebuilding expanded expression trees.
• Replace positional List lookups with indexed access or iterators.
• Restore single-pass (Once) lowering and correct the CSE test comments.

Added regression coverage for discarded failing aggregates, projection-inlined updates, nested extraction, generated names, single-pass lowering, and 100/500-field interpreted and codegen results.

Validation: 372 tests passed across 11 suites using freshly compiled sources against the cached SBT test classpath. Benchmarks were not rerun for this commit; previously reported results are unchanged measurements, not measurements of this revision.

[HyukjinKwon]

Review summary

This revision addresses the six findings from the previous review. A struct(...) input now lowers to a CreateNamedStruct that holds only the surviving values, with no With, so planner walkers no longer compute dropped aggregates or Python UDFs. The SimplifyExtractValueOps chain guard is back. UpdateFieldsExpression.dataType and doGenCode use indexed arrays and iterators instead of positional List access. Optimizer.scala is no longer changed, so the lowering batch runs Once again. The SPARK-58902 test comments now describe the bound-reference rule. Single nested-path updates keep the legacy generated names instead of updatenestedfield(...). The names of chained nested calls still differ, and that is raised below as a question rather than a defect.

One new non-blocking defect remains. The new usePrettyExpression case asks a failing UpdateFields for its dataType. So an invalid nested-path withField or dropFields through a non-struct value throws ClassCastException while its type-mismatch error is being built, instead of failing analysis with DATATYPE_MISMATCH. The existing error tests use paths too short to reach this case. Guarding the data type in that case and adding tests for the deeper error paths would close it.

Cross-cutting observations

• Both remaining items come from the generated-name preservation added in 049f52e. The new usePrettyExpression and sqlString rendering crashes when it renders an ill-typed nested update, and it still gives chained nested calls different names than the merge target.
  Separation: The crash concerns invalid inputs and has a local fix; the chained names concern valid inputs and depend on the author's choice between keeping the old names and documenting the new ones. Neither resolution settles the other, so the shared origin does not change grouping, severity, or repair scope.

Findings

1 total: 0 P0, 0 P1, 1 P2, 0 P3.

Non-blocking (P2)

• Invalid nested-path updates crash with ClassCastException while the error is rendered — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/package.scala:95 — see inline.

Re-review status

Prior AI findings: 6 addressed, 0 still present; additional unresolved findings in this review: 1.

New attribution: 1 newly introduced, 0 late catch, 0 previously raised, 0 unattributed.

Remaining prior AI findings

No prior AI findings remain.

Existing discussions

• existing discussion — All six fixes are present on the head (013eaf9, whose only later change is a no-op match cleanup). struct(...) inputs lower to a CreateNamedStruct of the surviving values with no With. The SimplifyExtractValueOps chain guard is restored. Nested updates render with the legacy names. dataType and doGenCode use the indexed arrays and iterators. Optimizer.scala is no longer changed, so the lowering batch is Once again. The CSE test comments describe the bound-reference rule. The naming fix introduced a new defect: the usePrettyExpression case it added asks a failing UpdateFields for its dataType, so invalid nested paths through non-struct values crash with ClassCastException while their type-mismatch error is being built.
• existing discussion — UpdateFields.sql and the new usePrettyExpression case render a single nested update like the legacy expansion, and the new test pins those names and checks that different parents get distinct names. The comment's compatibility check, that generated names of unaliased nested-path withField and dropFields match the merge-target names, still fails for chained nested calls. The analyzer merges them into one UpdateNestedField and renders it against the top-level input. That remaining gap is raised as a decision challenge.

PR description suggestions

• The title is missing the closing bracket after the JIRA ID. It should read [SPARK-60067][SQL] Keep UpdateFields compact through execution.
• "Does this PR introduce any user-facing change? No." understates the change. withField and dropFields on a nondeterministic struct input now evaluate it once per row instead of once per copied field, so the copied fields agree with each other, as the new tests assert. Please describe that change.

Decision challenges

Chained nested withField/dropFields calls get new generated column names

Chained nested-path calls now get different generated column names than in the merge target. Take df.select(col("s").withField("a.x", lit(1)).withField("a.y", lit(2))).columns. The merge target named it update_fields(s, WithField(update_fields(update_fields(s, WithField(update_fields(s.a, WithField(1)))).a, WithField(2)))), because the second nested value read field a of the first update. With this PR, OptimizeUpdateFields merges both operations into one UpdateNestedField(a, ...) during analysis, before aliases are generated. sqlString then renders every nested field against the top-level input, so the name becomes update_fields(s, WithField(update_fields(s.a, WithField(1), WithField(2)))).

The same happens with dropFields("a.b", "a.c"), which Column.dropFields folds into a chain, and with withField("_1.x", ...).withField("_2.y", ...). It also affects any nested-path call that follows another update in the chain. These names appear in df.columns, in the schema of unaliased writes, and through Spark Connect. The new "nested field updates should preserve generated column names" test pins only single nested calls.

Is this naming change intended? If it is, could the PR description mention it (it currently says there is no user-facing change), and could a test pin the chained names? If the old names should stay, each nested operation would need to render against the struct produced by the operations before it, including the merge target's last-wins deduplication.

Generated by Omnigent on Databricks.

[HyukjinKwon]

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:

    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.

[bhollis-dbx]

Fixed in 719f8ac. Error rendering now checks the update’s input types before reading its data type. Added regressions for withField/dropFields on non-struct inputs and deeper paths through non-struct intermediate fields, checking the classified analysis error and query context.

[HyukjinKwon]

Thanks, verified on the current head. usePrettyExpression now reads dataType only after checkInputDataTypes succeeds, so invalid nested paths through a non-struct value fail with DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE instead of ClassCastException. The new test covers two-part and deeper paths for both withField and dropFields, including the query context.

[HyukjinKwon]

Thanks, I checked these on the current head. struct(...) inputs now lower to a CreateNamedStruct that holds only the surviving values, with no With. The SimplifyExtractValueOps guard is back, nested updates keep the legacy generated names, dataType and doGenCode no longer index Lists by position, the lowering batch is Once again, and the CSE test comments describe the bound-reference rule. One regression comes from the naming change: the new usePrettyExpression case asks the failing UpdateFields for its dataType while the type-mismatch error is being built. A nested path through a non-struct value, such as $"id".withField("a.b", lit(1)), now throws a ClassCastException instead of an AnalysisException. Details are in this review.

[bhollis-dbx]

Is this naming change intended?

Keeping the compact chained names intentionally. Preserving the old names would require retaining operation history discarded during analysis. Added tests for same-parent, different-parent, overwritten, and mixed flat/nested updates, plus chained drops. Single nested calls retain their previous names. The PR description will document this change.

Changes

Pushed 719f8ac. Invalid nested updates now report DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE rather than ClassCastException. Added coverage for invalid paths and pinned the intentional compact names for chained nested updates.

All 374 tests passed across 11 suites using freshly compiled sources against the cached SBT test classpath. Benchmarks were not rerun.

[HyukjinKwon]

Review summary

The previous review's one finding is fixed. usePrettyExpression now renders the SQL text first and reads dataType only when checkInputDataTypes succeeds, so an invalid nested-path withField or dropFields through a non-struct value fails analysis with DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE instead of throwing ClassCastException. The new test covers two-part paths on a non-struct column and deeper paths through a non-struct intermediate for both operations. The previous round also asked whether the compact generated names for chained nested calls are intended. The new test 'chained nested field updates should use compact generated column names' pins them, which settles that choice in the code; the PR description still needs to mention it (see below).

This round adds three non-blocking findings. Nested-path updates re-match the parent field by name when they are lowered. If spark.sql.caseSensitive changes between analysis and execution of a classic DataFrame or plan-storing temp view, the lowered struct silently loses the update (P2). A positional replacement in fieldSources fixes this. The new With-based lowering for paths of three or more parts is exercised only on the default whole-stage path, so a nested-scope regression on the interpreted or non-whole-stage codegen path would go unnoticed (P3). optimizeFieldOps re-optimizes the accumulated list on every same-parent merge, so long programmatic chains analyze in quadratic time (P3).

Findings

3 total: 0 P0, 0 P1, 1 P2, 2 P3.

Non-blocking (P2)

• Nested updates re-match the parent field by name at lowering, so a case-sensitivity change between analysis and execution drops the update — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala:1130 — see inline.

Nit (P3)

• Nested With scopes from paths of three or more parts are not tested on the interpreted or non-whole-stage codegen paths — General
  A nested path of three or more parts on an input that is not a CreateNamedStruct lowers to nested With scopes. The outer With(structExpr) wraps an UpdateFieldsExpression whose updated value is UpdateFields(GetStructField(ref, i), Seq(UpdateNestedField(...))), and lowering that value creates a second With whose definition reads the outer reference. onEachEvalPath exists because With memoization is implemented separately for interpretation and for codegen. However, every new onEachEvalPath test in ColumnExpressionSuite uses a flat or two-part path (_1.copy, _1.d, a.x, ...), which produces a single With. The existing three-part tests, such as withField should add field to deeply nested struct, run only on the default whole-stage path.

  So a regression in binding the inner definition to the outer reference would pass every test under NO_CODEGEN, or under CODEGEN_ONLY with whole-stage codegen off, including when UpdateFieldsExpression.doGenCode moves updated values into split functions. Those are the paths used when codegen falls back for the large expressions this PR targets. An onEachEvalPath test with a three-part path on a nondeterministic UDF or attribute input, asserting the result and single evaluation of the input, would cover the shape.

  Verification:

  • Behavior: Nested updates three levels deep return the expected struct and evaluate a nondeterministic input once per row on all three evaluation paths.

• optimizeFieldOps re-optimizes the accumulated list on every same-parent merge, so long nested chains analyze in quadratic time — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/UpdateFields.scala:65 — see inline.

Re-review status

Prior AI findings: 1 addressed, 0 still present; additional unresolved findings in this review: 3.

New attribution: 0 newly introduced, 3 late catch, 0 previously raised, 0 unattributed.

Remaining prior AI findings

No prior AI findings remain.

Verification

• The previous review's ClassCastException finding is fixed: usePrettyExpression reads dataType only after a successful type check, and the new test asserts the AnalysisException for both invalid path shapes.

PR description suggestions

• "Does this PR introduce any user-facing change? No." understates three observable changes, and the answer should list them. (1) Unaliased chained nested-path withField/dropFields calls get compact generated column names, for example col("s").withField("_1.x", lit(1)).withField("_1.y", lit(2)) is now named update_fields(s, WithField(update_fields(s._1, WithField(1), WithField(2)))), which shows up in df.columns and in the schema of unaliased writes. (2) A nested path through a non-struct value now fails with DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE instead of the previous extraction error. (3) A nondeterministic struct input is evaluated once per row, so its copied fields agree with each other.

Generated by Omnigent on Databricks.

[HyukjinKwon]

Non-blocking (P2): fieldSources finds the parent field with the UpdateNestedField's own resolver (ExtractValue(..., nested.resolver) a few lines up). It then applies the update through this new WithField, whose resolver is SQLConf.get.resolver evaluated when it is constructed (line 710). fieldSources runs once at analysis (newFieldSources, which backs dataType and newFields) and again when ReplaceUpdateFieldsExpression lowers the node (evalExpr -> buildExpression). buildExpression pairs the recomputed ordinals with the analysis-time newFields.

So the two passes disagree when spark.sql.caseSensitive changes between analysis and execution. Take a struct column s with field a under case-insensitive resolution: val df2 = df.select(col("s").withField("A.x", lit(9))) analyzes to struct<A: struct<..., x>>. After spark.conf.set("spark.sql.caseSensitive", "true"), df2.collect() optimizes under the new setting. The nested lookup still finds ordinal 0, but WithField("A") no longer matches a, so it appends. That gives outputOrdinals = [0, -1] against a single output field. UpdateFieldsExpression.dataType zips them down to the original a, so x is silently dropped and the optimized type no longer matches the analyzed attribute. Going the other way on a struct that has both a and A overwrites both. A temp view that stores the analyzed plan reaches the same path. Before this PR, the nested WithField was built once at Column conversion, so this also needed an optimizer rewrite of that operation. This affects classic execution only; Spark Connect re-analyzes on every execution.

The parent's ordinal is already known at this point. Replacing current(field.ordinal) in place with StructField(nested.name, updated.dataType, updated.nullable) -> Updated(updated) would keep the analysis-time layout, and the rename to the call's casing, without a second name match.

Verification:

• Regression: Apply withField("A.x", ...) to a struct with field a under case-insensitive analysis, then execute after switching to case-sensitive: the result has the analyzed schema with x populated.
• Compatibility: Nested withField and dropFields results, schemas, and generated names under case-sensitive and case-insensitive resolution are unchanged when the setting does not change.

[HyukjinKwon]

Nit (P3): When a key already holds an UpdateNestedField, this runs optimizeFieldOps again on previous.fieldOps ++ nested.fieldOps. That re-maps every child already optimized and rebuilds the LinkedHashMap over the whole accumulated list. The flatten case first turns names.foldLeft(col("s"))((c, n) => c.withField(s"a.$n", v)) into a single UpdateFields(s, [UpdateNestedField(a, [f1]), ..., UpdateNestedField(a, [fm])]). The Substitution-batch application then does about m^2/2 operation visits, each with a key normalization, where the merge-target dedupe was a single linear pass. Programmatic updates of many fields under one parent are the wide-schema workload this PR targets, and the new linear-growth test bounds node counts rather than merge work. Only planning time is affected; results are unchanged.

Accumulating each key's nested operations in order and optimizing each key's list once, or merging into a per-key builder as operations arrive, would keep this linear. The output should stay exactly as today. That includes the point where a nested list containing DropField stops further deduplication, the handling of a preceding or following WithField on the same key, and the last-seen casing.

Verification:

• Compatibility: Optimized operation lists for mixed WithField/UpdateNestedField chains, including overwrites, nested lists containing DropField, and case-insensitive keys, are identical to the current output.
• Regression: A chain of thousands of same-parent nested withField calls analyzes with merge work that grows linearly with chain length.

[HyukjinKwon]

Review summary

All three findings from the previous review are addressed on this head. Nested updates now replace the parent at its resolved ordinal instead of re-matching it by name. A new test checks that the analyzed schema and rows survive a spark.sql.caseSensitive change, for both the DataFrame and a temp view. A new onEachEvalPath test runs a three-part nested withField plus dropFields on a nondeterministic input with a method-split threshold of 1. That covers the nested With scopes on the interpreted, non-whole-stage codegen, and whole-stage paths. optimizeFieldOps now collects each parent's nested chunks and optimizes them once, and a test bounds the merge to linear operation visits.

This round adds one non-blocking finding. The new scaladoc on UpdateFields.Updated says every computed value comes from a withField expression. fieldSources also stores the nested updates it synthesizes as Updated, and on the With path those read the bound input. Rewording that one line fixes it.

Findings

1 total: 0 P0, 0 P1, 0 P2, 1 P3.

Nit (P3)

• The Updated scaladoc says every computed field comes from withField — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala:1066 — see inline.

Re-review status

Prior AI findings: 3 addressed, 0 still present; additional unresolved findings in this review: 1.

New attribution: 0 newly introduced, 1 late catch, 0 previously raised, 0 unattributed.

Remaining prior AI findings

No prior AI findings remain.

Verification

• fieldSources replaces the resolved parent ordinal with current.updated, and a new test checks the analyzed schema and rows across case-sensitivity changes for a DataFrame and a temp view.
• A three-part nested withField/dropFields test on a nondeterministic input runs under onEachEvalPath with CODEGEN_METHOD_SPLIT_THRESHOLD set to 1.
• optimizeFieldOps accumulates nested chunks per parent key and optimizes them once; a test bounds merge visits linearly for 2000 same-parent updates.

PR description suggestions

• Replace "No." under "Does this PR introduce any user-facing change?" with the observable changes: (1) unaliased chained nested-path withField/dropFields calls get compact generated column names, for example col("s").withField("_1.x", lit(1)).withField("_1.y", lit(2)) is now named update_fields(s, WithField(update_fields(s._1, WithField(1), WithField(2)))), which shows up in df.columns and in the schema of unaliased writes; (2) a nested path through a non-struct value now fails with DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE; and (3) a nondeterministic struct input is evaluated once per row, so its copied fields agree with each other.

Generated by Omnigent on Databricks.

[HyukjinKwon]

Nit (P3): This says every Updated value is a withField expression, but fieldSources also stores Updated(updated) for each UpdateNestedField (line 1131), where updated is an UpdateFields built over the current parent field. That includes nested dropFields paths such as col("s").dropFields("a.b"), which have no withField value at all. On the With lowering path these values hold GetStructField(structRef, ...), so they depend on the bound input. A reader who trusts this comment could treat them as independent user expressions and move them out of that scope. Only the comment is affected; runtime behavior is unchanged. Maybe something like:

Suggested change

	/** Computes the field from an expression supplied to `withField`. */
	/** Computes the field from a `withField` value or a nested update of the current field. */

Verification:

• Inspection: The Updated doc matches every place fieldSources constructs an Updated source, including the UpdateNestedField branch.