zahed1994 opened a new pull request, #58348:
URL: https://github.com/apache/spark/pull/58348

   ### What changes were proposed in this pull request?
   This PR fixes a `java.util.NoSuchElementException` when 
`PushProjectionThroughUnion` optimizes a `Project` over a `Union` where project 
expressions contain subqueries or outer references.
   
   Specifically:
   - Updated `PushProjectionThroughUnion.pushToRight` to use 
`rewrites.getOrElse(a, a)` instead of `rewrites(a)` when looking up attribute 
references. This prevents `NoSuchElementException` when encountering 
subquery-internal or unmapped attributes.
   - Added `updateOuterReferencesInSubquery` helper in 
`PushProjectionThroughUnion` to recurse into subquery plans 
(`PlanExpression.plan`) and rewrite inner `OuterReference(a)` references to 
match the right child's corresponding attributes.
   
   ### Why are the changes needed?
   When a query contains a correlated scalar subquery (or any expression 
referencing subquery-internal attributes) over a `UNION ALL`, 
`PushProjectionThroughUnion` attempts to push project expressions to the 
children of the `Union`. 
   
   Prior to this fix:
   1. `pushToRight` called `rewrites(a)` directly on all attributes. 
Subquery-internal attributes (e.g. `r.x#7` in `WHERE r.x = u.a`) were absent 
from `rewrites` (which maps `left.output -> right.output`), throwing 
`java.util.NoSuchElementException: key not found: x#7`.
   2. Outer references inside subquery plans (`OuterReference(a)`) were skipped 
by `e transform` because `OuterReference` extends `LeafExpression` and 
`PlanExpression.plan` is a `LogicalPlan` rather than an `Expression` child.
   
   ### Does this PR introduce _any_ user-facing change?
   No. Fixes an optimization failure exception for correlated subqueries over 
`UNION ALL`.
   
   ### How was this patch tested?
   - Added catalyst optimizer unit test in `PushProjectThroughUnionSuite`.
   - Added end-to-end SQL query test in `SubquerySuite`.
   
   ### Was this patch authored or co-authored using generative AI tooling?
   No.
   


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

Reply via email to