Bug fix: Attribute nullability for set operations

Author: Hari Kishore Chaparala

src/it/scala/io/github/spark_redshift_community/spark/redshift/pushdown/PushdownLogicalPlanOperatorSuite.scala

@@ -1225,6 +1225,47 @@ abstract class PushdownLogicalPlanOperatorSuite extends IntegrationPushdownSuite
     doTest(sqlContext, testUnion13)
   }
 
+  test("Test UNION ALL attribute nullability") {
+    withTwoTempRedshiftTables("tableA", "tableB") { (tableA, tableB) =>
+      redshiftWrapper.executeUpdate(conn,
+        s"create table $tableA (id INT not null, category varchar(50))")
+      redshiftWrapper.executeUpdate(conn,
+        s"create table $tableB (id INT, category varchar(50))")
+
+      redshiftWrapper.executeUpdate(conn,
+        s"insert into $tableA VALUES (1, 'A'), (1, 'B'), (2, 'A')")
+      redshiftWrapper.executeUpdate(conn,
+        s"insert into $tableB VALUES (1, 'C'), (1, 'D'), (NULL, 'A')")
+
+      read.option("dbtable", tableA).load.createOrReplaceTempView(tableA)
+      read.option("dbtable", tableB).load.createOrReplaceTempView(tableB)
+
+      val strQuery =
+        s"select id, count(*) as cnt from " +
+        s"(select id, category from $tableA union all select id, category from $tableB) " +
+        s"group by rollup (id) order by id, cnt"
+
+      // If the nullability is wrong, the connector will return a row of (0, 1) instead of (null, 1)
+      // because Spark will misapply the non-nullability of the first table to the second table and
+      // convert the null column value into a zero [Redshift-87788]
+      checkAnswer(
+        sqlContext.sql(strQuery),
+        Seq(Row(null, 1),
+            Row(null, 6),
+            Row(1, 4),
+            Row(2, 1))
+      )
+
+      // We don't expect the group by rollup expression to be pushed down.
+      checkSqlStatement(
+        s"""( SELECT ( "SQ_0"."ID" ) AS "SQ_1_COL_0" FROM
+          | ( SELECT * FROM "PUBLIC"."$tableA" AS "RCQ_ALIAS" ) AS "SQ_0" ) UNION ALL
+          | ( SELECT ( "SQ_0"."ID" ) AS "SQ_1_COL_0" FROM
+          | ( SELECT * FROM "PUBLIC"."$tableB" AS "RCQ_ALIAS" ) AS "SQ_0" )""".stripMargin
+      )
+    }
+  }
+
   // No push down for except
   test("Test EXCEPT logical plan operator") {
     // "Column name" and result set

src/main/scala/io/github/spark_redshift_community/spark/redshift/pushdown/querygeneration/RedshiftQuery.scala

@@ -789,17 +789,53 @@ case class SetQuery(children: Seq[LogicalPlan],
     new QueryBuilder(child).treeRoot
   }
 
-  if(queries.contains(null)) {
+  if (queries.contains(null)) {
     throw new RedshiftPushdownUnsupportedException(
       RedshiftFailMessage.FAIL_PUSHDOWN_STATEMENT,
       setOperation,
       "Not all " + setOperation + " children query supported",
       false)
   }
+
+  private val childrenOutputs: Seq[Seq[Attribute]] = queries.map(_.helper.output)
+
+  private val resolvedAttributes: Seq[Attribute] = setOperation match {
+    case "UNION ALL" =>
+      // For UNION ALL: resolved nullable = OR of all child nullabilities
+      childrenOutputs.head.indices.map { i =>
+        val childNullabilities = childrenOutputs.map(_.apply(i).nullable)
+        val resolvedNullable = childNullabilities.contains(true)
+        val baseAttr = childrenOutputs.head(i)
+        baseAttr.withNullability(resolvedNullable)
+      }
+
+    case "EXCEPT" =>
+      // For EXCEPT: resolved nullability = from the first query
+      childrenOutputs.head
+
+    case "INTERSECT" =>
+      // For INTERSECT: resolved nullable = AND of all child nullabilities
+      // If any child is non-nullable (false), resolved is non-nullable (false).
+      childrenOutputs.head.indices.map { i =>
+        val childNullabilities = childrenOutputs.map(_.apply(i).nullable)
+        val resolvedNullable = childNullabilities.forall(_ == true)
+        val baseAttr = childrenOutputs.head(i)
+        baseAttr.withNullability(resolvedNullable)
+      }
+
+    case other =>
+      throw new RedshiftPushdownUnsupportedException(
+        "Unsupported set query pushdown",
+        other,
+        "Only UNION ALL, INTERSECT and EXCEPT are supported",
+        true
+      )
+  }
+
   override val helper: QueryHelper =
     QueryHelper(
       children = queries,
-      outputAttributes = Some(queries.head.helper.output),
+      outputAttributes = Some(resolvedAttributes),
       alias = alias,
       visibleAttributeOverride =
         Some(queries.foldLeft(Seq.empty[Attribute])((x, y) => x ++ y.helper.output).map(