improve nullability inference for Aggregator

Author: koertkuipers

sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala

@@ -38,10 +38,11 @@ object TypedAggregateExpression {
       bufferSerializer.map(_.toAttribute))
 
     val outputEncoder = encoderFor[OUT]
-    val outputType = if (outputEncoder.flat) {
-      outputEncoder.schema.head.dataType
+    val (outputType, outputNullable) = if (outputEncoder.flat) {
+      val sf = outputEncoder.schema.head
+      (sf.dataType, sf.nullable)
     } else {
-      outputEncoder.schema
+      (outputEncoder.schema, true)
     }
 
     new TypedAggregateExpression(
@@ -51,7 +52,8 @@ object TypedAggregateExpression {
       bufferDeserializer,
       outputEncoder.serializer,
       outputEncoder.deserializer.dataType,
-      outputType)
+      outputType,
+      outputNullable)
   }
 }
 
@@ -65,9 +67,8 @@ case class TypedAggregateExpression(
     bufferDeserializer: Expression,
     outputSerializer: Seq[Expression],
     outputExternalType: DataType,
-    dataType: DataType) extends DeclarativeAggregate with NonSQLExpression {
-
-  override def nullable: Boolean = true
+    dataType: DataType,
+    nullable: Boolean) extends DeclarativeAggregate with NonSQLExpression {
 
   override def deterministic: Boolean = true
 

sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala

@@ -271,4 +271,13 @@ class DatasetAggregatorSuite extends QueryTest with SharedSQLContext {
       "RowAgg(org.apache.spark.sql.Row)")
     assert(df.groupBy($"j").agg(RowAgg.toColumn as "agg1").columns.last == "agg1")
   }
+
+  test("spark-15204 improve nullability inference for Aggregator") {
+    val ds1 = Seq(1, 3, 2, 5).toDS()
+    assert(ds1.select(typed.sum((i: Int) => i)).schema.head.nullable === false)
+    val ds2 = Seq(AggData(1, "a"), AggData(2, "a")).toDS()
+    assert(ds2.groupByKey(_.b).agg(SeqAgg.toColumn).schema(1).nullable === true)
+    val ds3 = sql("SELECT 'Some String' AS b, 1279869254 AS a").as[AggData]
+    assert(ds3.groupByKey(_.a).agg(NameAgg.toColumn).schema(1).nullable === true)
+  }
 }