[SPARK-59042][SQL] Fix PushProjectionThroughUnion failure on correlated scalar subquery over UNION ALL

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala

@@ -1114,14 +1114,29 @@ object PushProjectionThroughUnion extends Rule[LogicalPlan] {
     AttributeMap(left.output.zip(right.output))
   }
 
+  private def updateOuterReferencesInSubquery(
+      plan: LogicalPlan,
+      rewrites: AttributeMap[Attribute]): LogicalPlan = {
+    plan.transformDown { case currentFragment =>
+      currentFragment.transformExpressions {
+        case OuterReference(a: Attribute) =>
+          OuterReference(rewrites.getOrElse(a, a))
+        case pe: PlanExpression[LogicalPlan @unchecked] =>
+          pe.withNewPlan(updateOuterReferencesInSubquery(pe.plan, rewrites))
+      }
+    }
+  }
+
   /**
    * Rewrites an expression so that it can be pushed to the right side of a
    * Union or Except operator. This method relies on the fact that the output attributes
    * of a union/intersect/except are always equal to the left child's output.
    */
   private def pushToRight[A <: Expression](e: A, rewrites: AttributeMap[Attribute]) = {
     val result = e transform {
-      case a: Attribute => rewrites(a)
+      case a: Attribute => rewrites.getOrElse(a, a)
+      case pe: PlanExpression[LogicalPlan @unchecked] =>
+        pe.withNewPlan(updateOuterReferencesInSubquery(pe.plan, rewrites))
     } match {
       // Make sure exprId is unique in each child of Union.
       case Alias(child, alias) => Alias(child, alias)()

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PushProjectThroughUnionSuite.scala

@@ -19,6 +19,7 @@ package org.apache.spark.sql.catalyst.optimizer
 
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.dsl.plans._
+import org.apache.spark.sql.catalyst.expressions.ScalarSubquery
 import org.apache.spark.sql.catalyst.plans.PlanTest
 import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan}
 import org.apache.spark.sql.catalyst.rules.RuleExecutor
@@ -51,4 +52,27 @@ class PushProjectThroughUnionSuite extends PlanTest {
 
     comparePlans(optimized, expected)
   }
+
+  test("SPARK-59042: PushProjectionThroughUnion handles subquery attributes cleanly") {
+    val testRelation1 = LocalRelation($"a".int)
+    val testRelation2 = LocalRelation($"d".int)
+    val subqueryRelation = LocalRelation($"x".int)
+    val subquery = ScalarSubquery(subqueryRelation.where($"x" === $"a").select($"x"))
+
+    val query = testRelation1
+      .union(testRelation2)
+      .select($"a", subquery.as("sub"))
+      .analyze
+    val optimized = Optimize.execute(query)
+
+    val expectedChild2Sub = ScalarSubquery(subqueryRelation.where($"x" === $"d").select($"x"))
+
+    val expected = testRelation1
+      .select($"a", subquery.as("sub"))
+      .union(testRelation2
+        .select($"d", expectedChild2Sub.as("sub")))
+      .analyze
+
+    comparePlans(optimized, expected)
+  }
 }

sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala

@@ -2678,4 +2678,22 @@ class SubquerySuite extends SharedSparkSession
 
     assert(exposedAttribute.exprId == outerReferenceAttribute.exprId)
   }
+
+  test("SPARK-59042: PushProjectionThroughUnion fails on a correlated scalar subquery " +
+    "over UNION ALL") {
+    val df = sql(
+      """
+        |SELECT u.a,
+        |       (SELECT max(r.x)
+        |        FROM (VALUES (1), (2), (3), (NULL)) AS r(x)
+        |        WHERE r.x = u.a) AS m
+        |FROM (
+        |  SELECT a FROM (VALUES (1), (2)) AS l(a)
+        |  UNION ALL
+        |  SELECT a FROM (VALUES (2), (3)) AS q(a)
+        |) u
+        |ORDER BY u.a, m
+        |""".stripMargin)
+    checkAnswer(df, Row(1, 1) :: Row(2, 2) :: Row(2, 2) :: Row(3, 3) :: Nil)
+  }
 }