[SPARK-12032] [SQL] Re-order inner joins to do join with conditions first

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

@@ -18,14 +18,12 @@
 package org.apache.spark.sql.catalyst.optimizer
 
 import scala.collection.immutable.HashSet
+
 import org.apache.spark.sql.catalyst.analysis.{CleanupAliases, EliminateSubQueries}
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate._
-import org.apache.spark.sql.catalyst.plans.Inner
-import org.apache.spark.sql.catalyst.plans.FullOuter
-import org.apache.spark.sql.catalyst.plans.LeftOuter
-import org.apache.spark.sql.catalyst.plans.RightOuter
-import org.apache.spark.sql.catalyst.plans.LeftSemi
+import org.apache.spark.sql.catalyst.planning.FilterAndInnerJoins
+import org.apache.spark.sql.catalyst.plans.{FullOuter, Inner, LeftOuter, LeftSemi, RightOuter}
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules._
 import org.apache.spark.sql.types._
@@ -44,6 +42,7 @@ object DefaultOptimizer extends Optimizer {
       // Operator push down
       SetOperationPushDown,
       SamplePushDown,
+      ReorderJoin,
       PushPredicateThroughJoin,
       PushPredicateThroughProject,
       PushPredicateThroughGenerate,
@@ -711,6 +710,49 @@ object PushPredicateThroughAggregate extends Rule[LogicalPlan] with PredicateHel
   }
 }
 
+/**
+  * Reorder the joins so that the bottom ones have at least one condition.
+  */
+object ReorderJoin extends Rule[LogicalPlan] with PredicateHelper {
+
+  /**
+    * Reorder the joins so that the bottom ones have at least one condition.
+    */
+  def reorder(
+      input: LogicalPlan,
+      joins: Seq[LogicalPlan],
+      conditions: Seq[Expression]): LogicalPlan = {
+    // filter out the conditions that could be pushed down to `joined`
+    val otherConditions = conditions.filterNot { cond =>
+      cond.references.subsetOf(input.outputSet)
+    }
+    if (joins.isEmpty) {
+      input
+    } else if (otherConditions.isEmpty) {
+      // no condition for these joins, so put them in original order
+      (Seq(input) ++ joins).reduceLeft(Join(_, _, Inner, None))
+    } else {
+      // find out the first join that have at least one condition
+      val conditionalJoin = joins.find { plan =>
+        val refs = input.outputSet ++ plan.outputSet
+        otherConditions.exists(cond => cond.references.subsetOf(refs))
+      }
+      assert(conditionalJoin.isDefined)
+      val picked = conditionalJoin.get
+      val joined = Join(input, picked, Inner, None)
+      reorder(joined, joins.filterNot(_ eq picked), otherConditions)
+    }
+  }
+
+  def apply(plan: LogicalPlan): LogicalPlan = plan transform {
+    // TODO: support outer join
+    case FilterAndInnerJoins(input, joins, filterConditions) if joins.size > 1 =>
+      assert(filterConditions.nonEmpty)
+      val joined = reorder(input, joins, filterConditions)
+      Filter(filterConditions.reduceLeft(And), joined)
+  }
+}
+
 /**
  * Pushes down [[Filter]] operators where the `condition` can be
  * evaluated using only the attributes of the left or right side of a join.  Other

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/planning/patterns.scala

@@ -21,7 +21,6 @@ import org.apache.spark.Logging
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.catalyst.plans.logical._
-import org.apache.spark.sql.catalyst.trees.TreeNodeRef
 
 /**
  * A pattern that matches any number of project or filter operations on top of another relational
@@ -132,6 +131,38 @@ object ExtractEquiJoinKeys extends Logging with PredicateHelper {
   }
 }
 
+/**
+  * A pattern that collects the filter and inner joins.
+  *
+  *          Filter
+  *            |
+  *        inner Join
+  *          /    \            ---->      (filters, Seq(plan1, plan2), input)
+  *  inner join    plan2
+  *      /    \
+  *   input    plan1
+  */
+object FilterAndInnerJoins extends PredicateHelper {
+  def unapply(plan: LogicalPlan): Option[(LogicalPlan, Seq[LogicalPlan], Seq[Expression])] =
+    plan match {
+      case f @ Filter(filterCondition, j @ Join(left, right, Inner, None)) =>
+
+        // flatten all inner joins, which are next to each other and has no condition
+        def flattenJoin(plan: LogicalPlan): (LogicalPlan, Seq[LogicalPlan]) = plan match {
+          case Join(left, right, Inner, None) =>
+            val (input, joins) = flattenJoin(left)
+            (input, joins ++ Seq(right))
+          case _ => (plan, Seq())
+        }
+
+        val allConditions = splitConjunctivePredicates(filterCondition)
+        val (input, joins) = flattenJoin(j)
+        Some((input, joins, allConditions))
+
+    case _ => None
+  }
+}
+
 /**
  * A pattern that collects all adjacent unions and returns their children as a Seq.
  */

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

@@ -38,6 +38,7 @@ class FilterPushdownSuite extends PlanTest {
         CombineFilters,
         PushPredicateThroughProject,
         BooleanSimplification,
+        ReorderJoin,
         PushPredicateThroughJoin,
         PushPredicateThroughGenerate,
         PushPredicateThroughAggregate,
@@ -548,6 +549,25 @@ class FilterPushdownSuite extends PlanTest {
     comparePlans(optimized, analysis.EliminateSubQueries(correctAnswer))
   }
 
+  test("joins: reorder inner joins") {
+    val x = testRelation.subquery('x)
+    val y = testRelation1.subquery('y)
+    val z = testRelation.subquery('z)
+
+    val originalQuery = {
+      x.join(y).join(z)
+        .where(("x.b".attr === "z.b".attr) && ("y.d".attr === "z.a".attr))
+    }
+
+    val optimized = Optimize.execute(originalQuery.analyze)
+    val correctAnswer =
+      x.join(z, condition = Some("x.b".attr === "z.b".attr))
+        .join(y, condition = Some("y.d".attr === "z.a".attr))
+        .analyze
+
+    comparePlans(optimized, analysis.EliminateSubQueries(correctAnswer))
+  }
+
   val testRelationWithArrayType = LocalRelation('a.int, 'b.int, 'c_arr.array(IntegerType))
 
   test("generate: predicate referenced no generated column") {
@@ -750,4 +770,5 @@ class FilterPushdownSuite extends PlanTest {
 
     comparePlans(optimized, correctAnswer)
   }
+
 }