You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "bhollis-dbx (via GitHub)" <gi...@apache.org> on 2026/09/14 23:37:14 UTC
[PR] [SPARK-59501][SQL] Reduce allocations in unchanged TreeNode traversals [spark]
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59501][SQL] Reduce allocations in unchanged TreeNode traversals [spark]
Posted by "bhollis-dbx (via GitHub)" <gi...@apache.org>.
bhollis-dbx commented on code in PR #58800:
URL: https://github.com/apache/spark/pull/58800#discussion_r4075549971
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala:
##########
@@ -559,15 +580,12 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]]
if (!cond.apply(this) || isRuleIneffective(ruleId)) {
return this
}
- val afterRuleOnChildren = mapChildren(_.transformUpWithPruning(cond, ruleId)(rule))
- val newNode = if (this fastEquals afterRuleOnChildren) {
- CurrentOrigin.withOrigin(origin) {
- rule.applyOrElse(this, identity[BaseType])
- }
+ val newNode = if (children.isEmpty) {
Review Comment:
Fixed in e87af127: gate the leaf fast path on LeafLike, with regression coverage for specialized nodes.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59501][SQL] Reduce allocations in unchanged TreeNode traversals [spark]
Posted by "gengliangwang (via GitHub)" <gi...@apache.org>.
gengliangwang closed pull request #58800: [SPARK-59501][SQL] Reduce allocations in unchanged TreeNode traversals
URL: https://github.com/apache/spark/pull/58800
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59501][SQL] Reduce allocations in unchanged TreeNode traversals [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on PR #58800:
URL: https://github.com/apache/spark/pull/58800#issuecomment-5771224071
Thanks for the ping. I reviewed the design of the generic TreeNode traversal changes; the rule-origin handling and child-rebuild behavior remain equivalent to the previous implementation.
<!-- SPARK_DEV_REVIEW_REPLY {"feedback_id":"conversation:5768350023","thread_id":"conversation:5768350023","verdict_sha256":"cc87f2c8f068c13a3e3bd824ff22fcd68f211f88bb922b0b17d4b6edcc8cfb3f"} -->
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59501][SQL] Reduce allocations in unchanged TreeNode traversals [spark]
Posted by "bhollis-dbx (via GitHub)" <gi...@apache.org>.
bhollis-dbx commented on code in PR #58800:
URL: https://github.com/apache/spark/pull/58800#discussion_r4075567492
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/trees/TreeNodeSuite.scala:
##########
@@ -180,6 +181,76 @@ class TreeNodeSuite extends SparkFunSuite with SQLHelper {
assert(actual === expect)
}
+ test("mapChildren returns the original node when all children are fast equal") {
+ val expression = Coalesce(Seq(Literal(1), Literal(2)))
+ val visited = new ArrayBuffer[Int]()
+ val result = expression.mapChildren {
+ case literal @ Literal(value: Int, _) =>
+ visited += value
+ literal
+ case other => other
+ }
+
+ assert(result eq expression)
+ assert(visited == Seq(1, 2))
+ val leaf = Dummy(None)
+ assert(leaf.mapChildren(identity) eq leaf)
+ }
+
+ test("mapChildren returns the original node when every child is equal but a distinct copy") {
+ val expression = Coalesce(Seq(Literal(1), Literal(2)))
+ // Return a fresh, structurally-equal (not reference-equal) copy for *every* child: this
+ // fills `equalCopies` yet must still return `this`, since no child materially changes.
+ val result = expression.mapChildren {
+ case Literal(value: Int, dt) => Literal(value, dt)
+ case other => other
+ }
+ assert(result eq expression)
+ }
+
+ test("mapChildren retains an equal replacement when another child changes") {
+ val tag = TreeNodeTag[String]("equal-copy")
+ val expression = Coalesce(Seq(Literal(1), Literal(2)))
+ val equalCopy = Literal(1)
+ equalCopy.setTagValue(tag, "retained")
+
+ val result = expression.mapChildren {
+ case Literal(1, _) => equalCopy
+ case Literal(2, _) => Literal(3)
+ case other => other
+ }
+
+ assert(result.children.head eq equalCopy)
+ assert(result.children.head.getTagValue(tag).contains("retained"))
+ assert(result.children(1) == Literal(3))
+ }
+
+ test("mapChildren retains non-adjacent equal replacements when a later child changes") {
+ val c0 = Literal(10)
+ val c1 = Literal(11)
+ val c2 = Literal(12)
+ val c3 = Literal(13)
+ val expression = Coalesce(Seq(c0, c1, c2, c3))
+ val copy0 = Literal(10)
+ val copy2 = Literal(12)
+
+ // Equal-but-distinct copies at non-adjacent indices 0 and 2, an unchanged same instance at 1,
+ // and a material change at 3. Exercises the replay loop's index bookkeeping across a gap.
Review Comment:
The test now keeps c0 reference-equal before copy1, exercising prefix backfill. It also asserts all five child identities/order and retains a later material change at c4.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59501][SQL] Reduce allocations in unchanged TreeNode traversals [spark]
Posted by "bhollis-dbx (via GitHub)" <gi...@apache.org>.
bhollis-dbx commented on code in PR #58800:
URL: https://github.com/apache/spark/pull/58800#discussion_r4075560403
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/trees/TreeNodeSuite.scala:
##########
@@ -180,6 +181,76 @@ class TreeNodeSuite extends SparkFunSuite with SQLHelper {
assert(actual === expect)
}
+ test("mapChildren returns the original node when all children are fast equal") {
+ val expression = Coalesce(Seq(Literal(1), Literal(2)))
+ val visited = new ArrayBuffer[Int]()
+ val result = expression.mapChildren {
+ case literal @ Literal(value: Int, _) =>
+ visited += value
+ literal
+ case other => other
+ }
+
+ assert(result eq expression)
+ assert(visited == Seq(1, 2))
+ val leaf = Dummy(None)
+ assert(leaf.mapChildren(identity) eq leaf)
+ }
+
+ test("mapChildren returns the original node when every child is equal but a distinct copy") {
+ val expression = Coalesce(Seq(Literal(1), Literal(2)))
+ // Return a fresh, structurally-equal (not reference-equal) copy for *every* child: this
+ // fills `equalCopies` yet must still return `this`, since no child materially changes.
Review Comment:
Added a reference-equal prefix, full identity/order assertions, and removed the stale equalCopies wording.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59501][SQL] Reduce allocations in unchanged TreeNode traversals [spark]
Posted by "gengliangwang (via GitHub)" <gi...@apache.org>.
gengliangwang commented on PR #58800:
URL: https://github.com/apache/spark/pull/58800#issuecomment-6283535708
**Merge Summary:**
- merged into master https://github.com/apache/spark/commit/42e7c6d59d4c0fd054c50d8541e02cfc4a4c6874
- merged into branch-4.x https://github.com/apache/spark/commit/46fbf5bb3ab2bcacb6f34dc60dd5eb9f047ecd0b
*Posted by `merge_spark_pr.py`*
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59501][SQL] Reduce allocations in unchanged TreeNode traversals [spark]
Posted by "bhollis-dbx (via GitHub)" <gi...@apache.org>.
bhollis-dbx commented on PR #58800:
URL: https://github.com/apache/spark/pull/58800#issuecomment-6276162964
Addressed in 3d72d9d13:
• Removed TreeNodeIdentity and its cast; moved applyRule beside the transforms.
• Restored equality short-circuiting after the first material child change.
• Added coverage for rebuild origin, tags, and exactly one mapping call per child in order.
Dropped the LeafLike fast path after a JDK 17 microbenchmark on a 1,501-node expression tree. Median traversal times with/without it were 138/96 µs with one thread and 207/91 µs with four threads. It saved about 24 KB per traversal but was slower. These are traversal measurements, not optimizer timings for the original view.
All 47 TreeNodeSuite tests pass through SBT on Java 17.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59501][SQL] Reduce allocations in unchanged TreeNode traversals [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #58800:
URL: https://github.com/apache/spark/pull/58800#discussion_r4309699823
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala:
##########
@@ -57,6 +57,14 @@ import org.apache.spark.util.collection.BitSet
/** Used by [[TreeNode.getNodeNumbered]] when traversing the tree for a given number */
private class MutableInt(var i: Int)
+// A single shared identity function used as the `PartialFunction.applyOrElse` fallback during tree
+// traversal, so a fresh `identity` closure is not allocated at every node. It is typed `Any => Any`
Review Comment:
**Nit (P3):** This rationale doesn't hold: the old fallback wasn't allocated per node. With Scala 2.13 (2.13.18 per the pom), `rule.applyOrElse(this, identity[BaseType])` compiles to a non-capturing `invokedynamic` lambda (`apply:()Lscala/Function1;` with a static body), and the JVM links that call site to one cached instance. I compiled the merge-base pattern with 2.13.18 and every `applyOrElse` call received the same default object. The per-node allocation came only from `withOrigin`'s by-name thunk, which captures the node and the rule, and `applyRule` already removes it. So this object, its comment, and the unchecked `asInstanceOf[BaseType => BaseType]` in `applyRule` add nothing. Passing `identity[BaseType]` directly in `applyRule` gives the same single instance without the cast.
**Verification:**
- **Behavior:** A transform rule that does not match a node returns the original node instance, so unchanged-tree transforms still return the original tree and mark rules ineffective as before.
- **Inspection:** The traversal code has no shared identity singleton or unchecked cast for the rule fallback, and no comment attributes an allocation saving to the identity fallback.
- **Compile:** The catalyst module compiles after the shared identity object is removed.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala:
##########
@@ -807,10 +825,53 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]]
* Returns a copy of this node where `f` has been applied to all the nodes in `children`.
*/
def mapChildren(f: BaseType => BaseType): BaseType = {
- if (containsChild.nonEmpty) {
- withNewChildren(children.map(f))
- } else {
+ val oldChildren = children
+ if (oldChildren.isEmpty) return this
+
+ // Invoke `f` once per child and defer allocating the replacement array until `f` returns a
+ // distinct instance. Backfill the reference-equal prefix when that first happens. `changed`
+ // records whether any change was material (`!fastEquals`). If no child changed materially we
+ // return `this` and drop the array -- an equal-but-distinct replacement alone does not require
+ // rebuilding the parent.
+ var newChildren: Array[BaseType] = null
+ var changed = false
+ var index = 0
+ val childIterator = oldChildren.iterator
+ while (childIterator.hasNext) {
+ val oldChild = childIterator.next()
+ val newChild = f(oldChild)
+ if (oldChild ne newChild) {
+ if (newChildren eq null) {
+ // First distinct replacement: allocate the result array and backfill the
+ // reference-equal prefix.
+ newChildren = new Array[TreeNode[_]](oldChildren.size).asInstanceOf[Array[BaseType]]
+ val originalIterator = oldChildren.iterator
+ var priorIndex = 0
+ while (priorIndex < index) {
+ newChildren(priorIndex) = originalIterator.next()
+ priorIndex += 1
+ }
+ }
+ if (!oldChild.fastEquals(newChild)) changed = true
Review Comment:
**Nit (P3):** Once `changed` is true this comparison can't affect anything: `changed` only decides between returning `this` and rebuilding, and every distinct child is stored in `newChildren` either way. The path this replaces (`withNewChildren` -> `childrenFastEquals`) returned at the first mismatch. Here, on generic n-ary nodes where several children change (e.g. resolving every `Union` branch, or casting many `Concat`/`CreateNamedStruct` arguments), each later distinct child pays a structural `==` walk it did not pay before. The arity traits do compare every child, but only because they substitute the original for equal copies; this loop keeps the copies regardless. Guarding the check restores the short-circuit with identical results:
```suggestion
if (!changed && !oldChild.fastEquals(newChild)) changed = true
```
**Verification:**
- **Behavior:** Generic mapChildren still returns the original node when every child is reference-equal or an equal-but-distinct copy, and still rebuilds with every mapped child, including retained equal copies, when any child changes materially.
- **Regression:** When two or more children of a generic n-ary node change materially, the rebuilt node holds every changed child at its original position and keeps the original node's tags and origin.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala:
##########
@@ -807,10 +825,53 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]]
* Returns a copy of this node where `f` has been applied to all the nodes in `children`.
*/
def mapChildren(f: BaseType => BaseType): BaseType = {
- if (containsChild.nonEmpty) {
- withNewChildren(children.map(f))
- } else {
+ val oldChildren = children
+ if (oldChildren.isEmpty) return this
+
+ // Invoke `f` once per child and defer allocating the replacement array until `f` returns a
+ // distinct instance. Backfill the reference-equal prefix when that first happens. `changed`
+ // records whether any change was material (`!fastEquals`). If no child changed materially we
+ // return `this` and drop the array -- an equal-but-distinct replacement alone does not require
+ // rebuilding the parent.
+ var newChildren: Array[BaseType] = null
+ var changed = false
+ var index = 0
+ val childIterator = oldChildren.iterator
+ while (childIterator.hasNext) {
+ val oldChild = childIterator.next()
+ val newChild = f(oldChild)
+ if (oldChild ne newChild) {
+ if (newChildren eq null) {
+ // First distinct replacement: allocate the result array and backfill the
+ // reference-equal prefix.
+ newChildren = new Array[TreeNode[_]](oldChildren.size).asInstanceOf[Array[BaseType]]
+ val originalIterator = oldChildren.iterator
+ var priorIndex = 0
+ while (priorIndex < index) {
+ newChildren(priorIndex) = originalIterator.next()
+ priorIndex += 1
+ }
+ }
+ if (!oldChild.fastEquals(newChild)) changed = true
+ newChildren(index) = newChild
+ } else if (newChildren ne null) {
+ newChildren(index) = oldChild
+ }
+ index += 1
+ }
+ if (!changed) {
+ // Every distinct replacement was only equal-but-distinct, so the parent is unchanged --
+ // matching `withNewChildren(children.map(f))`, which also returns `this` here.
this
+ } else {
+ // A child changed materially; rebuild. Call `withNewChildrenInternal` directly (as the
+ // arity traits do) rather than `withNewChildren`, whose size-assert and `childrenFastEquals`
+ // recheck are redundant here.
+ CurrentOrigin.withOrigin(origin) {
Review Comment:
**Nit (P3):** For generic n-ary nodes, this branch now provides the origin and tag handling that `withNewChildren` used to, but TreeNodeSuite would not notice if it regressed. 'preserves origin' uses `Add`, which goes through `BinaryLike.mapChildren`, and the new Coalesce rebuild tests assert only child identity and equality. If this `withOrigin(origin)` wrapper were dropped, a rebuilt `Coalesce` would silently take the ambient origin (e.g. the enclosing plan node's during `transformExpressions`) and lose its own SQL query context. Likewise, nothing checks the once-per-child `f` guarantee the comment above states for the rebuild path. Could one of the Coalesce rebuild tests also assert the rebuilt node's origin and the number of `f` calls?
**Verification:**
- **Regression:** A generic n-ary node rebuilt by mapChildren while an unrelated origin is current keeps its own origin, and the test fails if the rebuild picks up the ambient origin.
- **Regression:** On the rebuild path, including when a reference-equal prefix is backfilled, the mapping function is called exactly once per child in order, and the test fails on any extra call.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala:
##########
@@ -292,6 +300,23 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]]
value
}
+ // Applies `rule` to `node` with this node's origin active, restoring the previous origin
+ // afterward even if the rule throws. Equivalent to `CurrentOrigin.withOrigin`, but written as an
+ // explicit save/set/restore to avoid allocating that method's by-name closure at every node.
+ private def applyRule(
Review Comment:
**Nit (P3):** `applyRule` is only used by the transform methods (`transformDownWithReferenceEquality`, `transformDownWithPruning`, `transformUpWithPruning`, roughly 250 lines below). Placed here, it splits the cached-property/equality block between `hashCode` and `fastEquals`. Elsewhere this file keeps private helpers next to their callers, e.g. `childrenFastEquals` and `asIndexedSeq` sit right before `withNewChildren`. Could you move it into the transform section, for example just before the first transform method that calls it? This is purely a readability change: the helper that defines the transforms' origin and fallback behavior would sit next to its callers.
**Verification:**
- **Inspection:** The rule-application helper sits with the transform methods that call it rather than inside the cached-property and equality block, and its body is unchanged.
- **Compile:** The catalyst module compiles with the relocated helper.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59501][SQL] Reduce allocations in unchanged TreeNode traversals [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on PR #58800:
URL: https://github.com/apache/spark/pull/58800#issuecomment-5672778572
Reviewed as part of a batch pass (read-only). LGTM — careful, well-tested optimization.
Walked through the semantics against the originals:
- `mapChildren`: the single-pass lazy buffer reproduces `withNewChildren(children.map(f))`
exactly. It returns `this` when no child changes *materially* (matching
`childrenFastEquals`, and preserving the existing behavior of discarding
equal-but-distinct copies when nothing else changed); when a material change occurs the
buffer holds precisely what `children.map(f)` would produce (originals for `eq`
children, replacements otherwise), retaining equal-but-distinct replacements at earlier
and non-adjacent indices. The backfill runs only once (at the first distinct child), and
`children.isEmpty` is equivalent to the old `containsChild.nonEmpty` guard since
`containsChild = children.toSet`. Calling `withNewChildrenInternal` directly is safe
because the buffer size equals the child count by construction.
- `applyRule` is behaviorally identical to `CurrentOrigin.withOrigin(origin){ ... }`
(save/set/try-finally-restore), and sharing `TreeNodeIdentity.fn` is safe since identity
ignores its argument's type.
- The `transformUpWithPruning` reorder only short-circuits `mapChildren` for leaves (where
it returned `this` anyway); pruning and ineffective-rule tracking are unchanged.
Test coverage is thorough — the non-adjacent equal-copy replay case and the origin
save/restore-on-throw case are exactly the ones worth pinning. No findings.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59501][SQL] Reduce allocations in unchanged TreeNode traversals [spark]
Posted by "gengliangwang (via GitHub)" <gi...@apache.org>.
gengliangwang commented on PR #58800:
URL: https://github.com/apache/spark/pull/58800#issuecomment-6225779243
@bhollis-dbx could you resolve the conflicts?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59501][SQL] Reduce allocations in unchanged TreeNode traversals [spark]
Posted by "bhollis-dbx (via GitHub)" <gi...@apache.org>.
bhollis-dbx commented on PR #58800:
URL: https://github.com/apache/spark/pull/58800#issuecomment-5768350023
@cloud-fan I'd appreciate your eyes on this as well
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59501][SQL] Reduce allocations in unchanged TreeNode traversals [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59501][SQL] Reduce allocations in unchanged TreeNode traversals [spark]
Posted by "bhollis-dbx (via GitHub)" <gi...@apache.org>.
bhollis-dbx commented on PR #58800:
URL: https://github.com/apache/spark/pull/58800#issuecomment-5783046885
Replaced the ArrayBuffer plus toIndexedSeq with one exact-sized array wrapped without copying, avoiding the second cardinality-sized allocation.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59501][SQL] Reduce allocations in unchanged TreeNode traversals [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #58800:
URL: https://github.com/apache/spark/pull/58800#discussion_r4068454149
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/trees/TreeNodeSuite.scala:
##########
@@ -180,6 +181,76 @@ class TreeNodeSuite extends SparkFunSuite with SQLHelper {
assert(actual === expect)
}
+ test("mapChildren returns the original node when all children are fast equal") {
+ val expression = Coalesce(Seq(Literal(1), Literal(2)))
+ val visited = new ArrayBuffer[Int]()
+ val result = expression.mapChildren {
+ case literal @ Literal(value: Int, _) =>
+ visited += value
+ literal
+ case other => other
+ }
+
+ assert(result eq expression)
+ assert(visited == Seq(1, 2))
+ val leaf = Dummy(None)
+ assert(leaf.mapChildren(identity) eq leaf)
+ }
+
+ test("mapChildren returns the original node when every child is equal but a distinct copy") {
+ val expression = Coalesce(Seq(Literal(1), Literal(2)))
+ // Return a fresh, structurally-equal (not reference-equal) copy for *every* child: this
+ // fills `equalCopies` yet must still return `this`, since no child materially changes.
+ val result = expression.mapChildren {
+ case Literal(value: Int, dt) => Literal(value, dt)
+ case other => other
+ }
+ assert(result eq expression)
+ }
+
+ test("mapChildren retains an equal replacement when another child changes") {
+ val tag = TreeNodeTag[String]("equal-copy")
+ val expression = Coalesce(Seq(Literal(1), Literal(2)))
+ val equalCopy = Literal(1)
+ equalCopy.setTagValue(tag, "retained")
+
+ val result = expression.mapChildren {
+ case Literal(1, _) => equalCopy
+ case Literal(2, _) => Literal(3)
+ case other => other
+ }
+
+ assert(result.children.head eq equalCopy)
+ assert(result.children.head.getTagValue(tag).contains("retained"))
+ assert(result.children(1) == Literal(3))
+ }
+
+ test("mapChildren retains non-adjacent equal replacements when a later child changes") {
+ val c0 = Literal(10)
+ val c1 = Literal(11)
+ val c2 = Literal(12)
+ val c3 = Literal(13)
+ val expression = Coalesce(Seq(c0, c1, c2, c3))
+ val copy0 = Literal(10)
+ val copy2 = Literal(12)
+
+ // Equal-but-distinct copies at non-adjacent indices 0 and 2, an unchanged same instance at 1,
+ // and a material change at 3. Exercises the replay loop's index bookkeeping across a gap.
Review Comment:
**Non-blocking (P2):** `copy0` makes the first distinct result occur at index 0, so the `priorIndex < index` prefix-backfill loop is never exercised. Please keep at least one reference-equal child before the first distinct replacement and assert the full resulting order/identity, while retaining a later material change, so regressions in this new bookkeeping have a failure signal.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala:
##########
@@ -559,15 +580,12 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]]
if (!cond.apply(this) || isRuleIneffective(ruleId)) {
return this
}
- val afterRuleOnChildren = mapChildren(_.transformUpWithPruning(cond, ruleId)(rule))
- val newNode = if (this fastEquals afterRuleOnChildren) {
- CurrentOrigin.withOrigin(origin) {
- rule.applyOrElse(this, identity[BaseType])
- }
+ val newNode = if (children.isEmpty) {
Review Comment:
**Non-blocking (P2):** `children.isEmpty` runs before virtual `mapChildren` dispatch, so a fresh specialized non-leaf initializes its lazy child `IndexedSeq` even though the arity override reads the direct child fields. The previous transform-up path avoided that per-node allocation. Please gate the leaf fast path with the existing `LeafLike` marker and keep the `mapChildren` fallback for other nodes, with coverage that detects premature child materialization.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/trees/TreeNodeSuite.scala:
##########
@@ -193,6 +264,58 @@ class TreeNodeSuite extends SparkFunSuite with SQLHelper {
assert(transformed.origin.startPosition.isDefined)
}
+ test("transform rules see node origins and restore the previous origin") {
+ val nodeOrigin = Origin(line = Some(1))
+ val previousOrigin = Origin(line = Some(2))
+ val expression = CurrentOrigin.withOrigin(nodeOrigin) {
Review Comment:
**Non-blocking (P2):** The parent and both children are constructed under the same origin, so `CurrentOrigin.get == e.origin` also passes if the parent origin remains installed for child callbacks. Please construct the parent and children under distinct origins and retain the success/exception restoration assertions, so this test fails on an ancestor-origin leak.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala:
##########
@@ -745,10 +763,52 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]]
* Returns a copy of this node where `f` has been applied to all the nodes in `children`.
*/
def mapChildren(f: BaseType => BaseType): BaseType = {
- if (containsChild.nonEmpty) {
- withNewChildren(children.map(f))
- } else {
+ val oldChildren = children
+ if (oldChildren.isEmpty) return this
+
+ // Single pass, allocating nothing until `f` returns a distinct instance for some child.
Review Comment:
**Nit (P3):** This is not literally allocation-free until the first distinct child or a single traversal: `childIterator` is created up front, and a late first distinct child makes `originalIterator` walk the prefix again. Please describe the guarantees the code does provide—one `f` invocation per child and deferred replacement-buffer allocation with prefix backfill.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/trees/TreeNodeSuite.scala:
##########
@@ -180,6 +181,76 @@ class TreeNodeSuite extends SparkFunSuite with SQLHelper {
assert(actual === expect)
}
+ test("mapChildren returns the original node when all children are fast equal") {
+ val expression = Coalesce(Seq(Literal(1), Literal(2)))
+ val visited = new ArrayBuffer[Int]()
+ val result = expression.mapChildren {
+ case literal @ Literal(value: Int, _) =>
+ visited += value
+ literal
+ case other => other
+ }
+
+ assert(result eq expression)
+ assert(visited == Seq(1, 2))
+ val leaf = Dummy(None)
+ assert(leaf.mapChildren(identity) eq leaf)
+ }
+
+ test("mapChildren returns the original node when every child is equal but a distinct copy") {
+ val expression = Coalesce(Seq(Literal(1), Literal(2)))
+ // Return a fresh, structurally-equal (not reference-equal) copy for *every* child: this
+ // fills `equalCopies` yet must still return `this`, since no child materially changes.
Review Comment:
**Nit (P3):** There is no `equalCopies` structure in this implementation; distinct mapped children are accumulated in `newChildren`. Please describe the equal-but-distinct behavior directly, or use the current name, so the test explanation matches the code it protects.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59501][SQL] Reduce allocations in unchanged TreeNode traversals [spark]
Posted by "bhollis-dbx (via GitHub)" <gi...@apache.org>.
bhollis-dbx commented on code in PR #58800:
URL: https://github.com/apache/spark/pull/58800#discussion_r4075557757
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala:
##########
@@ -745,10 +763,52 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]]
* Returns a copy of this node where `f` has been applied to all the nodes in `children`.
*/
def mapChildren(f: BaseType => BaseType): BaseType = {
- if (containsChild.nonEmpty) {
- withNewChildren(children.map(f))
- } else {
+ val oldChildren = children
+ if (oldChildren.isEmpty) return this
+
+ // Single pass, allocating nothing until `f` returns a distinct instance for some child.
Review Comment:
Updated the comment to state one f call per child and deferred buffer allocation with prefix backfill.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59501][SQL] Reduce allocations in unchanged TreeNode traversals [spark]
Posted by "bhollis-dbx (via GitHub)" <gi...@apache.org>.
bhollis-dbx commented on code in PR #58800:
URL: https://github.com/apache/spark/pull/58800#discussion_r4075562978
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/trees/TreeNodeSuite.scala:
##########
@@ -193,6 +264,58 @@ class TreeNodeSuite extends SparkFunSuite with SQLHelper {
assert(transformed.origin.startPosition.isDefined)
}
+ test("transform rules see node origins and restore the previous origin") {
+ val nodeOrigin = Origin(line = Some(1))
+ val previousOrigin = Origin(line = Some(2))
+ val expression = CurrentOrigin.withOrigin(nodeOrigin) {
Review Comment:
Updated the test so the parent and children have distinct origins.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org