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: [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