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