[SPARK-11949][SQL] Check bitmasks to set nullable property

viirya · yhuai · commit 3c4938e26185 · 2015-12-01T21:51:47.000-08:00
Following up apache#10038. We can use bitmasks to determine which grouping expressions need to be set as nullable. cc yhuai Author: Liang-Chi Hsieh <viirya@appier.com> Closes apache#10067 from viirya/fix-cube-following. (cherry picked from commit 0f37d1d) Signed-off-by: Yin Huai <yhuai@databricks.com>
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
@@ -224,10 +224,15 @@ class Analyzer(
           case other => Alias(other, other.toString)()
         }
 
-        // TODO: We need to use bitmasks to determine which grouping expressions need to be
-        // set as nullable. For example, if we have GROUPING SETS ((a,b), a), we do not need
-        // to change the nullability of a.
-        val attributeMap = groupByAliases.map(a => (a -> a.toAttribute.withNullability(true))).toMap
+        val nonNullBitmask = x.bitmasks.reduce(_ & _)
+
+        val attributeMap = groupByAliases.zipWithIndex.map { case (a, idx) =>
+          if ((nonNullBitmask & 1 << idx) == 0) {
+            (a -> a.toAttribute.withNullability(true))
+          } else {
+            (a -> a.toAttribute)
+          }
+        }.toMap
 
         val aggregations: Seq[NamedExpression] = x.aggregations.map {
           // If an expression is an aggregate (contains a AggregateExpression) then we dont change
