Fix NULL handling for aggregation (#130)

* Modify COUNT and SUM to correctly handle NULL values

* Change average to support NULL values

* Fix

Author: wzheng

src/main/scala/edu/berkeley/cs/rise/opaque/Utils.scala

@@ -1196,30 +1196,47 @@ object Utils extends Logging {
       case avg @ Average(child) =>
         val sum = avg.aggBufferAttributes(0)
         val count = avg.aggBufferAttributes(1)
+        val dataType = child.dataType
+
+        val sumInitValue = child.nullable match {
+          case true => Literal.create(null, dataType)
+          case false => Cast(Literal(0), dataType)
+        }
+        val sumExpr = child.nullable match {
+          case true => If(IsNull(child), sum, If(IsNull(sum), Cast(child, dataType), Add(sum, Cast(child, dataType))))
+          case false => Add(sum, Cast(child, dataType))
+        }
+        val countExpr = If(IsNull(child), count, Add(count, Literal(1L)))
 
-        // TODO: support aggregating null values
         // TODO: support DecimalType to match Spark SQL behavior
         tuix.AggregateExpr.createAggregateExpr(
           builder,
           tuix.AggregateExpr.createInitialValuesVector(
             builder,
             Array(
-              /* sum = */ flatbuffersSerializeExpression(builder, Literal(0.0), input),
+              /* sum = */ flatbuffersSerializeExpression(builder, sumInitValue, input),
               /* count = */ flatbuffersSerializeExpression(builder, Literal(0L), input))),
           tuix.AggregateExpr.createUpdateExprsVector(
             builder,
             Array(
               /* sum = */ flatbuffersSerializeExpression(
-                builder, Add(sum, Cast(child, DoubleType)), concatSchema),
+                builder, sumExpr, concatSchema),
               /* count = */ flatbuffersSerializeExpression(
-                builder, Add(count, Literal(1L)), concatSchema))),
+                builder, countExpr, concatSchema))),
           flatbuffersSerializeExpression(
-            builder, Divide(sum, Cast(count, DoubleType)), aggSchema))
+            builder, Divide(Cast(sum, DoubleType), Cast(count, DoubleType)), aggSchema))
 
       case c @ Count(children) =>
         val count = c.aggBufferAttributes(0)
+        // COUNT(*) should count NULL values
+        // COUNT(expr) should return the number or rows for which the supplied expressions are non-NULL
+
+        val nullableChildren = children.filter(_.nullable)
+        val countExpr = nullableChildren.isEmpty match {
+          case true => Add(count, Literal(1L))
+          case false => If(nullableChildren.map(IsNull).reduce(Or), count, Add(count, Literal(1L)))
+        }
 
-        // TODO: support skipping null values
         tuix.AggregateExpr.createAggregateExpr(
           builder,
           tuix.AggregateExpr.createInitialValuesVector(
@@ -1230,7 +1247,7 @@ object Utils extends Logging {
             builder,
             Array(
               /* count = */ flatbuffersSerializeExpression(
-                builder, Add(count, Literal(1L)), concatSchema))),
+                builder, countExpr, concatSchema))),
           flatbuffersSerializeExpression(
             builder, count, aggSchema))
 
@@ -1316,22 +1333,31 @@ object Utils extends Logging {
 
       case s @ Sum(child) =>
         val sum = s.aggBufferAttributes(0)
-
         val sumDataType = s.dataType
+        // If any value is not NULL, return a non-NULL value
+        // If all values are NULL, return NULL
+
+        val initValue = child.nullable match {
+          case true => Literal.create(null, sumDataType)
+          case false => Cast(Literal(0), sumDataType)
+        }
+        val sumExpr = child.nullable match {
+          case true => If(IsNull(child), sum, If(IsNull(sum), Cast(child, sumDataType), Add(sum, Cast(child, sumDataType))))
+          case false => Add(sum, Cast(child, sumDataType))
+        }
 
-        // TODO: support aggregating null values
         tuix.AggregateExpr.createAggregateExpr(
           builder,
           tuix.AggregateExpr.createInitialValuesVector(
             builder,
             Array(
               /* sum = */ flatbuffersSerializeExpression(
-                builder, Cast(Literal(0), sumDataType), input))),
+                builder, initValue, input))),
           tuix.AggregateExpr.createUpdateExprsVector(
             builder,
             Array(
               /* sum = */ flatbuffersSerializeExpression(
-                builder, Add(sum, Cast(child, sumDataType)), concatSchema))),
+                builder, sumExpr, concatSchema))),
           flatbuffersSerializeExpression(
             builder, sum, aggSchema))
 

src/test/scala/edu/berkeley/cs/rise/opaque/OpaqueOperatorTests.scala

@@ -122,6 +122,30 @@ trait OpaqueOperatorTests extends FunSuite with BeforeAndAfterAll { self =>
     }
   }
 
+  /** Modified from https://stackoverflow.com/questions/33193958/change-nullable-property-of-column-in-spark-dataframe
+    * and https://stackoverflow.com/questions/32585670/what-is-the-best-way-to-define-custom-methods-on-a-dataframe
+    * Set nullable property of column.
+    * @param cn is the column name to change
+    * @param nullable is the flag to set, such that the column is  either nullable or not
+    */
+  object ExtraDFOperations {
+    implicit class AlternateDF(df : DataFrame) {
+      def setNullableStateOfColumn(cn: String, nullable: Boolean) : DataFrame = {
+        // get schema
+        val schema = df.schema
+        // modify [[StructField] with name `cn`
+        val newSchema = StructType(schema.map {
+          case StructField( c, t, _, m) if c.equals(cn) => StructField( c, t, nullable = nullable, m)
+          case y: StructField => y
+        })
+        // apply new schema
+        df.sqlContext.createDataFrame( df.rdd, newSchema )
+      }
+    }
+  }
+
+  import ExtraDFOperations._
+
   testAgainstSpark("Interval SQL") { securityLevel =>
     val data = Seq(Tuple2(1, new java.sql.Date(new java.util.Date().getTime())))
     val df = makeDF(data, securityLevel, "index", "time")
@@ -375,17 +399,28 @@ trait OpaqueOperatorTests extends FunSuite with BeforeAndAfterAll { self =>
   }
 
   testAgainstSpark("aggregate average") { securityLevel =>
-    val data = for (i <- 0 until 256) yield (i, abc(i), i.toDouble)
+    val data = (0 until 256).map{ i =>
+      if (i % 3 == 0 || (i + 1) % 6 == 0)
+        (i, abc(i), None)
+      else
+        (i, abc(i), Some(i.toDouble))
+    }.toSeq
     val words = makeDF(data, securityLevel, "id", "category", "price")
+    words.setNullableStateOfColumn("price", true)
 
-    words.groupBy("category").agg(avg("price").as("avgPrice"))
-      .collect.sortBy { case Row(category: String, _) => category }
+    val result = words.groupBy("category").agg(avg("price").as("avgPrice"))
+    result.collect.sortBy { case Row(category: String, _) => category }
   }
 
   testAgainstSpark("aggregate count") { securityLevel =>
-    val data = for (i <- 0 until 256) yield (i, abc(i), 1)
+    val data = (0 until 256).map{ i =>
+      if (i % 3 == 0 || (i + 1) % 6 == 0)
+        (i, abc(i), None)
+      else 
+        (i, abc(i), Some(i))
+    }.toSeq
     val words = makeDF(data, securityLevel, "id", "category", "price")
-
+    words.setNullableStateOfColumn("price", true)
     words.groupBy("category").agg(count("category").as("itemsInCategory"))
       .collect.sortBy { case Row(category: String, _) => category }
   }
@@ -423,8 +458,15 @@ trait OpaqueOperatorTests extends FunSuite with BeforeAndAfterAll { self =>
   }
 
   testAgainstSpark("aggregate sum") { securityLevel =>
-    val data = for (i <- 0 until 256) yield (i, abc(i), 1)
+    val data = (0 until 256).map{ i =>
+      if (i % 3 == 0 || i % 4 == 0)
+        (i, abc(i), None)
+      else
+        (i, abc(i), Some(i.toDouble))
+    }.toSeq
+
     val words = makeDF(data, securityLevel, "id", "word", "count")
+    words.setNullableStateOfColumn("count", true)
 
     words.groupBy("word").agg(sum("count").as("totalCount"))
       .collect.sortBy { case Row(word: String, _) => word }