[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]

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]

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]

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]

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]

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]

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]

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.

[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.