[SPARK-37448] Multiple optimizations related to CurrentOrigin.withOrigin

• Type: Improvement
• Status: In Progress
• Priority: Major
• Resolution: Unresolved
• Affects Version/s: 4.0.0
• Component/s: Optimizer
• Labels: pull-request-available
• Assignee: Josh Rosen
• Reporter: Josh Rosen
• Created: 23/Nov/21 08:42
• Updated: 07/Jun/24 09:33

Description

This PR / Jira proposes multiple optimizations related to CurrentOrigin.withOrigin, including:

Avoid CurrentOrigin.withOrigin calls in tree transformation methods when partial function isn't defined

`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.

Small optimizations to withOrigin itself to save a thread-local write.

See PR for full details.

Issue Links

• links to: GitHub Pull Request #46908