optimize Subquery

Author: zml1206

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

@@ -1820,8 +1820,7 @@ object PushPredicateThroughNonJoin extends Rule[LogicalPlan] with PredicateHelpe
     case Filter(condition, project @ Project(fields, grandChild))
       if fields.forall(_.deterministic) && canPushThroughCondition(grandChild, condition) =>
       val aliasMap = getAliasMap(project)
-      val replacedByWith = rewriteConditionByWith(condition, aliasMap)
-      project.copy(child = Filter(replaceAlias(replacedByWith, aliasMap), grandChild))
+      project.copy(child = Filter(rewriteCondition(condition, aliasMap), grandChild))
 
     // We can push down deterministic predicate through Aggregate, including throwable predicate.
     // If we can push down a filter through Aggregate, it means the filter only references the
@@ -1841,8 +1840,7 @@ object PushPredicateThroughNonJoin extends Rule[LogicalPlan] with PredicateHelpe
       }
 
       if (pushDown.nonEmpty) {
-        val replacedByWith = rewriteConditionByWith(pushDown, aliasMap)
-        val replaced = replaceAlias(replacedByWith.reduce(And), aliasMap)
+        val replaced = rewriteCondition(pushDown.reduce(And), aliasMap)
         val newAggregate = aggregate.copy(child = Filter(replaced, aggregate.child))
         // If there is no more filter to stay up, just eliminate the filter.
         // Otherwise, create "Filter(stayUp) <- Aggregate <- Filter(pushDownPredicate)".
@@ -1989,51 +1987,53 @@ object PushPredicateThroughNonJoin extends Rule[LogicalPlan] with PredicateHelpe
     }
   }
 
+  private def rewriteCondition(
+      cond: Expression,
+      aliasMap: AttributeMap[Alias]): Expression = {
+    replaceAlias(rewriteConditionByWith(cond, aliasMap), aliasMap)
+  }
+
   /**
    * Use [[With]] to rewrite condition which contains attribute that are not cheap and be consumed
    * multiple times. Each predicate generates one or 0 With. For facilitates subsequent merge
    * [[With]], use the same CommonExpressionDef ids for different [[With]].
    */
   private def rewriteConditionByWith(
-      cond: Seq[Expression],
-      aliasMap: AttributeMap[Alias]): Seq[Expression] = {
+      cond: Expression,
+      aliasMap: AttributeMap[Alias]): Expression = {
     if (!SQLConf.get.getConf(SQLConf.ALWAYS_INLINE_COMMON_EXPR)) {
-      val canRewriteCond = cond.filter(canRewriteByWith)
-      if (canRewriteCond.nonEmpty) {
-        val replaceWithMap = canRewriteCond
-          .flatMap(_.collect {case a: Attribute => a })
-          .groupBy(identity)
-          .transform((_, v) => v.size)
-          .filter(m => aliasMap.contains(m._1) && m._2 > 1)
-          .map(m => m._1 -> trimAliases(aliasMap.getOrElse(m._1, m._1)))
-          .filter(m => !CollapseProject.isCheap(m._2))
+      // SubqueryExpression can't contain common expression ref, replace alias for it first.
+      val newCond = replaceAliasForSubqueryExpression(cond, aliasMap)
+      val replaceWithMap = newCond.collect {case a: Attribute => a }
+        .groupBy(identity)
+        .transform((_, v) => v.size)
+        .filter(m => aliasMap.contains(m._1) && m._2 > 1)
+        .map(m => m._1 -> trimAliases(aliasMap.getOrElse(m._1, m._1)))
+        .filter(m => !CollapseProject.isCheap(m._2))
+      if (replaceWithMap.isEmpty) {
+        newCond
+      } else {
         val defsMap = AttributeMap(replaceWithMap.map(m => m._1 -> CommonExpressionDef(m._2)))
         val refsMap = AttributeMap(defsMap.map(m => m._1 -> new CommonExpressionRef(m._2)))
-        cond.map(rewriteByWith(_, defsMap, refsMap))
-      } else cond
+        splitConjunctivePredicates(newCond)
+          .map(rewriteByWith(_, defsMap, refsMap))
+          .reduce(And)
+      }
     } else cond
   }
 
-  private def rewriteConditionByWith(
-      cond: Expression,
+  private def replaceAliasForSubqueryExpression(
+      expr: Expression,
       aliasMap: AttributeMap[Alias]): Expression = {
-    if (!SQLConf.get.getConf(SQLConf.ALWAYS_INLINE_COMMON_EXPR)) {
-      rewriteConditionByWith(splitConjunctivePredicates(cond), aliasMap).reduce(And)
-    } else cond
-  }
-
-  // With does not support inline subquery
-  private def canRewriteByWith(expr: Expression): Boolean = {
-    !expr.containsPattern(PLAN_EXPRESSION)
+    expr.transform {
+      case s: SubqueryExpression => replaceAlias(s, aliasMap)
+    }
   }
 
   private def rewriteByWith(
       expr: Expression,
       defsMap: AttributeMap[CommonExpressionDef],
       refsMap: AttributeMap[CommonExpressionRef]): Expression = {
-    if (!canRewriteByWith(expr)) {
-      return expr
-    }
     val defs = mutable.HashSet.empty[CommonExpressionDef]
     val replaced = expr.transform {
       case a: Attribute if refsMap.contains(a) =>