join propagation

Author: sameeragarwal

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

@@ -97,27 +97,24 @@ case class Filter(condition: Expression, child: LogicalPlan) extends UnaryNode {
 }
 
 abstract class SetOperation(left: LogicalPlan, right: LogicalPlan) extends BinaryNode {
+  final override lazy val resolved: Boolean =
+    childrenResolved &&
+      left.output.length == right.output.length &&
+      left.output.zip(right.output).forall { case (l, r) => l.dataType == r.dataType }
 
-  override def output: Seq[Attribute] =
-    left.output.zip(right.output).map { case (leftAttr, rightAttr) =>
-      leftAttr.withNullability(leftAttr.nullable || rightAttr.nullable)
-    }
+  override def extractConstraintsFromChild(child: QueryPlan[LogicalPlan]): Set[Expression] = {
+    child.constraints.filter(_.references.subsetOf(child.outputSet))
+  }
 
   protected def leftConstraints: Set[Expression] = extractConstraintsFromChild(left)
 
   protected def rightConstraints: Set[Expression] = {
     require(left.output.size == right.output.size)
     val attributeRewrites = AttributeMap(right.output.zip(left.output))
-    println(extractConstraintsFromChild(right), attributeRewrites)
     extractConstraintsFromChild(right).map(_ transform {
       case a: Attribute => attributeRewrites(a)
     })
   }
-
-  final override lazy val resolved: Boolean =
-    childrenResolved &&
-      left.output.length == right.output.length &&
-      left.output.zip(right.output).forall { case (l, r) => l.dataType == r.dataType }
 }
 
 private[sql] object SetOperation {
@@ -176,6 +173,10 @@ case class Union(children: Seq[LogicalPlan]) extends LogicalPlan {
     Statistics(sizeInBytes = sizeInBytes)
   }
 
+  override def extractConstraintsFromChild(child: QueryPlan[LogicalPlan]): Set[Expression] = {
+    child.constraints.filter(_.references.subsetOf(child.outputSet))
+  }
+
   def rewriteConstraints(
       planA: LogicalPlan,
       planB: LogicalPlan,
@@ -214,19 +215,43 @@ case class Join(
     }
   }
 
+  def extractNullabilityConstraintsFromJoinCondition(): Set[Expression] = {
+    var constraints = Set.empty[Expression]
+    if (condition.isDefined) {
+      splitConjunctivePredicates(condition.get).foreach {
+        case EqualTo(l, r) =>
+          constraints = constraints.union(Set(IsNotNull(l), IsNotNull(r)))
+        case GreaterThan(l, r) =>
+          constraints = constraints.union(Set(IsNotNull(l), IsNotNull(r)))
+        case GreaterThanOrEqual(l, r) =>
+          constraints = constraints.union(Set(IsNotNull(l), IsNotNull(r)))
+        case LessThan(l, r) =>
+          constraints = constraints.union(Set(IsNotNull(l), IsNotNull(r)))
+        case LessThanOrEqual(l, r) =>
+          constraints = constraints.union(Set(IsNotNull(l), IsNotNull(r)))
+      }
+    }
+    // Currently we only propagate constraints if the condition consists of equality
+    // and ranges. For all other cases, we return an empty set of constraints
+    constraints
+  }
+
   override def constraints: Set[Expression] = {
     joinType match {
       case Inner =>
+        extractConstraintsFromChild(left).union(extractConstraintsFromChild(right))
+          .union(extractNullabilityConstraintsFromJoinCondition())
       case LeftSemi =>
         extractConstraintsFromChild(left)
+          .union(extractNullabilityConstraintsFromJoinCondition())
       case LeftOuter =>
         extractConstraintsFromChild(left).union(extractConstraintsFromChild(right))
       case RightOuter =>
         extractConstraintsFromChild(left).union(extractConstraintsFromChild(right))
       case FullOuter =>
         extractConstraintsFromChild(left).union(extractConstraintsFromChild(right))
       case _ =>
-        extractConstraintsFromChild(left).union(extractConstraintsFromChild(right))
+        Set.empty
     }
   }
 

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/ConstraintPropagationSuite.scala

@@ -29,7 +29,7 @@ class ConstraintPropagationSuite extends SparkFunSuite {
   private def resolveColumn(tr: LocalRelation, columnName: String): Expression =
     tr.analyze.resolveQuoted(columnName, caseInsensitiveResolution).get
 
-  test("propagating constraints in filter") {
+  test("propagating constraints in filter/project") {
     val tr = LocalRelation('a.int, 'b.string, 'c.int)
     assert(tr.analyze.constraints.isEmpty)
     assert(tr.select('a.attr).analyze.constraints.isEmpty)
@@ -40,10 +40,26 @@ class ConstraintPropagationSuite extends SparkFunSuite {
   }
 
   test("propagating constraints in union") {
-    val tr1 = LocalRelation('a.int, 'b.string, 'c.int)
-    val tr2 = LocalRelation('a.int, 'b.string, 'c.int)
-    assert(tr1.analyze.constraints.isEmpty && tr2.analyze.constraints.isEmpty)
-    assert(tr1.where('a.attr > 10).unionAll(tr2.where('a.attr > 10))
-      .analyze.constraints == Set(resolveColumn(tr1, "a") > 10))
+    val tr1 = LocalRelation('a.int, 'b.int, 'c.int)
+    val tr2 = LocalRelation('d.int, 'e.int, 'f.int)
+    val tr3 = LocalRelation('g.int, 'h.int, 'i.int)
+    assert(tr1.where('a.attr > 10).unionAll(tr2.where('e.attr > 10)
+      .unionAll(tr3.where('i.attr > 10))).analyze.constraints.isEmpty)
+    assert(tr1.where('a.attr > 10).unionAll(tr2.where('d.attr > 10)
+      .unionAll(tr3.where('g.attr > 10))).analyze.constraints == Set(resolveColumn(tr1, "a") > 10))
+  }
+
+  test("propagating constraints in intersect") {
+    val tr1 = LocalRelation('a.int, 'b.int, 'c.int)
+    val tr2 = LocalRelation('a.int, 'b.int, 'c.int)
+    assert(tr1.where('a.attr > 10).intersect(tr2.where('b.attr < 100)).analyze.constraints ==
+      Set(resolveColumn(tr1, "a") > 10, resolveColumn(tr1, "b") < 100))
+  }
+
+  test("propagating constraints in except") {
+    val tr1 = LocalRelation('a.int, 'b.int, 'c.int)
+    val tr2 = LocalRelation('a.int, 'b.int, 'c.int)
+    assert(tr1.where('a.attr > 10).except(tr2.where('b.attr < 100)).analyze.constraints ==
+      Set(resolveColumn(tr1, "a") > 10))
   }
 }