[SPARK-37448] Avoid CurrentOrigin.withOrigin calls in tree transformation methods when partial function isn't defined

• Type: Improvement
• Status: In Progress
• Priority: Major
• Resolution: Unresolved
• Affects Version/s: 3.2.0
• Component/s: Optimizer
• Assignee: Josh Rosen
• Reporter: Josh Rosen
• Created: 23/Nov/21 08:42
• Updated: 23/Nov/21 08:50

Description

`CurrentOrigin.withOrigin { block }` performs a ThreadLocal.get() call to retrieve the current value of the `origin` thread local so it can be restored after `block` exits.

TreeNode has transform methods which apply PartialFunctions to every node in the tree. By checking `rule.isDefinedAt` before applying a rule we can skip the `CurrentOrigin.withOrigin` call in cases where we know the partial function isn't defined. This avoids the performance overhead associated with thread-local reads and sets.