Add withNullability to Attribute and use it to change nullabilities.

Author: ueshin

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala

@@ -52,6 +52,7 @@ case class UnresolvedAttribute(name: String) extends Attribute with trees.LeafNo
   override lazy val resolved = false
 
   override def newInstance = this
+  override def withNullability(newNullability: Boolean) = this
   override def withQualifiers(newQualifiers: Seq[String]) = this
 
   // Unresolved attributes are transient at compile time and don't get evaluated during execution.
@@ -95,6 +96,7 @@ case class Star(
   override lazy val resolved = false
 
   override def newInstance = this
+  override def withNullability(newNullability: Boolean) = this
   override def withQualifiers(newQualifiers: Seq[String]) = this
 
   def expand(input: Seq[Attribute]): Seq[NamedExpression] = {

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

@@ -33,14 +33,16 @@ case class BoundReference(ordinal: Int, baseReference: Attribute)
 
   type EvaluatedType = Any
 
-  def nullable = baseReference.nullable
-  def dataType = baseReference.dataType
-  def exprId = baseReference.exprId
-  def qualifiers = baseReference.qualifiers
-  def name = baseReference.name
+  override def nullable = baseReference.nullable
+  override def dataType = baseReference.dataType
+  override def exprId = baseReference.exprId
+  override def qualifiers = baseReference.qualifiers
+  override def name = baseReference.name
 
-  def newInstance = BoundReference(ordinal, baseReference.newInstance)
-  def withQualifiers(newQualifiers: Seq[String]) =
+  override def newInstance = BoundReference(ordinal, baseReference.newInstance)
+  override def withNullability(newNullability: Boolean) =
+    BoundReference(ordinal, baseReference.withNullability(newNullability))
+  override def withQualifiers(newQualifiers: Seq[String]) =
     BoundReference(ordinal, baseReference.withQualifiers(newQualifiers))
 
   override def toString = s"$baseReference:$ordinal"

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

@@ -57,6 +57,7 @@ abstract class NamedExpression extends Expression {
 abstract class Attribute extends NamedExpression {
   self: Product =>
 
+  def withNullability(newNullability: Boolean): Attribute
   def withQualifiers(newQualifiers: Seq[String]): Attribute
 
   def toAttribute = this
@@ -133,7 +134,7 @@ case class AttributeReference(name: String, dataType: DataType, nullable: Boolea
   /**
    * Returns a copy of this [[AttributeReference]] with changed nullability.
    */
-  def withNullability(newNullability: Boolean) = {
+  override def withNullability(newNullability: Boolean) = {
     if (nullable == newNullability) {
       this
     } else {

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

@@ -51,13 +51,7 @@ case class Generate(
       .map(a => generator.output.map(_.withQualifiers(a :: Nil)))
       .getOrElse(generator.output)
     if (join && outer) {
-      output.map {
-        case attr if !attr.resolved => attr
-        case attr if !attr.nullable =>
-          AttributeReference(
-            attr.name, attr.dataType, nullable = true)(attr.exprId, attr.qualifiers)
-        case attr => attr
-      }
+      output.map(_.withNullability(true))
     } else {
       output
     }
@@ -93,26 +87,17 @@ case class Join(
   condition: Option[Expression]) extends BinaryNode {
 
   override def references = condition.map(_.references).getOrElse(Set.empty)
-  override def output = {
-    def nullabilize(output: Seq[Attribute]) = {
-      output.map {
-        case attr if !attr.resolved => attr
-        case attr if !attr.nullable =>
-          AttributeReference(
-            attr.name, attr.dataType, nullable = true)(attr.exprId, attr.qualifiers)
-        case attr => attr
-      }
-    }
 
+  override def output = {
     joinType match {
       case LeftSemi =>
         left.output
       case LeftOuter =>
-        left.output ++ nullabilize(right.output)
+        left.output ++ right.output.map(_.withNullability(true))
       case RightOuter =>
-        nullabilize(left.output) ++ right.output
+        left.output.map(_.withNullability(true)) ++ right.output
       case FullOuter =>
-        nullabilize(left.output) ++ nullabilize(right.output)
+        left.output.map(_.withNullability(true)) ++ right.output.map(_.withNullability(true))
       case _ =>
         left.output ++ right.output
     }

sql/core/src/main/scala/org/apache/spark/sql/execution/Generate.scala

@@ -41,12 +41,7 @@ case class Generate(
 
   protected def generatorOutput: Seq[Attribute] = {
     if (join && outer) {
-      generator.output.map {
-        case attr if !attr.nullable =>
-          AttributeReference(
-            attr.name, attr.dataType, nullable = true)(attr.exprId, attr.qualifiers)
-        case attr => attr
-      }
+      generator.output.map(_.withNullability(true))
     } else {
       generator.output
     }

sql/core/src/main/scala/org/apache/spark/sql/execution/joins.scala

@@ -319,23 +319,14 @@ case class BroadcastNestedLoopJoin(
 
   override def otherCopyArgs = sqlContext :: Nil
 
-  def output = {
-    def nullabilize(output: Seq[Attribute]) = {
-      output.map {
-        case attr if !attr.nullable =>
-          AttributeReference(
-            attr.name, attr.dataType, nullable = true)(attr.exprId, attr.qualifiers)
-        case attr => attr
-      }
-    }
-
+  override def output = {
     joinType match {
       case LeftOuter =>
-        left.output ++ nullabilize(right.output)
+        left.output ++ right.output.map(_.withNullability(true))
       case RightOuter =>
-        nullabilize(left.output) ++ right.output
+        left.output.map(_.withNullability(true)) ++ right.output
       case FullOuter =>
-        nullabilize(left.output) ++ nullabilize(right.output)
+        left.output.map(_.withNullability(true)) ++ right.output.map(_.withNullability(true))
       case _ =>
         left.output ++ right.output
     }