Infer additional constraints from attribute equality

Author: sameeragarwal

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

@@ -32,6 +32,7 @@ abstract class QueryPlan[PlanType <: QueryPlan[PlanType]] extends TreeNode[PlanT
    */
   protected def getRelevantConstraints(constraints: Set[Expression]): Set[Expression] = {
     constraints
+      .union(inferAdditionalConstraints(constraints))
       .union(constructIsNotNullConstraints(constraints))
       .filter(constraint =>
         constraint.references.nonEmpty && constraint.references.subsetOf(outputSet))
@@ -61,6 +62,25 @@ abstract class QueryPlan[PlanType <: QueryPlan[PlanType]] extends TreeNode[PlanT
     }.foldLeft(Set.empty[Expression])(_ union _.toSet)
   }
 
+  /**
+   * Infers an additional set of constraints from a given set of equality constraints.
+   * For e.g., if an operator has constraints of the form (`a = 5`, `a = b`), this returns an
+   * additional constraint of the form `b = 5`
+   */
+  private def inferAdditionalConstraints(constraints: Set[Expression]): Set[Expression] = {
+    constraints.map {
+      case eq @ EqualTo(l: Attribute, r: Attribute) =>
+        (constraints -- Set(eq)).map(_ transform {
+          case a: Attribute if a.semanticEquals(l) => r
+        }).union(
+          (constraints -- Set(eq)).map(_ transform {
+            case a: Attribute if a.semanticEquals(r) => l
+          }))
+      case _ =>
+        Set.empty[Expression]
+    }.foldLeft(Set.empty[Expression])(_ union _) -- constraints
+  }
+
   /**
    * An [[ExpressionSet]] that contains invariants about the rows output by this operator. For
    * example, if this set contains the expression `a = 2` then that expression is guaranteed to

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

@@ -158,6 +158,7 @@ class ConstraintPropagationSuite extends SparkFunSuite {
         tr2.resolveQuoted("d", caseInsensitiveResolution).get < 100,
         tr1.resolveQuoted("a", caseInsensitiveResolution).get ===
           tr2.resolveQuoted("a", caseInsensitiveResolution).get,
+        tr2.resolveQuoted("a", caseInsensitiveResolution).get > 10,
         IsNotNull(tr2.resolveQuoted("a", caseInsensitiveResolution).get),
         IsNotNull(tr1.resolveQuoted("a", caseInsensitiveResolution).get),
         IsNotNull(tr2.resolveQuoted("d", caseInsensitiveResolution).get))))
@@ -203,4 +204,17 @@ class ConstraintPropagationSuite extends SparkFunSuite {
       .join(tr2.where('d.attr < 100), FullOuter, Some("tr1.a".attr === "tr2.a".attr))
       .analyze.constraints.isEmpty)
   }
+
+  test("infer additional constraints in filters") {
+    val tr = LocalRelation('a.int, 'b.int, 'c.int)
+
+    verifyConstraints(tr
+      .where('a.attr > 10 && 'a.attr === 'b.attr)
+      .analyze.constraints,
+      ExpressionSet(Seq(resolveColumn(tr, "a") > 10,
+        resolveColumn(tr, "b") > 10,
+        resolveColumn(tr, "a") === resolveColumn(tr, "b"),
+        IsNotNull(resolveColumn(tr, "a")),
+        IsNotNull(resolveColumn(tr, "b")))))
+  }
 }