[SPARK-19712][SQL] Pushdown LeftSemi/LeftAnti below join

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

@@ -66,6 +66,7 @@ abstract class Optimizer(sessionCatalog: SessionCatalog)
         PushPredicateThroughJoin,
         PushDownPredicate,
         PushDownLeftSemiAntiJoin,
+        PushLeftSemiLeftAntiThroughJoin,
         LimitPushDown,
         ColumnPruning,
         InferFiltersFromConstraints,

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

@@ -159,3 +159,106 @@ object PushDownLeftSemiAntiJoin extends Rule[LogicalPlan] with PredicateHelper {
     }
   }
 }
+
+/**
+ * This rule is a variant of [[PushPredicateThroughJoin]] which can handle
+ * pushing down Left semi and Left Anti joins below a join operator. The
+ * allowable join types are:
+ *  1) Inner
+ *  2) Cross
+ *  3) LeftOuter
+ *  4) RightOuter
+ */
+object PushLeftSemiLeftAntiThroughJoin extends Rule[LogicalPlan] with PredicateHelper {
+  /**
+   * Define an enumeration to identify whether a LeftSemi/LeftAnti join can be pushed down to
+   * the left leg or the right leg of the join.
+   */
+  object pushdownDirection extends Enumeration {
+    val toRightBranch, toLeftBranch, none = Value
+  }
+
+  /**
+   * LeftSemi/LeftAnti joins are pushed down when its left child is a join operator
+   * with a join type that is in the AllowedJoinTypes.
+   */
+  object AllowedJoinTypes {
+    def unapply(joinType: JoinType): Option[JoinType] = joinType match {
+      case Inner | Cross | LeftOuter | RightOuter => Some(joinType)
+      case _ => None
+    }
+  }
+
+  /**
+   * Determine which side of the join a LeftSemi/LeftAnti join can be pushed to.
+   */
+  private def pushTo(leftChild: Join, rightChild: LogicalPlan, joinCond: Option[Expression]) = {
+    val left = leftChild.left
+    val right = leftChild.right
+    val joinType = leftChild.joinType
+    val rightOutput = rightChild.outputSet
+
+    if (joinCond.nonEmpty) {
+      val noPushdown = (pushdownDirection.none, None)
+      val conditions = splitConjunctivePredicates(joinCond.get)
+      val (leftConditions, rest) =
+        conditions.partition(_.references.subsetOf(left.outputSet ++ rightOutput))
+      val (rightConditions, commonConditions) =
+        rest.partition(_.references.subsetOf(right.outputSet ++ rightOutput))
+
+      if (rest.isEmpty && leftConditions.nonEmpty) {
+        // When the join conditions can be computed based on the left leg of
+        // leftsemi/anti join then push the leftsemi/anti join to the left side.
+        (pushdownDirection.toLeftBranch, leftConditions.reduceLeftOption(And))
+      } else if (leftConditions.isEmpty && rightConditions.nonEmpty && commonConditions.isEmpty) {
+        // When the join conditions can be computed based on the attributes from right leg of
+        // leftsemi/anti join then push the leftsemi/anti join to the right side.
+        (pushdownDirection.toRightBranch, rightConditions.reduceLeftOption(And))
+      } else {
+        noPushdown
+      }
+    } else {
+      /**
+       * When the join condition is empty,
+       * 1) if this is a left outer join or inner join, push leftsemi/anti join down
+       *    to the left leg of join.
+       * 2) if a right outer join, to the right leg of join,
+       */
+      val action = joinType match {
+        case RightOuter =>
+          pushdownDirection.toRightBranch
+        case _: InnerLike | LeftOuter =>
+          pushdownDirection.toLeftBranch
+        case _ =>
+          pushdownDirection.none
+      }
+      (action, None)
+    }
+  }
+
+  def apply(plan: LogicalPlan): LogicalPlan = plan transform {
+    // push LeftSemi/LeftAnti down into the join below
+    case j @ Join(left @ Join(gLeft, gRight, AllowedJoinTypes(_), belowJoinCond, childHint),
+    right, LeftSemiOrAnti(joinType), joinCond, parentHint) =>
+      val belowJoinType = left.joinType
+      val (action, newJoinCond) = pushTo(left, right, joinCond)
+
+      action match {
+        case pushdownDirection.toLeftBranch
+          if (belowJoinType == LeftOuter || belowJoinType.isInstanceOf[InnerLike]) =>
+          // push down leftsemi/anti join to the left table
+          val newLeft = Join(gLeft, right, joinType, newJoinCond, parentHint)
+          Join(newLeft, gRight, belowJoinType, belowJoinCond, childHint)
+        case pushdownDirection.toRightBranch
+          if (belowJoinType == RightOuter || belowJoinType.isInstanceOf[InnerLike]) =>
+          // push down leftsemi/anti join to the right table
+          val newRight = Join(gRight, right, joinType, newJoinCond, parentHint)
+          Join(gLeft, newRight, belowJoinType, belowJoinCond, childHint)
+        case _ =>
+          // Do nothing
+          j
+      }
+  }
+}
+
+

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

@@ -37,13 +37,14 @@ class LeftSemiPushdownSuite extends PlanTest {
         CombineFilters,
         PushDownPredicate,
         PushDownLeftSemiAntiJoin,
+        PushLeftSemiLeftAntiThroughJoin,
         BooleanSimplification,
         CollapseProject) :: Nil
   }
 
   val testRelation = LocalRelation('a.int, 'b.int, 'c.int)
-
   val testRelation1 = LocalRelation('d.int)
+  val testRelation2 = LocalRelation('e.int)
 
   test("Project: LeftSemiAnti join pushdown") {
     val originalQuery = testRelation
@@ -314,4 +315,155 @@ class LeftSemiPushdownSuite extends PlanTest {
     val optimized = Optimize.execute(originalQuery.analyze)
     comparePlans(optimized, originalQuery.analyze)
   }
+
+  Seq(LeftSemi, LeftAnti).foreach { case outerJT =>
+    Seq(Inner, LeftOuter, Cross).foreach { case innerJT =>
+      test(s"$outerJT pushdown with empty join condition join type $innerJT") {
+        val joinedRelation = testRelation1.join(testRelation2, joinType = innerJT, None)
+        val originalQuery = joinedRelation.join(testRelation, joinType = outerJT, None)
+        val optimized = Optimize.execute(originalQuery.analyze)
+
+        val pushedDownJoin = testRelation1.join(testRelation, joinType = outerJT, None)
+        val correctAnswer = pushedDownJoin.join(testRelation2, joinType = innerJT, None)
+        comparePlans(optimized, correctAnswer)
+      }
+    }
+  }
+
+  Seq(LeftSemi, LeftAnti).foreach { case jt =>
+      test(s"$jt pushdown with empty join condition join type RightOuter") {
+        val joinedRelation = testRelation1.join(testRelation2, joinType = RightOuter, None)
+        val originalQuery = joinedRelation.join(testRelation, joinType = jt, None)
+        val optimized = Optimize.execute(originalQuery.analyze)
+
+        val pushedDownJoin = testRelation2.join(testRelation, joinType = jt, None)
+        val correctAnswer = testRelation1.join(pushedDownJoin, joinType = RightOuter, None)
+        comparePlans(optimized, correctAnswer)
+      }
+  }
+
+  Seq(LeftSemi, LeftAnti).foreach { case outerJT =>
+    Seq(Inner, LeftOuter, Cross).foreach { case innerJT =>
+      test(s"$outerJT pushdown with join condition referring to left leg of join type $innerJT") {
+        val joinedRelation = testRelation1.join(testRelation2, joinType = innerJT, None)
+        val originalQuery =
+          joinedRelation.join(testRelation, joinType = outerJT, condition = Some('a === 'd))
+        val optimized = Optimize.execute(originalQuery.analyze)
+
+        val pushedDownJoin =
+          testRelation1.join(testRelation, joinType = outerJT, condition = Some('a === 'd))
+        val correctAnswer = pushedDownJoin.join(testRelation2, joinType = innerJT, None).analyze
+        comparePlans(optimized, correctAnswer)
+      }
+    }
+  }
+
+  Seq(LeftSemi, LeftAnti).foreach { case outerJT =>
+    Seq(Inner, LeftOuter, Cross).foreach { case innerJT =>
+      test(s"$outerJT pushdown with outer and inner join condition for join type $innerJT") {
+        val joinedRelation =
+          testRelation1.join(testRelation2, joinType = innerJT, condition = Some('d === 'e))
+        val originalQuery =
+          joinedRelation.join(testRelation, joinType = outerJT, condition = Some('a === 'd))
+        val optimized = Optimize.execute(originalQuery.analyze)
+
+        val pushedDownJoin =
+          testRelation1.join(testRelation, joinType = outerJT, condition = Some('a === 'd))
+        val correctAnswer = pushedDownJoin
+          .join(testRelation2, joinType = innerJT, condition = Some('d === 'e))
+          .analyze
+        comparePlans(optimized, correctAnswer)
+      }
+    }
+  }
+
+  Seq(LeftSemi, LeftAnti).foreach { case jt =>
+    test(s"$jt no pushdown - join condition refers left leg - join type for RightOuter") {
+      val joinedRelation = testRelation1.join(testRelation2, joinType = RightOuter, None)
+      val originalQuery =
+        joinedRelation.join(testRelation, joinType = jt, condition = Some('a === 'd))
+      val optimized = Optimize.execute(originalQuery.analyze)
+      comparePlans(optimized, originalQuery.analyze)
+    }
+  }
+
+  Seq(LeftSemi, LeftAnti).foreach { case outerJT =>
+    Seq(Inner, RightOuter, Cross).foreach { case innerJT =>
+      test(s"$outerJT pushdown with join condition referring to right leg - join type $innerJT") {
+        val joinedRelation = testRelation1.join(testRelation2, joinType = innerJT, None)
+        val originalQuery =
+          joinedRelation.join(testRelation, joinType = outerJT, condition = Some('a === 'e))
+        val optimized = Optimize.execute(originalQuery.analyze)
+
+        val pushedDownJoin =
+          testRelation2.join(testRelation, joinType = outerJT, condition = Some('a === 'e))
+        val correctAnswer = testRelation1.join(pushedDownJoin, joinType = innerJT, None).analyze
+        comparePlans(optimized, correctAnswer)
+      }
+    }
+  }
+
+  Seq(LeftSemi, LeftAnti).foreach { case outerJT =>
+    Seq(Inner, RightOuter, Cross).foreach { case innerJT =>
+      test(s"$outerJT pushdown with outer and inner join conditions for join type $innerJT") {
+        val joinedRelation = testRelation1.
+          join(testRelation2, joinType = innerJT, condition = Some('e === 'd))
+        val originalQuery =
+          joinedRelation.join(testRelation, joinType = outerJT, condition = Some('a === 'e))
+        val optimized = Optimize.execute(originalQuery.analyze)
+
+        val pushedDownJoin =
+          testRelation2.join(testRelation, joinType = outerJT, condition = Some('a === 'e))
+        val correctAnswer = testRelation1.
+          join(pushedDownJoin, joinType = innerJT, condition = Some('e === 'd))
+          .analyze
+        comparePlans(optimized, correctAnswer)
+      }
+    }
+  }
+
+  Seq(LeftSemi, LeftAnti).foreach { case jt =>
+    test(s"$jt no pushdown - join condition refers right leg - join type for LeftOuter") {
+      val joinedRelation = testRelation1.join(testRelation2, joinType = LeftOuter, None)
+      val originalQuery =
+        joinedRelation.join(testRelation, joinType = jt, condition = Some('a === 'e))
+      val optimized = Optimize.execute(originalQuery.analyze)
+      comparePlans(optimized, originalQuery.analyze)
+    }
+  }
+
+  Seq(LeftSemi, LeftAnti).foreach { case outerJT =>
+    Seq(Inner, LeftOuter, RightOuter, Cross).foreach { case innerJT =>
+      test(s"$outerJT no pushdown - join condition refers both leg - join type $innerJT") {
+        val joinedRelation = testRelation1.join(testRelation2, joinType = innerJT, None)
+        val originalQuery = joinedRelation
+          .join(testRelation, joinType = outerJT, condition = Some('a === 'd && 'a === 'e))
+        val optimized = Optimize.execute(originalQuery.analyze)
+        comparePlans(optimized, originalQuery.analyze)
+      }
+    }
+  }
+
+  Seq(LeftSemi, LeftAnti).foreach { case outerJT =>
+    Seq(Inner, LeftOuter, RightOuter, Cross).foreach { case innerJT =>
+      test(s"$outerJT no pushdown - join condition refers none of the leg - join type $innerJT") {
+        val joinedRelation = testRelation1.join(testRelation2, joinType = innerJT, None)
+        val originalQuery = joinedRelation
+          .join(testRelation, joinType = outerJT, condition = Some('d + 'e === 'a))
+        val optimized = Optimize.execute(originalQuery.analyze)
+        comparePlans(optimized, originalQuery.analyze)
+      }
+    }
+  }
+
+  Seq(LeftSemi, LeftAnti).foreach { case jt =>
+    test(s"$jt no pushdown when child join type is FullOuter") {
+      val joinedRelation = testRelation1.join(testRelation2, joinType = FullOuter, None)
+      val originalQuery =
+        joinedRelation.join(testRelation, joinType = jt, condition = Some('a === 'e))
+      val optimized = Optimize.execute(originalQuery.analyze)
+      comparePlans(optimized, originalQuery.analyze)
+    }
+  }
+
 }