[SPARK-9372] [SQL] For joins, insert IS NOT NULL filters to children.

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala

@@ -389,7 +389,7 @@ class Analyzer(
         a.copy(aggregateExpressions = expanded)
 
       // Special handling for cases when self-join introduce duplicate expression ids.
-      case j @ Join(left, right, _, _) if !j.selfJoinResolved =>
+      case j @ Join(left, right, _, _, _) if !j.selfJoinResolved =>
         val conflictingAttributes = left.outputSet.intersect(right.outputSet)
         logDebug(s"Conflicting attributes ${conflictingAttributes.mkString(",")} in $j")
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala

@@ -19,7 +19,7 @@ package org.apache.spark.sql.catalyst.analysis
 
 import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateFunction, AggregateExpression}
+import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.types._
 
@@ -87,12 +87,12 @@ trait CheckAnalysis {
               s"filter expression '${f.condition.prettyString}' " +
                 s"of type ${f.condition.dataType.simpleString} is not a boolean.")
 
-          case j @ Join(_, _, _, Some(condition)) if condition.dataType != BooleanType =>
+          case j @ Join(_, _, _, Some(condition), _) if condition.dataType != BooleanType =>
             failAnalysis(
               s"join condition '${condition.prettyString}' " +
                 s"of type ${condition.dataType.simpleString} is not a boolean.")
 
-          case j @ Join(_, _, _, Some(condition)) =>
+          case j @ Join(_, _, _, Some(condition), _) =>
             def checkValidJoinConditionExprs(expr: Expression): Unit = expr match {
               case p: Predicate =>
                 p.asInstanceOf[Expression].children.foreach(checkValidJoinConditionExprs)
@@ -190,14 +190,15 @@ trait CheckAnalysis {
                  | ${exprs.map(_.prettyString).mkString(",")}""".stripMargin)
 
           // Special handling for cases when self-join introduce duplicate expression ids.
-          case j @ Join(left, right, _, _) if left.outputSet.intersect(right.outputSet).nonEmpty =>
-            val conflictingAttributes = left.outputSet.intersect(right.outputSet)
-            failAnalysis(
-              s"""
-                 |Failure when resolving conflicting references in Join:
-                 |$plan
-                 |Conflicting attributes: ${conflictingAttributes.mkString(",")}
-                 |""".stripMargin)
+          case j @ Join(left, right, _, _, _)
+            if left.outputSet.intersect(right.outputSet).nonEmpty =>
+              val conflictingAttributes = left.outputSet.intersect(right.outputSet)
+              failAnalysis(
+                s"""
+                   |Failure when resolving conflicting references in Join:
+                   |$plan
+                   |Conflicting attributes: ${conflictingAttributes.mkString(",")}
+                   |""".stripMargin)
 
           case o if !o.resolved =>
             failAnalysis(

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala

@@ -87,6 +87,13 @@ class EquivalentExpressions {
     equivalenceMap.values.map(_.toSeq).toSeq
   }
 
+  /**
+   * Returns true if e exists.
+   */
+  def contains(e: Expression): Boolean = {
+    equivalenceMap.contains(Expr(e))
+  }
+
   /**
    * Returns the state of the data structure as a string. If `all` is false, skips sets of
    * equivalent expressions with cardinality 1.

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

@@ -23,6 +23,7 @@ import org.apache.spark.sql.catalyst.analysis.{CleanupAliases, EliminateSubQueri
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.planning.ExtractFiltersAndInnerJoins
+import org.apache.spark.sql.catalyst.planning.ExtractNonNullableAttributes
 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._
@@ -52,6 +53,8 @@ object DefaultOptimizer extends Optimizer {
       ProjectCollapsing,
       CombineFilters,
       CombineLimits,
+      // Predicate inference
+      AddJoinKeyNullabilityFilters,
       // Constant folding
       NullPropagation,
       OptimizeIn,
@@ -233,7 +236,7 @@ object ColumnPruning extends Rule[LogicalPlan] {
           child))
 
     // Eliminate unneeded attributes from either side of a Join.
-    case Project(projectList, Join(left, right, joinType, condition)) =>
+    case Project(projectList, Join(left, right, joinType, condition, generated)) =>
       // Collect the list of all references required either above or to evaluate the condition.
       val allReferences: AttributeSet =
         AttributeSet(
@@ -243,15 +246,16 @@ object ColumnPruning extends Rule[LogicalPlan] {
       /** Applies a projection only when the child is producing unnecessary attributes */
       def pruneJoinChild(c: LogicalPlan): LogicalPlan = prunedChild(c, allReferences)
 
-      Project(projectList, Join(pruneJoinChild(left), pruneJoinChild(right), joinType, condition))
+      Project(projectList,
+          Join(pruneJoinChild(left), pruneJoinChild(right), joinType, condition, generated))
 
     // Eliminate unneeded attributes from right side of a LeftSemiJoin.
-    case Join(left, right, LeftSemi, condition) =>
+    case Join(left, right, LeftSemi, condition, generated) =>
       // Collect the list of all references required to evaluate the condition.
       val allReferences: AttributeSet =
         condition.map(_.references).getOrElse(AttributeSet(Seq.empty))
 
-      Join(left, prunedChild(right, allReferences), LeftSemi, condition)
+      Join(left, prunedChild(right, allReferences), LeftSemi, condition, generated)
 
     // Push down project through limit, so that we may have chance to push it further.
     case Project(projectList, Limit(exp, child)) =>
@@ -355,6 +359,51 @@ object LikeSimplification extends Rule[LogicalPlan] {
   }
 }
 
+/**
+ * This rule adds IsNotNull predicates based on join keys. If the join contains a condition
+ * `a` binaryOp *, a is non-nullable. This adds filters for those attributes to the children.
+ *
+ * To avoid the problem of repeatedly generating the IsNotNull predicates, the join operator
+ * remembers all the expressions it generated.
+ */
+object AddJoinKeyNullabilityFilters extends Rule[LogicalPlan] with PredicateHelper {
+  def apply(plan: LogicalPlan): LogicalPlan = plan transform {
+    case j @ Join(left, right, Inner, Some(condition), generated) => {
+      val nonNullableKeys = ExtractNonNullableAttributes.unapply(condition)
+      val newKeys = nonNullableKeys.filter { generated.isEmpty || !generated.get.contains(_) }
+
+      if (newKeys.isEmpty) {
+        j
+      } else {
+        val leftKeys = newKeys.filter { canEvaluate(_, left) }
+        val rightKeys = newKeys.filter { canEvaluate(_, right) }
+
+        if (leftKeys.nonEmpty || rightKeys.nonEmpty) {
+          val newGenerated =
+            if (j.generatedExpressions.isDefined) j.generatedExpressions.get
+            else new EquivalentExpressions
+
+          var newLeft: LogicalPlan = left
+          var newRight: LogicalPlan = right
+
+          if (leftKeys.nonEmpty) {
+            newLeft = Filter(leftKeys.map(IsNotNull(_)).reduce(And), left)
+            leftKeys.foreach { e => newGenerated.addExpr(e) }
+          }
+          if (rightKeys.nonEmpty) {
+            newRight = Filter(rightKeys.map(IsNotNull(_)).reduce(And), right)
+            rightKeys.foreach { e => newGenerated.addExpr(e) }
+          }
+
+          Join(newLeft, newRight, Inner, Some(condition), Some(newGenerated))
+        } else {
+          j
+        }
+      }
+    }
+  }
+}
+
 /**
  * Replaces [[Expression Expressions]] that can be statically evaluated with
  * equivalent [[Literal]] values. This rule is more specific with
@@ -784,7 +833,7 @@ object PushPredicateThroughJoin extends Rule[LogicalPlan] with PredicateHelper {
 
   def apply(plan: LogicalPlan): LogicalPlan = plan transform {
     // push the where condition down into join filter
-    case f @ Filter(filterCondition, Join(left, right, joinType, joinCondition)) =>
+    case f @ Filter(filterCondition, Join(left, right, joinType, joinCondition, generated)) =>
       val (leftFilterConditions, rightFilterConditions, commonFilterCondition) =
         split(splitConjunctivePredicates(filterCondition), left, right)
 
@@ -797,14 +846,14 @@ object PushPredicateThroughJoin extends Rule[LogicalPlan] with PredicateHelper {
             reduceLeftOption(And).map(Filter(_, right)).getOrElse(right)
           val newJoinCond = (commonFilterCondition ++ joinCondition).reduceLeftOption(And)
 
-          Join(newLeft, newRight, Inner, newJoinCond)
+          Join(newLeft, newRight, Inner, newJoinCond, generated)
         case RightOuter =>
           // push down the right side only `where` condition
           val newLeft = left
           val newRight = rightFilterConditions.
             reduceLeftOption(And).map(Filter(_, right)).getOrElse(right)
           val newJoinCond = joinCondition
-          val newJoin = Join(newLeft, newRight, RightOuter, newJoinCond)
+          val newJoin = Join(newLeft, newRight, RightOuter, newJoinCond, generated)
 
           (leftFilterConditions ++ commonFilterCondition).
             reduceLeftOption(And).map(Filter(_, newJoin)).getOrElse(newJoin)
@@ -814,15 +863,15 @@ object PushPredicateThroughJoin extends Rule[LogicalPlan] with PredicateHelper {
             reduceLeftOption(And).map(Filter(_, left)).getOrElse(left)
           val newRight = right
           val newJoinCond = joinCondition
-          val newJoin = Join(newLeft, newRight, joinType, newJoinCond)
+          val newJoin = Join(newLeft, newRight, joinType, newJoinCond, generated)
 
           (rightFilterConditions ++ commonFilterCondition).
             reduceLeftOption(And).map(Filter(_, newJoin)).getOrElse(newJoin)
         case FullOuter => f // DO Nothing for Full Outer Join
       }
 
     // push down the join filter into sub query scanning if applicable
-    case f @ Join(left, right, joinType, joinCondition) =>
+    case f @ Join(left, right, joinType, joinCondition, generated) =>
       val (leftJoinConditions, rightJoinConditions, commonJoinCondition) =
         split(joinCondition.map(splitConjunctivePredicates).getOrElse(Nil), left, right)
 
@@ -835,23 +884,23 @@ object PushPredicateThroughJoin extends Rule[LogicalPlan] with PredicateHelper {
             reduceLeftOption(And).map(Filter(_, right)).getOrElse(right)
           val newJoinCond = commonJoinCondition.reduceLeftOption(And)
 
-          Join(newLeft, newRight, joinType, newJoinCond)
+          Join(newLeft, newRight, joinType, newJoinCond, generated)
         case RightOuter =>
           // push down the left side only join filter for left side sub query
           val newLeft = leftJoinConditions.
             reduceLeftOption(And).map(Filter(_, left)).getOrElse(left)
           val newRight = right
           val newJoinCond = (rightJoinConditions ++ commonJoinCondition).reduceLeftOption(And)
 
-          Join(newLeft, newRight, RightOuter, newJoinCond)
+          Join(newLeft, newRight, RightOuter, newJoinCond, generated)
         case LeftOuter =>
           // push down the right side only join filter for right sub query
           val newLeft = left
           val newRight = rightJoinConditions.
             reduceLeftOption(And).map(Filter(_, right)).getOrElse(right)
           val newJoinCond = (leftJoinConditions ++ commonJoinCondition).reduceLeftOption(And)
 
-          Join(newLeft, newRight, LeftOuter, newJoinCond)
+          Join(newLeft, newRight, LeftOuter, newJoinCond, generated)
         case FullOuter => f
       }
   }

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

@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql.catalyst.planning
 
+import scala.collection.mutable
+
 import org.apache.spark.Logging
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans._
@@ -95,7 +97,7 @@ object ExtractEquiJoinKeys extends Logging with PredicateHelper {
     (JoinType, Seq[Expression], Seq[Expression], Option[Expression], LogicalPlan, LogicalPlan)
 
   def unapply(plan: LogicalPlan): Option[ReturnType] = plan match {
-    case join @ Join(left, right, joinType, condition) =>
+    case join @ Join(left, right, joinType, condition, _) =>
       logDebug(s"Considering join on: $condition")
       // Find equi-join predicates that can be evaluated before the join, and thus can be used
       // as join keys.
@@ -150,21 +152,21 @@ object ExtractFiltersAndInnerJoins extends PredicateHelper {
 
   // flatten all inner joins, which are next to each other
   def flattenJoin(plan: LogicalPlan): (Seq[LogicalPlan], Seq[Expression]) = plan match {
-    case Join(left, right, Inner, cond) =>
+    case Join(left, right, Inner, cond, _) =>
       val (plans, conditions) = flattenJoin(left)
       (plans ++ Seq(right), conditions ++ cond.toSeq)
 
-    case Filter(filterCondition, j @ Join(left, right, Inner, joinCondition)) =>
+    case Filter(filterCondition, j @ Join(left, right, Inner, joinCondition, _)) =>
       val (plans, conditions) = flattenJoin(j)
       (plans, conditions ++ splitConjunctivePredicates(filterCondition))
 
     case _ => (Seq(plan), Seq())
   }
 
   def unapply(plan: LogicalPlan): Option[(Seq[LogicalPlan], Seq[Expression])] = plan match {
-    case f @ Filter(filterCondition, j @ Join(_, _, Inner, _)) =>
+    case f @ Filter(filterCondition, j @ Join(_, _, Inner, _, _)) =>
       Some(flattenJoin(f))
-    case j @ Join(_, _, Inner, _) =>
+    case j @ Join(_, _, Inner, _, _) =>
       Some(flattenJoin(j))
     case _ => None
   }
@@ -184,3 +186,39 @@ object Unions {
     case other => other :: Nil
   }
 }
+
+/**
+ * A pattern that finds all attributes in `expr` that cannot be nullable.
+ */
+object ExtractNonNullableAttributes extends Logging with PredicateHelper {
+  def unapply(condition: Expression): Set[Attribute] = {
+    val predicates = splitConjunctivePredicates(condition)
+
+    val result = mutable.HashSet.empty[Attribute]
+    def extract(e: Expression): Unit = e match {
+      case IsNotNull(a: Attribute) => result.add(a)
+      case BinaryComparison(a: Attribute, b: Attribute) => {
+        if (!e.isInstanceOf[EqualNullSafe]) {
+          result.add(a)
+          result.add(b)
+        }
+      }
+      case BinaryComparison(Cast(a: Attribute, _), Cast(b: Attribute, _)) => {
+        if (!e.isInstanceOf[EqualNullSafe]) {
+          result.add(a)
+          result.add(b)
+        }
+      }
+      case BinaryComparison(a: Attribute, _) => if (!e.isInstanceOf[EqualNullSafe]) result.add(a)
+      case BinaryComparison(_, a: Attribute) => if (!e.isInstanceOf[EqualNullSafe]) result.add(a)
+      case BinaryComparison(Cast(a: Attribute, _), _) =>
+        if (!e.isInstanceOf[EqualNullSafe]) result.add(a)
+      case BinaryComparison(_, Cast(a: Attribute, _)) =>
+        if (!e.isInstanceOf[EqualNullSafe]) result.add(a)
+      case Not(child) => extract(child)
+      case _ =>
+    }
+    predicates.foreach { extract(_) }
+    result.toSet
+  }
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala

@@ -99,6 +99,13 @@ abstract class LogicalPlan extends QueryPlan[LogicalPlan] with Logging {
    */
   lazy val resolved: Boolean = expressions.forall(_.resolved) && childrenResolved
 
+  /**
+   * Returns true if the two plans are semantically equal. This should ignore state generated
+   * during planning to help the planning process.
+   * TODO: implement this as a pass that canonicalizes the plan tree instead?
+   */
+  def semanticEquals(other: LogicalPlan): Boolean = this == other
+
   override protected def statePrefix = if (!resolved) "'" else super.statePrefix
 
   /**

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala

@@ -122,11 +122,22 @@ case class Except(left: LogicalPlan, right: LogicalPlan) extends SetOperation(le
   override def output: Seq[Attribute] = left.output
 }
 
+object Join {
+  def apply(
+    left: LogicalPlan,
+    right: LogicalPlan,
+    joinType: JoinType,
+    condition: Option[Expression]): Join = {
+    Join(left, right, joinType, condition, None)
+  }
+}
+
 case class Join(
   left: LogicalPlan,
   right: LogicalPlan,
   joinType: JoinType,
-  condition: Option[Expression]) extends BinaryNode {
+  condition: Option[Expression],
+  generatedExpressions: Option[EquivalentExpressions]) extends BinaryNode {
 
   override def output: Seq[Attribute] = {
     joinType match {
@@ -152,6 +163,17 @@ case class Join(
       selfJoinResolved &&
       condition.forall(_.dataType == BooleanType)
   }
+
+  override  def simpleString: String = s"$nodeName $joinType, $condition".trim
+
+  override def semanticEquals(other: LogicalPlan): Boolean = {
+    other match {
+      case Join (l, r, joinType, condition, _) => {
+        l == left && r == right && this.joinType == joinType && this.condition == condition
+      }
+      case _ => false
+    }
+  }
 }
 
 /**