Move PullupCorrelatedPredicates and RewritePredicateSubquery after OptimizeSubqueries

This commit moves two rules right next to the rule OptimizeSubqueries.
 1. PullupCorrelatedPredicates:
    the rewrite of [Not] Exists and [Not] In (ListQuery) to PredicateSubquery
 2. RewritePredicateSubquery:
    the rewrite of PredicateSubquery to LeftSemi/LeftAnti

With this change, [Not] Exists/In subquery is now rewritten to LeftSemi/LeftAnti
at the beginning of Optimizer.

By moving rule PullupCorrelatedPredicates after rule OptimizerSubqueries, all
the rules from the nested call to the entire Optimizer on the plans in subqueries
will need to deal with (1). the correlated columns wrapped with OuterReference,
and (2) the SubqueryExpression.

We will block any push down of both types of expressions for the following reasons:

1. We do not want to push any correlated expressions further down the plan tree.
   Deep correlation is not yet supported in Spark, and, even when supported,
   deep correlation is more difficult to be unnested to a join.
2. We do not want to push any correlated subquery down because the correlated
   columns' ExprIds in the subquery may need to remap to different ExprIds from
   the plan below the current Filter that hosts the subquery.

Another side effect is we used to push down Exists/In subquery as if it is a
predicate in rule PushDownPredicate and rule PushPredicateThroughJoin. Now
Exists/In subquery is rewritten to LeftSemi/LeftAnti, we need to handle
the push down of LeftSemi/LeftAnti instead. This will be done in a followup
commit.

Another Todo is to merge the two-stage rewrite in rule PullupCorrelatedPredicates
and rule RewritePredicateSubquery into a single stage rewrite.

Author: nsyca

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

@@ -68,10 +68,10 @@ abstract class Optimizer(sessionCatalog: SessionCatalog, conf: CatalystConf)
     //   since the other rules might make two separate Unions operators adjacent.
     Batch("Union", Once,
       CombineUnions) ::
-    Batch("Pullup Correlated Expressions", Once,
-      PullupCorrelatedPredicates) ::
     Batch("Subquery", Once,
-      OptimizeSubqueries) ::
+      OptimizeSubqueries,
+      PullupCorrelatedPredicates,
+      RewritePredicateSubquery) ::
     Batch("Replace Operators", fixedPoint,
       ReplaceIntersectWithSemiJoin,
       ReplaceExceptWithAntiJoin,
@@ -130,10 +130,7 @@ abstract class Optimizer(sessionCatalog: SessionCatalog, conf: CatalystConf)
       ConvertToLocalRelation,
       PropagateEmptyRelation) ::
     Batch("OptimizeCodegen", Once,
-      OptimizeCodegen(conf)) ::
-    Batch("RewriteSubquery", Once,
-      RewritePredicateSubquery,
-      CollapseProject) :: Nil
+      OptimizeCodegen(conf)) :: Nil
   }
 
   /**
@@ -746,7 +743,9 @@ object PushDownPredicate extends Rule[LogicalPlan] with PredicateHelper {
     // state and all the input rows processed before. In another word, the order of input rows
     // matters for non-deterministic expressions, while pushing down predicates changes the order.
     case filter @ Filter(condition, project @ Project(fields, grandChild))
-      if fields.forall(_.deterministic) && canPushThroughCondition(grandChild, condition) =>
+      if fields.forall(_.deterministic) &&
+        !SubqueryExpression.hasCorrelatedSubquery(condition) &&
+        !SubExprUtils.containsOuter(condition) =>
 
       // Create a map of Aliases to their values from the child projection.
       // e.g., 'SELECT a + b AS c, d ...' produces Map(c -> a + b).
@@ -769,7 +768,9 @@ object PushDownPredicate extends Rule[LogicalPlan] with PredicateHelper {
         splitConjunctivePredicates(condition).span(_.deterministic)
 
       val (pushDown, rest) = candidates.partition { cond =>
-        cond.references.subsetOf(partitionAttrs)
+        cond.references.subsetOf(partitionAttrs) &&
+        !SubqueryExpression.hasCorrelatedSubquery(cond) &&
+        !SubExprUtils.containsOuter(cond)
       }
 
       val stayUp = rest ++ containingNonDeterministic
@@ -797,7 +798,9 @@ object PushDownPredicate extends Rule[LogicalPlan] with PredicateHelper {
 
       val (pushDown, rest) = candidates.partition { cond =>
         val replaced = replaceAlias(cond, aliasMap)
-        cond.references.nonEmpty && replaced.references.subsetOf(aggregate.child.outputSet)
+        cond.references.nonEmpty && replaced.references.subsetOf(aggregate.child.outputSet) &&
+        !SubqueryExpression.hasCorrelatedSubquery(cond) &&
+        !SubExprUtils.containsOuter(cond)
       }
 
       val stayUp = rest ++ containingNonDeterministic
@@ -815,7 +818,14 @@ object PushDownPredicate extends Rule[LogicalPlan] with PredicateHelper {
 
     case filter @ Filter(condition, union: Union) =>
       // Union could change the rows, so non-deterministic predicate can't be pushed down
-      val (pushDown, stayUp) = splitConjunctivePredicates(condition).span(_.deterministic)
+      val (candidates, containingNonDeterministic) =
+        splitConjunctivePredicates(condition).span(_.deterministic)
+
+      val (pushDown, rest) = candidates.partition { cond =>
+        !SubqueryExpression.hasCorrelatedSubquery(cond) &&
+        !SubExprUtils.containsOuter(cond)
+      }
+      val stayUp = rest ++ containingNonDeterministic
 
       if (pushDown.nonEmpty) {
         val pushDownCond = pushDown.reduceLeft(And)
@@ -839,7 +849,9 @@ object PushDownPredicate extends Rule[LogicalPlan] with PredicateHelper {
       }
 
     case filter @ Filter(condition, u: UnaryNode)
-        if canPushThrough(u) && u.expressions.forall(_.deterministic) =>
+        if canPushThrough(u) && u.expressions.forall(_.deterministic) &&
+          !SubqueryExpression.hasCorrelatedSubquery(condition) &&
+          !SubExprUtils.containsOuter(condition) =>
       pushDownPredicate(filter, u.child) { predicate =>
         u.withNewChildren(Seq(Filter(predicate, u.child)))
       }
@@ -887,20 +899,6 @@ object PushDownPredicate extends Rule[LogicalPlan] with PredicateHelper {
       filter
     }
   }
-
-  /**
-   * Check if we can safely push a filter through a projection, by making sure that predicate
-   * subqueries in the condition do not contain the same attributes as the plan they are moved
-   * into. This can happen when the plan and predicate subquery have the same source.
-   */
-  private def canPushThroughCondition(plan: LogicalPlan, condition: Expression): Boolean = {
-    val attributes = plan.outputSet
-    val matched = condition.find {
-      case s: SubqueryExpression => s.plan.outputSet.intersect(attributes).nonEmpty
-      case _ => false
-    }
-    matched.isEmpty
-  }
 }
 
 /**
@@ -927,13 +925,18 @@ object PushPredicateThroughJoin extends Rule[LogicalPlan] with PredicateHelper {
     // any deterministic expression that follows a non-deterministic expression. To achieve this,
     // we only consider pushing down those expressions that precede the first non-deterministic
     // expression in the condition.
-    val (pushDownCandidates, containingNonDeterministic) = condition.span(_.deterministic)
+    val (candidates, containingNonDeterministic) = condition.span(_.deterministic)
+    val (pushDownCandidates, subquery) = candidates.partition { cond =>
+        !SubqueryExpression.hasCorrelatedSubquery(cond) &&
+        !SubExprUtils.containsOuter(cond)
+    }
     val (leftEvaluateCondition, rest) =
       pushDownCandidates.partition(_.references.subsetOf(left.outputSet))
     val (rightEvaluateCondition, commonCondition) =
         rest.partition(expr => expr.references.subsetOf(right.outputSet))
 
-    (leftEvaluateCondition, rightEvaluateCondition, commonCondition ++ containingNonDeterministic)
+    (leftEvaluateCondition, rightEvaluateCondition,
+      subquery ++ commonCondition ++ containingNonDeterministic)
   }
 
   def apply(plan: LogicalPlan): LogicalPlan = plan transform {

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

@@ -445,7 +445,11 @@ object EliminateOuterJoin extends Rule[LogicalPlan] with PredicateHelper {
    * Returns whether the expression returns null or false when all inputs are nulls.
    */
   private def canFilterOutNull(e: Expression): Boolean = {
-    if (!e.deterministic || SubqueryExpression.hasCorrelatedSubquery(e)) return false
+    if (!e.deterministic ||
+        SubqueryExpression.hasCorrelatedSubquery(e) ||
+        SubExprUtils.containsOuter(e)) {
+      return false
+    }
     val attributes = e.references.toSeq
     val emptyRow = new GenericInternalRow(attributes.length)
     val boundE = BindReferences.bindReference(e, attributes)

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

@@ -799,7 +799,7 @@ class FilterPushdownSuite extends PlanTest {
     comparePlans(optimized, correctedAnswer)
   }
 
-  test("predicate subquery: push down simple") {
+  test("correlated subquery (simple): no push down") {
     val x = testRelation.subquery('x)
     val y = testRelation.subquery('y)
     val z = LocalRelation('a.int, 'b.int, 'c.int).subquery('z)
@@ -809,14 +809,14 @@ class FilterPushdownSuite extends PlanTest {
       .where(Exists(z.where("x.a".attr === "z.a".attr)))
       .analyze
     val answer = x
-      .where(Exists(z.where("x.a".attr === "z.a".attr)))
       .join(y, Inner, Option("x.a".attr === "y.a".attr))
+      .where(Exists(z.where("x.a".attr === "z.a".attr)))
       .analyze
     val optimized = Optimize.execute(Optimize.execute(query))
     comparePlans(optimized, answer)
   }
 
-  test("predicate subquery: push down complex") {
+  test("correlated subquery (complex): no push down") {
     val w = testRelation.subquery('w)
     val x = testRelation.subquery('x)
     val y = testRelation.subquery('y)
@@ -828,9 +828,9 @@ class FilterPushdownSuite extends PlanTest {
       .where(Exists(z.where("w.a".attr === "z.a".attr)))
       .analyze
     val answer = w
-      .where(Exists(z.where("w.a".attr === "z.a".attr)))
       .join(x, Inner, Option("w.a".attr === "x.a".attr))
       .join(y, LeftOuter, Option("x.a".attr === "y.a".attr))
+      .where(Exists(z.where("w.a".attr === "z.a".attr)))
       .analyze
     val optimized = Optimize.execute(Optimize.execute(query))
     comparePlans(optimized, answer)