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