address comments

allisonwang-db · allisonwang-db · commit 7ba19740668c · 2021-07-19T14:05:47.000-07:00
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala
@@ -437,24 +437,19 @@ abstract class QueryPlan[PlanType <: QueryPlan[PlanType]]
 
   /**
    * Returns a copy of this node where the given partial function has been recursively applied
-   * first to this node's children, then this node's subqueries, and finally this node itself
-   * (post-order). When the partial function does not apply to a given node, it is left unchanged.
+   * first to the subqueries in this node's children, then this node's children, and finally
+   * this node itself (post-order). When the partial function does not apply to a given node,
+   * it is left unchanged.
    */
   def transformUpWithSubqueries(f: PartialFunction[PlanType, PlanType]): PlanType = {
-    val g: PartialFunction[PlanType, PlanType] = new PartialFunction[PlanType, PlanType] {
-      override def isDefinedAt(x: PlanType): Boolean = true
-
-      override def apply(plan: PlanType): PlanType = {
-        val transformed = plan transformExpressionsUp {
-          case planExpression: PlanExpression[PlanType] =>
-            val newPlan = planExpression.plan.transformUpWithSubqueries(f)
-            planExpression.withNewPlan(newPlan)
-        }
-        f.applyOrElse[PlanType, PlanType](transformed, identity)
+    transformUp { case plan =>
+      val transformed = plan transformExpressionsUp {
+        case planExpression: PlanExpression[PlanType] =>
+          val newPlan = planExpression.plan.transformUpWithSubqueries(f)
+          planExpression.withNewPlan(newPlan)
       }
+      f.applyOrElse[PlanType, PlanType](transformed, identity)
     }
-
-    transformUp(g)
   }
 
   /**
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeOneRowRelationSubquerySuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeOneRowRelationSubquerySuite.scala
@@ -17,11 +17,12 @@
 
 package org.apache.spark.sql.catalyst.optimizer
 
+import org.apache.spark.sql.catalyst.analysis.CleanupAliases
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.dsl.plans._
-import org.apache.spark.sql.catalyst.expressions.{Alias, ScalarSubquery}
+import org.apache.spark.sql.catalyst.expressions.ScalarSubquery
 import org.apache.spark.sql.catalyst.plans._
-import org.apache.spark.sql.catalyst.plans.logical.{DomainJoin, LocalRelation, LogicalPlan, OneRowRelation, Project}
+import org.apache.spark.sql.catalyst.plans.logical.{DomainJoin, LocalRelation, LogicalPlan, OneRowRelation}
 import org.apache.spark.sql.catalyst.rules.RuleExecutor
 import org.apache.spark.sql.internal.SQLConf
 
@@ -46,7 +47,9 @@ class OptimizeOneRowRelationSubquerySuite extends PlanTest {
     val batches =
       Batch("Subquery", Once,
         OptimizeOneRowRelationSubquery,
-        PullupCorrelatedPredicates) :: Nil
+        PullupCorrelatedPredicates) ::
+      Batch("Cleanup", FixedPoint(10),
+        CleanupAliases) :: Nil
   }
 
   private def assertHasDomainJoin(plan: LogicalPlan): Unit = {
@@ -91,8 +94,8 @@ class OptimizeOneRowRelationSubquerySuite extends PlanTest {
     val inner = t0.select('a.as("a1"), 'b.as("b1")).select(('a1 + 'b1).as("c"))
     val query = t1.select(ScalarSubquery(inner).as("sub"))
     val optimized = Optimize.execute(query.analyze)
-    val correctAnswer = Project(Alias(Alias(a + b, "c")(), "sub")() :: Nil, t1)
-    comparePlans(optimized, correctAnswer)
+    val correctAnswer = t1.select(('a + 'b).as("c").as("sub"))
+    comparePlans(optimized, correctAnswer.analyze)
   }
 
   test("Optimize lateral subquery with multiple projects") {
@@ -111,8 +114,8 @@ class OptimizeOneRowRelationSubquerySuite extends PlanTest {
     val inner = t0.select('a.as("b")).select(ScalarSubquery(t0.select('b)).as("s"))
     val query = t1.select(ScalarSubquery(inner).as("sub"))
     val optimized = Optimize.execute(query.analyze)
-    val correctAnswer = Project(Alias(Alias(a, "s")(), "sub")() :: Nil, t1)
-    comparePlans(optimized, correctAnswer)
+    val correctAnswer = t1.select('a.as("s").as("sub"))
+    comparePlans(optimized, correctAnswer.analyze)
   }
 
   test("Batch should be idempotent") {
@@ -149,8 +152,9 @@ class OptimizeOneRowRelationSubquerySuite extends PlanTest {
     }
   }
 
-  test("Should not optimize subquery with nested subqueries") {
+  test("Should not optimize subquery with nested subqueries that can't be optimized") {
     // SELECT (SELECT (SELECT a WHERE a = 1) FROM (SELECT a AS a)) FROM t1
+    // Filter (a = 1) cannot be optimized.
     val inner = t0.select('a).where('a === 1)
     val subquery = t0.select('a.as("a"))
       .select(ScalarSubquery(inner).as("s")).select('s + 1)
