[SPARK-28610][SQL] Allow having a decimal buffer for long sum

Author: mgaido91

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

@@ -21,6 +21,7 @@ import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.util.TypeUtils
+import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types._
 
 @ExpressionDescription(
@@ -56,7 +57,11 @@ 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 if SQLConf.get.getConf(SQLConf.SUM_DECIMAL_BUFFER_FOR_LONG) =>
+      DecimalType.BigIntDecimal
+    case _ => resultType
+  }
 
   private lazy val sum = AttributeReference("sum", sumDataType)()
 
@@ -89,5 +94,11 @@ case class Sum(child: Expression) extends DeclarativeAggregate with ImplicitCast
     )
   }
 
-  override lazy val evaluateExpression: Expression = sum
+  override lazy val evaluateExpression: Expression = {
+    if (sumDataType == resultType) {
+      sum
+    } else {
+      Cast(sum, resultType)
+    }
+  }
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -1789,6 +1789,16 @@ object SQLConf {
       .booleanConf
       .createWithDefault(false)
 
+  val SUM_DECIMAL_BUFFER_FOR_LONG =
+    buildConf("spark.sql.sum.decimalBufferForLong")
+      .doc("If it is set to true, sum of long uses decimal type for the buffer. When false " +
+        "(default), long is used to the buffer. If spark.sql.arithmeticOperations.failOnOverFlow" +
+        " is turned on, having this config set to true allows operations which have temporary " +
+        "overflows to execute properly without the exception thrown when this flag is false.")
+      .internal()
+      .booleanConf
+      .createWithDefault(false)
+
   val LEGACY_HAVING_WITHOUT_GROUP_BY_AS_WHERE =
     buildConf("spark.sql.legacy.parser.havingWithoutGroupByAsWhere")
       .internal()

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

@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql
 
+import org.apache.spark.SparkException
+
 import scala.util.Random
 
 import org.scalatest.Matchers.the
@@ -927,4 +929,20 @@ class DataFrameAggregateSuite extends QueryTest with SharedSQLContext {
       assert(error.message.contains("function count_if requires boolean type"))
     }
   }
+
+  test("SPARK-28610: temporary overflow on sum of long should not fail") {
+    withSQLConf(SQLConf.ARITHMETIC_OPERATIONS_FAIL_ON_OVERFLOW.key -> "true",
+        SQLConf.SUM_DECIMAL_BUFFER_FOR_LONG.key -> "true") {
+      val df = sparkContext.parallelize(Seq(100L, Long.MaxValue, -1000L), 1).toDF("a")
+      checkAnswer(df.select(sum($"a")), Row(Long.MaxValue - 900L))
+    }
+    withSQLConf(SQLConf.ARITHMETIC_OPERATIONS_FAIL_ON_OVERFLOW.key -> "true",
+      SQLConf.SUM_DECIMAL_BUFFER_FOR_LONG.key -> "false") {
+      val df = sparkContext.parallelize(Seq(100L, Long.MaxValue, -1000L), 1).toDF("a")
+      val e = intercept[SparkException] {
+        df.select(sum($"a")).collect()
+      }
+      assert(e.getCause.isInstanceOf[ArithmeticException])
+    }
+  }
 }