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

["stevomitric (via GitHub)" <gi...@apache.org>]

stevomitric commented on code in PR #58348:
URL: https://github.com/apache/spark/pull/58348#discussion_r3951718327


##########
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))

Review Comment:
   consider inlining:
   ```scala
    case pe: PlanExpression[LogicalPlan @unchecked] =>
      pe.withNewPlan(pe.plan.transformAllExpressionsWithSubqueries {
        case OuterReference(a: Attribute) => OuterReference(rewrites.getOrElse(a, a))
      })
   ```



-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org