Skip to content

[SPARK-56869][SQL] Speed up TreeNode transforms when rule doesn't match #55889

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
gengliangwang wants to merge 1 commit into from
Closed

[SPARK-56869][SQL] Speed up TreeNode transforms when rule doesn't match #55889

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
70 changes: 47 additions & 23 deletions sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -491,23 +491,36 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]]
if (!cond.apply(this) || isRuleIneffective(ruleId)) {
return this
}
val afterRule = CurrentOrigin.withOrigin(origin) {
rule.applyOrElse(this, identity[BaseType])
}

// Check if unchanged and then possibly return old copy to avoid gc churn.
if (this fastEquals afterRule) {
// CurrentOrigin.withOrigin is only observable when the rule constructs new nodes
// (they pick up origin via `override val origin = CurrentOrigin.get`). When the rule
// doesn't fire, the wrapping is pure ThreadLocal thrash — profiling shows ~88% of
// transform CPU goes to ThreadLocal.set / get on no-match nodes. Skip it.
if (rule.isDefinedAt(this)) {
val afterRule = CurrentOrigin.withOrigin(origin) {
rule.apply(this)
}
// Check if unchanged and then possibly return old copy to avoid gc churn.
if (this fastEquals afterRule) {
val rewritten_plan = mapChildren(_.transformDownWithPruning(cond, ruleId)(rule))
if (this eq rewritten_plan) {
markRuleAsIneffective(ruleId)
this
} else {
rewritten_plan
}
} else {
// If the transform function replaces this node with a new one, carry over the tags.
afterRule.copyTagsFrom(this)
afterRule.mapChildren(_.transformDownWithPruning(cond, ruleId)(rule))
}
} else {
val rewritten_plan = mapChildren(_.transformDownWithPruning(cond, ruleId)(rule))
if (this eq rewritten_plan) {
markRuleAsIneffective(ruleId)
this
} else {
rewritten_plan
}
} else {
// If the transform function replaces this node with a new one, carry over the tags.
afterRule.copyTagsFrom(this)
afterRule.mapChildren(_.transformDownWithPruning(cond, ruleId)(rule))
}
}

Expand Down Expand Up @@ -544,14 +557,15 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]]
return this
}
val afterRuleOnChildren = mapChildren(_.transformUpWithPruning(cond, ruleId)(rule))
val newNode = if (this fastEquals afterRuleOnChildren) {
// Skip the CurrentOrigin.withOrigin wrap when the rule doesn't match — see
// transformDownWithPruning above for the same optimization.
val target = if (this fastEquals afterRuleOnChildren) this else afterRuleOnChildren
val newNode = if (rule.isDefinedAt(target)) {
CurrentOrigin.withOrigin(origin) {
rule.applyOrElse(this, identity[BaseType])
rule.apply(target)
}
} else {
CurrentOrigin.withOrigin(origin) {
rule.applyOrElse(afterRuleOnChildren, identity[BaseType])
}
target
}
if (this eq newNode) {
markRuleAsIneffective(ruleId)
Expand Down Expand Up @@ -588,8 +602,15 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]]
}
val afterRuleOnChildren =
mapChildren(_.transformUpWithBeforeAndAfterRuleOnChildren(cond, ruleId)(rule))
val newNode = CurrentOrigin.withOrigin(origin) {
rule.applyOrElse((this, afterRuleOnChildren), { t: (BaseType, BaseType) => t._2 })
// Skip the CurrentOrigin.withOrigin wrap when the rule doesn't match — see
// transformDownWithPruning above. The default would return afterRuleOnChildren anyway.
val key = (this, afterRuleOnChildren)
val newNode = if (rule.isDefinedAt(key)) {
CurrentOrigin.withOrigin(origin) {
rule.apply(key)
}
} else {
afterRuleOnChildren
}
if (this eq newNode) {
this.markRuleAsIneffective(ruleId)
Expand Down Expand Up @@ -685,12 +706,15 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]]
// alternatives. I.e. the "multiTransformDown is lazy" test case in `TreeNodeSuite` would fail.
// Please note that this behaviour has a downside as well that we can only mark the rule on the
// original node ineffective if the rule didn't match.
var ruleApplied = true
val afterRules = CurrentOrigin.withOrigin(origin) {
rule.applyOrElse(this, (_: BaseType) => {
ruleApplied = false
Seq.empty
})
// Skip the CurrentOrigin.withOrigin wrap when the rule doesn't match — see
// transformDownWithPruning above. Also lets us drop the side-effecting default.
val ruleApplied = rule.isDefinedAt(this)
val afterRules: Seq[BaseType] = if (ruleApplied) {
CurrentOrigin.withOrigin(origin) {
rule.apply(this)
}
} else {
Seq.empty
}

val afterRulesLazyList = if (afterRules.isEmpty) {
Expand Down