init import.

apache · Mar 17, 2016 · c2f9b05 · c2f9b05
1 parent 8ef3399
commit c2f9b05
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 0 deletions.
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/planning/patterns.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/planning/patterns.scala
@@ -202,3 +202,25 @@ object Unions {
     }
   }
 }
+
+/**
+ * A pattern that finds the original expression from a sequence of casts.
+ */
+object Casts {
+  def unapply(expr: Expression): Option[Attribute] = expr match {
+    case c: Cast => collectCasts(expr)
+    case _ => None
+  }
+
+  private def collectCasts(e: Expression): Option[Attribute] = {
+    if (e.isInstanceOf[Cast]) {
+      collectCasts(e.children(0))
+    } else {
+      if (e.isInstanceOf[Attribute]) {
+        Some(e.asInstanceOf[Attribute])
+      } else {
+        None
+      }
+    }
+  }
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala
@@ -18,6 +18,7 @@
 package org.apache.spark.sql.catalyst.plans
 
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.planning.Casts
 import org.apache.spark.sql.catalyst.trees.TreeNode
 import org.apache.spark.sql.types.{DataType, StructType}
 
@@ -36,6 +37,13 @@ abstract class QueryPlan[PlanType <: QueryPlan[PlanType]] extends TreeNode[PlanT
       .union(constructIsNotNullConstraints(constraints))
       .filter(constraint =>
         constraint.references.nonEmpty && constraint.references.subsetOf(outputSet))
+      .map(_.transform {
+        case n @ IsNotNull(c) =>
+          c match {
+            case Casts(a) if outputSet.contains(a) => IsNotNull(a)
+            case _ => n
+          }
+      })
   }
 
   /**