HyukjinKwon commented on PR #58800:
URL: https://github.com/apache/spark/pull/58800#issuecomment-5673023985
<!-- ai-code-review -->
**Review — [SPARK-59501][SQL] Reduce allocations in unchanged TreeNode
traversals**
LGTM. This is a careful, allocation-reducing rewrite of a very hot Catalyst
path, and I read it as semantically equivalent to the original:
- **`TreeNodeIdentity.fn` + `applyRule`** replace the per-node
`identity[BaseType]` closure and the by-name `CurrentOrigin.withOrigin`
closure. `applyRule`'s save/set/`try`/`finally`-restore is exactly
`CurrentOrigin.withOrigin`'s body, so the origin is restored even when the rule
throws (the new "transform rules see node origins and restore the previous
origin" test, including its `intercept[RuntimeException]`, pins this down). The
`Any => Any` → `BaseType => BaseType` cast is safe since `identity` is
type-agnostic.
- **`transformUpWithPruning`** — the `children.isEmpty` fast path is
equivalent to the old code, because `mapChildren` on a leaf returns `this` and
never invokes `f` anyway; the non-empty branch is unchanged in effect.
- **`mapChildren` rewrite** — the rebuild branch is exactly
`withNewChildren`'s rebuild branch (`withOrigin(origin){
withNewChildrenInternal(...); copyTagsFrom(this) }`), skipping only the
size-assert and the `childrenFastEquals` recheck, which the inline `changed`
flag subsumes. The reconstructed child sequence is identical to
`children.map(f)`: `eq` positions store the reference-equal original (==
`f(child)`), `ne` positions store `f(child)`. `changed` ⟺ some child is not
`fastEquals`, matching `withNewChildren`'s return-`this` condition, so no
allocation happens unless a child materially changes.
The tests are the strong point — they exercise exactly the subtle cases:
all-`fastEquals` (returns `this`), every-child equal-but-distinct copy (still
returns `this`, buffer dropped), and retention of equal-but-distinct copies
(including a non-adjacent gap) when another child changes materially. That
retention matches the previous base `mapChildren` behavior. Note it differs
from the arity traits (`BinaryLike` etc.), which normalize an
equal-but-distinct child back to the original — but that difference is
pre-existing and out of scope here; this PR only touches the general-arity base
path.
No public API surface changes (all new members are `private`/`protected`
internal), so no Mima/`@since` considerations. Style is clean (100-col, ASCII).
No correctness concerns found.
--
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]
