Generalize pattern for planning hash joins.

Author: marmbrus

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

@@ -19,7 +19,10 @@ package org.apache.spark.sql.catalyst.planning
 
 import scala.annotation.tailrec
 
+import org.apache.spark.sql.Logging
+
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.catalyst.plans.logical._
 
 /**
@@ -101,6 +104,55 @@ object PhysicalOperation extends PredicateHelper {
   }
 }
 
+/**
+ * A pattern that finds joins with equality conditions that can be evaluated using hashing
+ * techniques.  For inner joins, any filters on top of the join operator are also matched.
+ */
+object HashFilteredJoin extends Logging with PredicateHelper {
+  /** (joinType, rightKeys, leftKeys, condition, left, right) */
+  type ReturnType =
+    (JoinType, Seq[Expression], Seq[Expression], Option[Expression], LogicalPlan, LogicalPlan)
+
+  def unapply(plan: LogicalPlan): Option[ReturnType] = plan match {
+    // All predicates can be evaluated for inner join (i.e., those that are in the ON
+    // clause and WHERE clause.)
+    case FilteredOperation(predicates, join @ Join(left, right, Inner, condition)) =>
+      logger.debug(s"Considering hash inner join on: ${predicates ++ condition}")
+      splitPredicates(predicates ++ condition, join)
+    case join @ Join(left, right, joinType, condition) =>
+      logger.debug(s"Considering hash join on: $condition")
+      splitPredicates(condition.toSeq, join)
+    case _ => None
+  }
+
+  // Find equi-join predicates that can be evaluated before the join, and thus can be used
+  // as join keys.
+  def splitPredicates(allPredicates: Seq[Expression], join: Join): Option[ReturnType] = {
+    val Join(left, right, joinType, _) = join
+    val (joinPredicates, otherPredicates) = allPredicates.partition {
+      case Equals(l, r) if (canEvaluate(l, left) && canEvaluate(r, right)) ||
+        (canEvaluate(l, right) && canEvaluate(r, left)) => true
+      case _ => false
+    }
+
+    val joinKeys = joinPredicates.map {
+      case Equals(l,r) if canEvaluate(l, left) && canEvaluate(r, right) => (l, r)
+      case Equals(l,r) if canEvaluate(l, right) && canEvaluate(r, left) => (r, l)
+    }
+
+    // Do not consider this strategy if there are no join keys.
+    if (joinKeys.nonEmpty) {
+      val leftKeys = joinKeys.map(_._1)
+      val rightKeys = joinKeys.map(_._2)
+
+      Some((joinType, leftKeys, rightKeys, otherPredicates.reduceOption(And), left, right))
+    } else {
+      logger.debug(s"Avoiding hash join with no join keys.")
+      None
+    }
+  }
+}
+
 /**
  * A pattern that collects all adjacent unions and returns their children as a Seq.
  */

sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala

@@ -28,42 +28,12 @@ import org.apache.spark.sql.parquet._
 abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
   self: SQLContext#SparkPlanner =>
 
-  object HashJoin extends Strategy {
+  object HashJoin extends Strategy with PredicateHelper {
     def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
-      case FilteredOperation(predicates, logical.Join(left, right, Inner, condition)) =>
-        logger.debug(s"Considering join: ${predicates ++ condition}")
-        // Find equi-join predicates that can be evaluated before the join, and thus can be used
-        // as join keys. Note we can only mix in the conditions with other predicates because the
-        // match above ensures that this is and Inner join.
-        val (joinPredicates, otherPredicates) = (predicates ++ condition).partition {
-          case Equals(l, r) if (canEvaluate(l, left) && canEvaluate(r, right)) ||
-                               (canEvaluate(l, right) && canEvaluate(r, left)) => true
-          case _ => false
-        }
-
-        val joinKeys = joinPredicates.map {
-          case Equals(l,r) if canEvaluate(l, left) && canEvaluate(r, right) => (l, r)
-          case Equals(l,r) if canEvaluate(l, right) && canEvaluate(r, left) => (r, l)
-        }
-
-        // Do not consider this strategy if there are no join keys.
-        if (joinKeys.nonEmpty) {
-          val leftKeys = joinKeys.map(_._1)
-          val rightKeys = joinKeys.map(_._2)
-
-          val joinOp = execution.HashJoin(
-            leftKeys, rightKeys, BuildRight, planLater(left), planLater(right))
-
-          // Make sure other conditions are met if present.
-          if (otherPredicates.nonEmpty) {
-            execution.Filter(combineConjunctivePredicates(otherPredicates), joinOp) :: Nil
-          } else {
-            joinOp :: Nil
-          }
-        } else {
-          logger.debug(s"Avoiding spark join with no join keys.")
-          Nil
-        }
+      case HashFilteredJoin(Inner, leftKeys, rightKeys, condition, left, right) =>
+        val hashJoin =
+          execution.HashJoin(leftKeys, rightKeys, BuildRight, planLater(left), planLater(right))
+        condition.map(Filter(_, hashJoin)).getOrElse(hashJoin) :: Nil
       case _ => Nil
     }
   }