Support filter at lower data type

Author: wangyum

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

@@ -68,21 +68,46 @@ trait ConstraintHelper {
         val candidateConstraints = binaryComparisons - eq
         inferredConstraints ++= replaceConstraints(candidateConstraints, l, r)
         inferredConstraints ++= replaceConstraints(candidateConstraints, r, l)
-      case eq @ EqualTo(l: Cast, r: Attribute) =>
-        inferredConstraints ++= replaceConstraints(binaryComparisons - eq, r, l)
-      case eq @ EqualTo(l: Attribute, r: Cast) =>
-        inferredConstraints ++= replaceConstraints(binaryComparisons - eq, l, r)
+      case eq @ EqualTo(l @ Cast(lc: Attribute, _, tz), r: Attribute) =>
+        val candidateConstraints = binaryComparisons - eq
+        val bridge = Cast(r, lc.dataType, tz)
+        inferredConstraints ++= replaceConstraints(candidateConstraints, r, l)
+        inferredConstraints ++= replaceConstraints(candidateConstraints, lc, bridge)
+      case eq @ EqualTo(l: Attribute, r @ Cast(rc: Attribute, _, tz)) =>
+        val candidateConstraints = binaryComparisons - eq
+        val bridge = Cast(l, rc.dataType, tz)
+        inferredConstraints ++= replaceConstraints(candidateConstraints, l, r)
+        inferredConstraints ++= replaceConstraints(candidateConstraints, rc, bridge)
       case _ => // No inference
     }
     inferredConstraints -- constraints
   }
 
+  private def replaceConstraint(
+      constraint: Expression,
+      source: Expression,
+      destination: Expression): Expression = constraint transform {
+    case e: Expression if e.semanticEquals(source) => destination
+  }
+
   private def replaceConstraints(
       constraints: Set[Expression],
       source: Expression,
-      destination: Expression): Set[Expression] = constraints.map(_ transform {
-    case e: Expression if e.semanticEquals(source) => destination
-  })
+      dest: Expression): Set[Expression] = {
+    constraints.map {
+      case b @ BinaryComparison(left, right) =>
+        (replaceConstraint(left, source, dest), replaceConstraint(right, source, dest)) match {
+          case (Cast(Cast(child, _, _), dt, _), replacedRight)
+            if dt == child.dataType && child.dataType == replacedRight.dataType =>
+            b.makeCopy(Array(child, replacedRight))
+          case (replacedLeft, Cast(Cast(child, _, _), dt, _))
+            if dt == child.dataType && child.dataType == replacedLeft.dataType =>
+            b.makeCopy(Array(replacedLeft, child))
+          case (replacedLeft, replacedRight) =>
+            b.makeCopy(Array(replacedLeft, replacedRight))
+        }
+    }
+  }
 
   /**
    * Infers a set of `isNotNull` constraints from null intolerant expressions as well as

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

@@ -24,7 +24,7 @@ import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules._
 import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.types.LongType
+import org.apache.spark.sql.types.{IntegerType, LongType}
 
 class InferFiltersFromConstraintsSuite extends PlanTest {
 
@@ -47,8 +47,8 @@ class InferFiltersFromConstraintsSuite extends PlanTest {
       y: LogicalPlan,
       expectedLeft: LogicalPlan,
       expectedRight: LogicalPlan,
-      joinType: JoinType) = {
-    val condition = Some("x.a".attr === "y.a".attr)
+      joinType: JoinType,
+      condition: Option[Expression] = Some("x.a".attr === "y.a".attr)) = {
     val originalQuery = x.join(y, joinType, condition).analyze
     val correctAnswer = expectedLeft.join(expectedRight, joinType, condition).analyze
     val optimized = Optimize.execute(originalQuery)
@@ -265,7 +265,22 @@ class InferFiltersFromConstraintsSuite extends PlanTest {
     testConstraintsAfterJoin(x, y, x.where(IsNotNull('a)), y, RightOuter)
   }
 
-  test("Constraints should be inferred from cast equality constraint") {
+  test("Constraints should be inferred from cast equality constraint(filter at lower data type)") {
+    val testRelation1 = LocalRelation('a.int)
+    val testRelation2 = LocalRelation('b.long)
+    val originalLeft = testRelation1.where('a === 1).subquery('left)
+    val originalRight = testRelation2.subquery('right)
+
+    val left = testRelation1.where(IsNotNull('a) && 'a === 1).subquery('left)
+    val right = testRelation2.where(IsNotNull('b) && 'b.cast(IntegerType) === 1).subquery('right)
+
+    Seq(Some("left.a".attr.cast(LongType) === "right.b".attr),
+      Some("right.b".attr === "left.a".attr.cast(LongType))).foreach { condition =>
+      testConstraintsAfterJoin(originalLeft, originalRight, left, right, Inner, condition)
+    }
+  }
+
+  test("Constraints should be inferred from cast equality constraint(filter at higher data type)") {
     val testRelation1 = LocalRelation('a.int)
     val testRelation2 = LocalRelation('b.long)
     val originalLeft = testRelation1.subquery('left)
@@ -276,9 +291,7 @@ class InferFiltersFromConstraintsSuite extends PlanTest {
 
     Seq(Some("left.a".attr.cast(LongType) === "right.b".attr),
       Some("right.b".attr === "left.a".attr.cast(LongType))).foreach { condition =>
-      val optimized = Optimize.execute(originalLeft.join(originalRight, Inner, condition).analyze)
-      val correctAnswer = left.join(right, Inner, condition).analyze
-      comparePlans(optimized, correctAnswer)
+      testConstraintsAfterJoin(originalLeft, originalRight, left, right, Inner, condition)
     }
   }
 }