[SPARK-12957][SQL] Initial support for constraint propagation in SparkSQL

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala

@@ -17,15 +17,30 @@
 
 package org.apache.spark.sql.catalyst.plans
 
-import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeSet, Expression, VirtualColumn}
+import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.trees.TreeNode
 import org.apache.spark.sql.types.{DataType, StructType}
 
-abstract class QueryPlan[PlanType <: TreeNode[PlanType]] extends TreeNode[PlanType] {
+abstract class QueryPlan[PlanType <: TreeNode[PlanType]]
+  extends TreeNode[PlanType] with PredicateHelper {
   self: PlanType =>
 
   def output: Seq[Attribute]
 
+  /**
+   * Extracts the output property from a given child.
+   */
+  def extractConstraintsFromChild(child: QueryPlan[PlanType]): Set[Expression] = {
+    child.constraints.filter(_.references.subsetOf(outputSet))
+  }
+
+  /**
+   * An sequence of expressions that describes the data property of the output rows of this
+   * operator. For example, if the output of this operator is column `a`, an example `constraints`
+   * can be `Set(a > 10, a < 20)`.
+   */
+  def constraints: Set[Expression] = Set.empty
+
   /**
    * Returns the set of attributes that are output by this node.
    */

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

@@ -301,10 +301,14 @@ abstract class LeafNode extends LogicalPlan {
 /**
  * A logical plan node with single child.
  */
-abstract class UnaryNode extends LogicalPlan {
+abstract class UnaryNode extends LogicalPlan with PredicateHelper {
   def child: LogicalPlan
 
   override def children: Seq[LogicalPlan] = child :: Nil
+
+  override def constraints: Set[Expression] = {
+    extractConstraintsFromChild(child)
+  }
 }
 
 /**

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

@@ -88,13 +88,33 @@ case class Generate(
 
 case class Filter(condition: Expression, child: LogicalPlan) extends UnaryNode {
   override def output: Seq[Attribute] = child.output
+
+  override def constraints: Set[Expression] = {
+    val newConstraint = splitConjunctivePredicates(condition).filter(
+      _.references.subsetOf(outputSet)).toSet
+    newConstraint.union(extractConstraintsFromChild(child))
+  }
 }
 
 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 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))
+    extractConstraintsFromChild(right).map(_ transform {
+      case a: Attribute => attributeRewrites(a)
+    })
+  }
 }
 
 private[sql] object SetOperation {
@@ -107,11 +127,17 @@ case class Intersect(left: LogicalPlan, right: LogicalPlan) extends SetOperation
     left.output.zip(right.output).map { case (leftAttr, rightAttr) =>
       leftAttr.withNullability(leftAttr.nullable && rightAttr.nullable)
     }
+
+  override def constraints: Set[Expression] = {
+    leftConstraints.union(rightConstraints)
+  }
 }
 
 case class Except(left: LogicalPlan, right: LogicalPlan) extends SetOperation(left, right) {
   /** We don't use right.output because those rows get excluded from the set. */
   override def output: Seq[Attribute] = left.output
+
+  override def constraints: Set[Expression] = leftConstraints
 }
 
 /** Factory for constructing new `Union` nodes. */
@@ -146,6 +172,26 @@ case class Union(children: Seq[LogicalPlan]) extends LogicalPlan {
     val sizeInBytes = children.map(_.statistics.sizeInBytes).sum
     Statistics(sizeInBytes = sizeInBytes)
   }
+
+  override def extractConstraintsFromChild(child: QueryPlan[LogicalPlan]): Set[Expression] = {
+    child.constraints.filter(_.references.subsetOf(child.outputSet))
+  }
+
+  def rewriteConstraints(
+      planA: LogicalPlan,
+      planB: LogicalPlan,
+      constraints: Set[Expression]): Set[Expression] = {
+    require(planA.output.size == planB.output.size)
+    val attributeRewrites = AttributeMap(planB.output.zip(planA.output))
+    constraints.map(_ transform {
+      case a: Attribute => attributeRewrites(a)
+    })
+  }
+
+  override def constraints: Set[Expression] = {
+    children.map(child => rewriteConstraints(children.head, child,
+      extractConstraintsFromChild(child))).reduce(_ intersect _)
+  }
 }
 
 case class Join(
@@ -169,6 +215,52 @@ case class Join(
     }
   }
 
+  override def constraints: Set[Expression] = {
+    var constraints = Set.empty[Expression]
+
+    // Currently we only propagate constraints if the condition consists of equality
+    // and ranges. For all other cases, we return an empty set of constraints
+    def extractIsNotNullConstraints(condition: Expression): Set[Expression] = {
+      splitConjunctivePredicates(condition).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)))
+      }
+      constraints
+    }
+
+    def extractIsNullConstraints(plan: LogicalPlan): Set[Expression] = {
+      constraints = constraints.union(plan.output.map(IsNull).toSet)
+      constraints
+    }
+
+    constraints = joinType match {
+      case Inner if condition.isDefined =>
+        extractConstraintsFromChild(left).union(extractConstraintsFromChild(right))
+          .union(extractIsNotNullConstraints(condition.get))
+      case LeftSemi if condition.isDefined =>
+        extractConstraintsFromChild(left).union(extractConstraintsFromChild(right))
+          .union(extractIsNotNullConstraints(condition.get))
+      case LeftOuter =>
+        extractConstraintsFromChild(left).union(extractIsNullConstraints(right))
+      case RightOuter =>
+        extractConstraintsFromChild(right).union(extractIsNullConstraints(left))
+      case FullOuter =>
+        extractIsNullConstraints(left).union(extractIsNullConstraints(right))
+      case _ =>
+        Set.empty
+    }
+
+    constraints.filter(_.references.subsetOf(outputSet))
+  }
+
   def selfJoinResolved: Boolean = left.outputSet.intersect(right.outputSet).isEmpty
 
   // Joins are only resolved if they don't introduce ambiguous expression ids.

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

@@ -0,0 +1,125 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.plans
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.catalyst.analysis._
+import org.apache.spark.sql.catalyst.dsl.expressions._
+import org.apache.spark.sql.catalyst.dsl.plans._
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.plans.logical._
+
+class ConstraintPropagationSuite extends SparkFunSuite {
+
+  private def resolveColumn(tr: LocalRelation, columnName: String): Expression =
+    tr.analyze.resolveQuoted(columnName, caseInsensitiveResolution).get
+
+  private def verifyConstraints(a: Set[Expression], b: Set[Expression]): Unit = {
+    assert(a.forall(i => b.map(_.semanticEquals(i)).reduce(_ || _)))
+    assert(b.forall(i => a.map(_.semanticEquals(i)).reduce(_ || _)))
+  }
+
+  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)
+    assert(tr.where('a.attr > 10).select('c.attr, 'b.attr).analyze.constraints.isEmpty)
+    verifyConstraints(tr.where('a.attr > 10).analyze.constraints, Set(resolveColumn(tr, "a") > 10))
+    verifyConstraints(tr.where('a.attr > 10).select('c.attr, 'a.attr).where('c.attr < 100)
+      .analyze.constraints, Set(resolveColumn(tr, "a") > 10, resolveColumn(tr, "c") < 100))
+  }
+
+  test("propagating constraints in union") {
+    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)
+    verifyConstraints(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)
+    verifyConstraints(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)
+    verifyConstraints(tr1.where('a.attr > 10).except(tr2.where('b.attr < 100)).analyze.constraints,
+      Set(resolveColumn(tr1, "a") > 10))
+  }
+
+  test("propagating constraints in inner join") {
+    val tr1 = LocalRelation('a.int, 'b.int, 'c.int).subquery('tr1)
+    val tr2 = LocalRelation('a.int, 'd.int, 'e.int).subquery('tr2)
+    verifyConstraints(tr1.where('a.attr > 10).join(tr2.where('d.attr < 100), Inner,
+      Some("tr1.a".attr === "tr2.a".attr)).analyze.constraints,
+      Set(tr1.resolveQuoted("a", caseInsensitiveResolution).get > 10,
+        tr2.resolveQuoted("d", caseInsensitiveResolution).get < 100,
+        IsNotNull(tr2.resolveQuoted("a", caseInsensitiveResolution).get),
+        IsNotNull(tr1.resolveQuoted("a", caseInsensitiveResolution).get)))
+  }
+
+  test("propagating constraints in left-semi join") {
+    val tr1 = LocalRelation('a.int, 'b.int, 'c.int).subquery('tr1)
+    val tr2 = LocalRelation('a.int, 'd.int, 'e.int).subquery('tr2)
+    verifyConstraints(tr1.where('a.attr > 10).join(tr2.where('d.attr < 100), LeftSemi,
+      Some("tr1.a".attr === "tr2.a".attr)).analyze.constraints,
+      Set(tr1.resolveQuoted("a", caseInsensitiveResolution).get > 10,
+        IsNotNull(tr1.resolveQuoted("a", caseInsensitiveResolution).get)))
+  }
+
+  test("propagating constraints in left-outer join") {
+    val tr1 = LocalRelation('a.int, 'b.int, 'c.int).subquery('tr1)
+    val tr2 = LocalRelation('a.int, 'd.int, 'e.int).subquery('tr2)
+    verifyConstraints(tr1.where('a.attr > 10).join(tr2.where('d.attr < 100), LeftOuter,
+      Some("tr1.a".attr === "tr2.a".attr)).analyze.constraints,
+      Set(tr1.resolveQuoted("a", caseInsensitiveResolution).get > 10,
+        IsNull(tr2.resolveQuoted("a", caseInsensitiveResolution).get),
+        IsNull(tr2.resolveQuoted("d", caseInsensitiveResolution).get),
+        IsNull(tr2.resolveQuoted("e", caseInsensitiveResolution).get)))
+  }
+
+  test("propagating constraints in right-outer join") {
+    val tr1 = LocalRelation('a.int, 'b.int, 'c.int).subquery('tr1)
+    val tr2 = LocalRelation('a.int, 'd.int, 'e.int).subquery('tr2)
+    verifyConstraints(tr1.where('a.attr > 10).join(tr2.where('d.attr < 100), RightOuter,
+      Some("tr1.a".attr === "tr2.a".attr)).analyze.constraints,
+      Set(tr2.resolveQuoted("d", caseInsensitiveResolution).get < 100,
+        IsNull(tr1.resolveQuoted("a", caseInsensitiveResolution).get),
+        IsNull(tr1.resolveQuoted("b", caseInsensitiveResolution).get),
+        IsNull(tr1.resolveQuoted("c", caseInsensitiveResolution).get)))
+  }
+
+  test("propagating constraints in full-outer join") {
+    val tr1 = LocalRelation('a.int, 'b.int, 'c.int).subquery('tr1)
+    val tr2 = LocalRelation('a.int, 'd.int, 'e.int).subquery('tr2)
+    verifyConstraints(tr1.where('a.attr > 10).join(tr2.where('d.attr < 100), FullOuter,
+      Some("tr1.a".attr === "tr2.a".attr)).analyze.constraints,
+      Set(IsNull(tr1.resolveQuoted("a", caseInsensitiveResolution).get),
+        IsNull(tr1.resolveQuoted("b", caseInsensitiveResolution).get),
+        IsNull(tr1.resolveQuoted("c", caseInsensitiveResolution).get),
+        IsNull(tr2.resolveQuoted("a", caseInsensitiveResolution).get),
+        IsNull(tr2.resolveQuoted("d", caseInsensitiveResolution).get),
+        IsNull(tr2.resolveQuoted("e", caseInsensitiveResolution).get)))
+  }
+}