use larger intermediate buffer for sum

Author: mgaido91

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala

@@ -45,7 +45,13 @@ case class Sum(child: Expression) extends DeclarativeAggregate with ImplicitCast
     case _ => DoubleType
   }
 
-  private lazy val sumDataType = resultType
+  private lazy val sumDataType = child.dataType match {
+    case LongType => DecimalType.BigIntDecimal
+    case _ => resultType
+  }
+
+  private lazy val castToResultType: (Expression) => Expression =
+    if (sumDataType == resultType) (e: Expression) => e else (e: Expression) => Cast(e, resultType)
 
   private lazy val sum = AttributeReference("sum", sumDataType)()
 
@@ -78,5 +84,5 @@ case class Sum(child: Expression) extends DeclarativeAggregate with ImplicitCast
     )
   }
 
-  override lazy val evaluateExpression: Expression = sum
+  override lazy val evaluateExpression: Expression = castToResultType(sum)
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/types/Decimal.scala

@@ -215,7 +215,8 @@ final class Decimal extends Ordered[Decimal] with Serializable {
     if (decimalVal.eq(null)) {
       longVal / POW_10(_scale)
     } else {
-      decimalVal.longValue()
+      // This will throw an exception if overflow occurs
+      decimalVal.toLongExact
     }
   }
 

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CastSuite.scala

@@ -922,4 +922,10 @@ class CastSuite extends SparkFunSuite with ExpressionEvalHelper {
     val ret6 = cast(Literal.create((1, Map(1 -> "a", 2 -> "b", 3 -> "c"))), StringType)
     checkEvaluation(ret6, "[1, [1 -> a, 2 -> b, 3 -> c]]")
   }
+
+  test("SPARK-24598: Cast to long should fail on overflow") {
+    checkExceptionInExpression[ArithmeticException](
+      cast(Literal.create(Decimal(Long.MaxValue) + Decimal(1)), LongType), "Overflow")
+    checkEvaluation(cast(Literal.create(Decimal(Long.MaxValue)), LongType), Long.MaxValue)
+  }
 }

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

@@ -114,11 +114,7 @@ class DataFrameRangeSuite extends QueryTest with SharedSQLContext with Eventuall
     val random = new Random(seed)
 
     def randomBound(): Long = {
-      val n = if (random.nextBoolean()) {
-        random.nextLong() % (Long.MaxValue / (100 * MAX_NUM_STEPS))
-      } else {
-        random.nextLong() / 2
-      }
+      val n = random.nextLong() % (Long.MaxValue / (100 * MAX_NUM_STEPS))
       if (random.nextBoolean()) n else -n
     }
 

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

@@ -17,6 +17,7 @@
 
 package org.apache.spark.sql
 
+import org.apache.spark.SparkException
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.expressions.Aggregator
 import org.apache.spark.sql.expressions.scalalang.typed
@@ -333,4 +334,12 @@ class DatasetAggregatorSuite extends QueryTest with SharedSQLContext {
       df.groupBy($"i").agg(VeryComplexResultAgg.toColumn),
       Row(1, Row(Row(1, "a"), Row(1, "a"))) :: Row(2, Row(Row(2, "bc"), Row(2, "bc"))) :: Nil)
   }
+
+  test("SPARK-24598: sum throws exception instead of silently overflow") {
+    val df1 = Seq(Long.MinValue, -10, Long.MaxValue).toDF("i")
+    checkAnswer(df1.agg(sum($"i")), Row(-11))
+    val df2 = Seq(Long.MinValue, -10, 8).toDF("i")
+    val e = intercept[SparkException](df2.agg(sum($"i")).collect())
+    assert(e.getCause.isInstanceOf[ArithmeticException])
+  }
 }