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]
