bhollis-dbx opened a new pull request, #58800: URL: https://github.com/apache/spark/pull/58800
### What changes were proposed in this pull request? This PR reduces per-node allocations when Catalyst rules traverse plan and expression trees that they leave unchanged. It touches only generic `TreeNode` traversal and is behavior-preserving. - `TreeNode.mapChildren` no longer eagerly builds `children.map(f)`. It visits each child once, returns the original node when every mapped child is `fastEquals` to its original, and defers allocating a replacement collection until the first materially changed child. When a rebuild is required it retains the earlier equal-but-distinct mapped copies (preserving their tags) and calls `withNewChildrenInternal` directly, skipping the redundant `fastEquals` recheck that `withNewChildren` would repeat. - `transformUpWithPruning` applies the rule directly to leaf nodes after the existing pruning checks, skipping the child-mapping closure and collection for leaves. - The `PartialFunction.applyOrElse` identity fallback is shared through a single `TreeNodeIdentity.fn` instead of allocating an identity function per node, and `CurrentOrigin.withOrigin` is replaced by an explicit save/set/restore helper (`applyRule`) that drops the per-node by-name closure while preserving origin semantics on both normal and exceptional returns. `mapChildrenWithReferenceEquality`, the arity-trait (`LeafLike`/`UnaryLike`/`BinaryLike`/`TernaryLike`/`QuaternaryLike`) `mapChildren` overrides, and the `BestEffortLazyVal`-backed property caches (`treePatternBits`, `containsChild`, `height`, `hashCode`) are unchanged. ### Why are the changes needed? Catalyst rules repeatedly traverse large, immutable plan and expression trees that they mostly leave unchanged. Even for an unmodified node, the generic `TreeNode` traversal still allocates on every node: a replacement child collection, a per-node identity fallback, a `CurrentOrigin` by-name closure, and eager memoizer state for cached tree properties whether or not those properties are ever read. This cost is paid by every applicable analyzer and optimizer rule, on every pass, and scales with tree size, so queries whose analyzed plan expands into hundreds of branches amplify it. On a synthetic 500-branch compatibility view (4,001 analyzed plan nodes, ~378K expression occurrences) this unchanged-tree allocation overhead dominated optimizer time. See SPARK-59501. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? New unit tests in `TreeNodeSuite`: - `mapChildren` returns the original node when every mapped child is `fastEquals`, including when each child is an equal-but-distinct copy, and retains equal replacements at non-adjacent indices when a later child changes. - Leaf and non-leaf rule visitation, and that pruning and ineffective-rule tracking are unchanged. - Transform rules observe the node origin and the previous origin is restored afterward, including when a rule throws. The full suite passes locally: `build/sbt 'catalyst/testOnly org.apache.spark.sql.catalyst.trees.TreeNodeSuite'` (42 tests, 0 failures). ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Isaac This pull request and its description were written by Isaac. -- 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]
