[SPARK-21417][SQL] Infer join conditions using propagated constraints

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

@@ -0,0 +1,66 @@
+/*
+ * 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.expressions
+
+import scala.collection.mutable
+
+import org.apache.spark.sql.catalyst.expressions.EquivalentExpressionMap.SemanticallyEqualExpr
+
+/**
+ * A class that allows you to map an expression into a set of equivalent expressions. The keys are
+ * handled based on their semantic meaning and ignoring cosmetic differences. The values are
+ * represented as [[ExpressionSet]]s.
+ *
+ * The underlying representation of keys depends on the [[Expression.semanticHash]] and
+ * [[Expression.semanticEquals]] methods.
+ *
+ * {{{
+ *   val map = new EquivalentExpressionMap()
+ *
+ *   map.put(1 + 2, a)
+ *   map.put(rand(), b)
+ *
+ *   map.get(2 + 1) => Set(a) // 1 + 2 and 2 + 1 are semantically equivalent
+ *   map.get(1 + 2) => Set(a) // 1 + 2 and 2 + 1 are semantically equivalent
+ *   map.get(rand()) => Set() // non-deterministic expressions are not equivalent
+ * }}}
+ */
+class EquivalentExpressionMap {
+
+  private val equivalenceMap = mutable.HashMap.empty[SemanticallyEqualExpr, ExpressionSet]
+
+  def put(expression: Expression, equivalentExpression: Expression): Unit = {
+    val equivalentExpressions = equivalenceMap.getOrElseUpdate(expression, ExpressionSet.empty)
+    equivalenceMap(expression) = equivalentExpressions + equivalentExpression
+  }
+
+  def get(expression: Expression): Set[Expression] =
+    equivalenceMap.getOrElse(expression, ExpressionSet.empty)
+}
+
+object EquivalentExpressionMap {
+
+  private implicit class SemanticallyEqualExpr(val expr: Expression) {
+    override def equals(obj: Any): Boolean = obj match {
+      case other: SemanticallyEqualExpr => expr.semanticEquals(other.expr)
+      case _ => false
+    }
+
+    override def hashCode: Int = expr.semanticHash()
+  }
+}

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

@@ -27,6 +27,8 @@ object ExpressionSet {
     expressions.foreach(set.add)
     set
   }
+
+  val empty: ExpressionSet = ExpressionSet(Nil)
 }
 
 /**

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

@@ -18,6 +18,7 @@
 package org.apache.spark.sql.catalyst.optimizer
 
 import scala.annotation.tailrec
+import scala.collection.mutable
 
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.planning.ExtractFiltersAndInnerJoins
@@ -160,8 +161,9 @@ object EliminateOuterJoin extends Rule[LogicalPlan] with PredicateHelper {
  * conditions would potentially shuffle children as child node's partitioning won't satisfy the JOIN
  * node's requirements which otherwise could have.
  *
- * For instance, if there is a CROSS join, where the left relation has 'a = 1' and the right
- * relation has 'b = 1', the rule infers 'a = b' as a join predicate.
+ * For instance, given a CROSS join with the constraint 'a = 1' from the left child and the
+ * constraint 'b = 1' from the right child, this rule infers a new join predicate 'a = b' and
+ * converts it to an Inner join.
  */
 object EliminateCrossJoin extends Rule[LogicalPlan] with PredicateHelper {
 
@@ -174,7 +176,7 @@ object EliminateCrossJoin extends Rule[LogicalPlan] with PredicateHelper {
   }
 
   private def eliminateCrossJoin(plan: LogicalPlan): LogicalPlan = plan transform {
-    case join@Join(leftPlan, rightPlan, Cross, None) =>
+    case join @ Join(leftPlan, rightPlan, Cross, None) =>
       val leftConstraints = join.constraints.filter(_.references.subsetOf(leftPlan.outputSet))
       val rightConstraints = join.constraints.filter(_.references.subsetOf(rightPlan.outputSet))
       val inferredJoinPredicates = inferJoinPredicates(leftConstraints, rightConstraints)
@@ -184,67 +186,29 @@ object EliminateCrossJoin extends Rule[LogicalPlan] with PredicateHelper {
 
   private def inferJoinPredicates(
       leftConstraints: Set[Expression],
-      rightConstraints: Set[Expression]): Set[EqualTo] = {
-
-    // iterate through the left constraints and build a hash map that points semantically
-    // equivalent expressions into attributes
-    val emptyEquivalenceMap = Map.empty[SemanticExpression, Set[Attribute]]
-    val equivalenceMap = leftConstraints.foldLeft(emptyEquivalenceMap) { case (map, constraint) =>
-      constraint match {
-        case EqualTo(attr: Attribute, expr: Expression) =>
-          updateEquivalenceMap(map, attr, expr)
-        case EqualTo(expr: Expression, attr: Attribute) =>
-          updateEquivalenceMap(map, attr, expr)
-        case _ => map
-      }
-    }
-
-    // iterate through the right constraints and infer join conditions using the equivalence map
-    rightConstraints.foldLeft(Set.empty[EqualTo]) { case (joinConditions, constraint) =>
-      constraint match {
-        case EqualTo(attr: Attribute, expr: Expression) =>
-          appendJoinConditions(attr, expr, equivalenceMap, joinConditions)
-        case EqualTo(expr: Expression, attr: Attribute) =>
-          appendJoinConditions(attr, expr, equivalenceMap, joinConditions)
-        case _ => joinConditions
-      }
-    }
-  }
+      rightConstraints: Set[Expression]): mutable.Set[EqualTo] = {
 
-  private def updateEquivalenceMap(
-      equivalenceMap: Map[SemanticExpression, Set[Attribute]],
-      attr: Attribute,
-      expr: Expression): Map[SemanticExpression, Set[Attribute]] = {
+    val equivalentExpressionMap = new EquivalentExpressionMap()
 
-    val equivalentAttrs = equivalenceMap.getOrElse(expr, Set.empty[Attribute])
-    if (equivalentAttrs.contains(attr)) {
-      equivalenceMap
-    } else {
-      equivalenceMap.updated(expr, equivalentAttrs + attr)
+    leftConstraints.foreach {
+      case EqualTo(attr: Attribute, expr: Expression) =>
+        equivalentExpressionMap.put(expr, attr)
+      case EqualTo(expr: Expression, attr: Attribute) =>
+        equivalentExpressionMap.put(expr, attr)
+      case _ =>
     }
-  }
 
-  private def appendJoinConditions(
-      attr: Attribute,
-      expr: Expression,
-      equivalenceMap: Map[SemanticExpression, Set[Attribute]],
-      joinConditions: Set[EqualTo]): Set[EqualTo] = {
+    val joinConditions = mutable.Set.empty[EqualTo]
 
-    equivalenceMap.get(expr) match {
-      case Some(equivalentAttrs) => joinConditions ++ equivalentAttrs.map(EqualTo(attr, _))
-      case None => joinConditions
+    rightConstraints.foreach {
+      case EqualTo(attr: Attribute, expr: Expression) =>
+        joinConditions ++= equivalentExpressionMap.get(expr).map(EqualTo(attr, _))
+      case EqualTo(expr: Expression, attr: Attribute) =>
+        joinConditions ++= equivalentExpressionMap.get(expr).map(EqualTo(attr, _))
+      case _ =>
     }
-  }
 
-  // the purpose of this class is to treat 'a === 1 and 1 === 'a as the same expressions
-  implicit class SemanticExpression(private val expr: Expression) {
-
-    override def hashCode(): Int = expr.semanticHash()
-
-    override def equals(other: Any): Boolean = other match {
-      case other: SemanticExpression => expr.semanticEquals(other.expr)
-      case _ => false
-    }
+    joinConditions
   }
 
 }

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressionMapSuite.scala

@@ -0,0 +1,56 @@
+/*
+ * 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.expressions
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.catalyst.dsl.expressions._
+
+class EquivalentExpressionMapSuite extends SparkFunSuite {
+
+  private val onePlusTwo = Literal(1) + Literal(2)
+  private val twoPlusOne = Literal(2) + Literal(1)
+  private val rand = Rand(10)
+
+  test("behaviour of the equivalent expression map") {
+    val equivalentExpressionMap = new EquivalentExpressionMap()
+    equivalentExpressionMap.put(onePlusTwo, 'a)
+    equivalentExpressionMap.put(Literal(1) + Literal(3), 'b)
+    equivalentExpressionMap.put(rand, 'c)
+
+    // 1 + 2 should be equivalent to 2 + 1
+    assertResult(ExpressionSet(Seq('a)))(equivalentExpressionMap.get(twoPlusOne))
+    // non-deterministic expressions should not be equivalent
+    assertResult(ExpressionSet.empty)(equivalentExpressionMap.get(rand))
+
+    // if the same (key, value) is added several times, the map still returns only one entry
+    equivalentExpressionMap.put(onePlusTwo, 'a)
+    equivalentExpressionMap.put(twoPlusOne, 'a)
+    assertResult(ExpressionSet(Seq('a)))(equivalentExpressionMap.get(twoPlusOne))
+
+    // get several equivalent attributes
+    equivalentExpressionMap.put(onePlusTwo, 'e)
+    assertResult(ExpressionSet(Seq('a, 'e)))(equivalentExpressionMap.get(onePlusTwo))
+    assertResult(2)(equivalentExpressionMap.get(onePlusTwo).size)
+
+    // several non-deterministic expressions should not be equivalent
+    equivalentExpressionMap.put(rand, 'd)
+    assertResult(ExpressionSet.empty)(equivalentExpressionMap.get(rand))
+    assertResult(0)(equivalentExpressionMap.get(rand).size)
+  }
+
+}

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/EliminateCrossJoinSuite.scala

@@ -19,7 +19,7 @@ package org.apache.spark.sql.catalyst.optimizer
 
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.dsl.plans._
-import org.apache.spark.sql.catalyst.expressions.{Cast, Expression, Literal, Not}
+import org.apache.spark.sql.catalyst.expressions.{Cast, Expression, Literal, Not, Rand}
 import org.apache.spark.sql.catalyst.plans.{Cross, Inner, JoinType, PlanTest}
 import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan}
 import org.apache.spark.sql.catalyst.rules.RuleExecutor
@@ -43,6 +43,7 @@ class EliminateCrossJoinSuite extends PlanTest {
       originalFilter = 'a === 1 && 'c === 1 && 'd === 1,
       originalJoinType = Cross,
       originalJoinCondition = None,
+      expectedFilter = None,
       expectedLeftRelationFilter = 'a === 1,
       expectedRightRelationFilter = 'c === 1 && 'd === 1,
       expectedJoinType = Inner,
@@ -54,6 +55,7 @@ class EliminateCrossJoinSuite extends PlanTest {
       originalFilter = 'a === 1 && 'b === 2 && 'd === 1,
       originalJoinType = Cross,
       originalJoinCondition = None,
+      expectedFilter = None,
       expectedLeftRelationFilter = 'a === 1 && 'b === 2,
       expectedRightRelationFilter = 'd === 1,
       expectedJoinType = Inner,
@@ -68,6 +70,7 @@ class EliminateCrossJoinSuite extends PlanTest {
       originalFilter = 'a === 1 && Literal(1) === 'd && 'd === 'a,
       originalJoinType = Cross,
       originalJoinCondition = None,
+      expectedFilter = None,
       expectedLeftRelationFilter = 'a === 1,
       expectedRightRelationFilter = Literal(1) === 'd,
       expectedJoinType = Cross,
@@ -80,6 +83,7 @@ class EliminateCrossJoinSuite extends PlanTest {
       originalFilter = 'a === Literal(1) * Literal(2) && Literal(2) * Literal(1) === 'c,
       originalJoinType = Cross,
       originalJoinCondition = None,
+      expectedFilter = None,
       expectedLeftRelationFilter = 'a === Literal(1) * Literal(2),
       expectedRightRelationFilter = Literal(2) * Literal(1) === 'c,
       expectedJoinType = Inner,
@@ -91,6 +95,7 @@ class EliminateCrossJoinSuite extends PlanTest {
       originalFilter = 'a === 1 && Literal(1) === 'a && 'c === 1,
       originalJoinType = Cross,
       originalJoinCondition = None,
+      expectedFilter = None,
       expectedLeftRelationFilter = 'a === 1 && Literal(1) === 'a,
       expectedRightRelationFilter = 'c === 1,
       expectedJoinType = Inner,
@@ -102,6 +107,7 @@ class EliminateCrossJoinSuite extends PlanTest {
       originalFilter = 'a === Cast("1", IntegerType) && 'c === Cast("1", IntegerType) && 'd === 1,
       originalJoinType = Cross,
       originalJoinCondition = None,
+      expectedFilter = None,
       expectedLeftRelationFilter = 'a === Cast("1", IntegerType),
       expectedRightRelationFilter = 'c === Cast("1", IntegerType) && 'd === 1,
       expectedJoinType = Inner,
@@ -114,6 +120,7 @@ class EliminateCrossJoinSuite extends PlanTest {
       originalFilter = (('a >= 1 && 'c === 1) || 'd === 10) && 'b === 10 && 'c === 1,
       originalJoinType = Cross,
       originalJoinCondition = None,
+      expectedFilter = None,
       expectedLeftRelationFilter = 'b === 10,
       expectedRightRelationFilter = 'c === 1,
       expectedJoinType = Cross,
@@ -125,6 +132,7 @@ class EliminateCrossJoinSuite extends PlanTest {
       originalFilter = 'a === 1 && 'c === 1 && Literal(1) === 'a && Literal(1) === 'c,
       originalJoinType = Cross,
       originalJoinCondition = None,
+      expectedFilter = None,
       expectedLeftRelationFilter = 'a === 1 && Literal(1) === 'a,
       expectedRightRelationFilter = 'c === 1 && Literal(1) === 'c,
       expectedJoinType = Inner,
@@ -137,6 +145,7 @@ class EliminateCrossJoinSuite extends PlanTest {
         originalFilter = 'a === 1 && 'c === 1 && 'd === 1,
         originalJoinType = Cross,
         originalJoinCondition = None,
+        expectedFilter = None,
         expectedLeftRelationFilter = 'a === 1,
         expectedRightRelationFilter = 'c === 1 && 'd === 1,
         expectedJoinType = Cross,
@@ -149,6 +158,7 @@ class EliminateCrossJoinSuite extends PlanTest {
       originalFilter = 'a >= 1 && 'c === 1 && 'd >= 1,
       originalJoinType = Cross,
       originalJoinCondition = None,
+      expectedFilter = None,
       expectedLeftRelationFilter = 'a >= 1,
       expectedRightRelationFilter = 'c === 1 && 'd >= 1,
       expectedJoinType = Cross,
@@ -160,6 +170,7 @@ class EliminateCrossJoinSuite extends PlanTest {
       originalFilter = Literal(1) === 'b && ('c === 1 || 'd === 1),
       originalJoinType = Cross,
       originalJoinCondition = None,
+      expectedFilter = None,
       expectedLeftRelationFilter = Literal(1) === 'b,
       expectedRightRelationFilter = 'c === 1 || 'd === 1,
       expectedJoinType = Cross,
@@ -171,6 +182,7 @@ class EliminateCrossJoinSuite extends PlanTest {
       originalFilter = Literal(1) === 'b && 'c === 1,
       originalJoinType = Cross,
       originalJoinCondition = Some('c === 'b),
+      expectedFilter = None,
       expectedLeftRelationFilter = Literal(1) === 'b,
       expectedRightRelationFilter = 'c === 1,
       expectedJoinType = Cross,
@@ -182,16 +194,30 @@ class EliminateCrossJoinSuite extends PlanTest {
       originalFilter = Not('a === 1) && 'd === 1,
       originalJoinType = Cross,
       originalJoinCondition = None,
+      expectedFilter = None,
       expectedLeftRelationFilter = Not('a === 1),
       expectedRightRelationFilter = 'd === 1,
       expectedJoinType = Cross,
       expectedJoinCondition = None)
   }
 
+  test("inability to detect join conditions (5)") {
+    checkJoinOptimization(
+      originalFilter = 'a === Rand(10) && 'b === 1 && 'd === Rand(10) && 'c === 3,
+      originalJoinType = Cross,
+      originalJoinCondition = None,
+      expectedFilter = Some('a === Rand(10) && 'd === Rand(10)),
+      expectedLeftRelationFilter = 'b === 1,
+      expectedRightRelationFilter = 'c === 3,
+      expectedJoinType = Cross,
+      expectedJoinCondition = None)
+  }
+
   private def checkJoinOptimization(
       originalFilter: Expression,
       originalJoinType: JoinType,
       originalJoinCondition: Option[Expression],
+      expectedFilter: Option[Expression],
       expectedLeftRelationFilter: Expression,
       expectedRightRelationFilter: Expression,
       expectedJoinType: JoinType,
@@ -204,7 +230,9 @@ class EliminateCrossJoinSuite extends PlanTest {
 
     val left = testRelation1.where(expectedLeftRelationFilter)
     val right = testRelation2.where(expectedRightRelationFilter)
-    val expectedQuery = left.join(right, expectedJoinType, expectedJoinCondition).analyze
+    val join = left.join(right, expectedJoinType, expectedJoinCondition)
+    val expectedQuery = expectedFilter.fold(join)(join.where(_)).analyze
+
     comparePlans(optimizedQuery, expectedQuery)
   }
 }