[SPARK-26218][SQL] Overflow on arithmetic operations returns incorrect result

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

@@ -56,7 +56,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)()
 
@@ -89,5 +95,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/catalyst/expressions/arithmetic.scala

@@ -117,6 +117,8 @@ case class Abs(child: Expression)
 
 abstract class BinaryArithmetic extends BinaryOperator with NullIntolerant {
 
+  protected val checkOverflow = SQLConf.get.getConf(SQLConf.ARITHMETIC_OPERATION_OVERFLOW_CHECK)
+
   override def dataType: DataType = left.dataType
 
   override lazy val resolved: Boolean = childrenResolved && checkInputDataTypes().isSuccess
@@ -129,17 +131,41 @@ abstract class BinaryArithmetic extends BinaryOperator with NullIntolerant {
   def calendarIntervalMethod: String =
     sys.error("BinaryArithmetics must override either calendarIntervalMethod or genCode")
 
+  def checkOverflowCode(result: String, op1: String, op2: String): String =
+    sys.error("BinaryArithmetics must override either checkOverflowCode or genCode")
+
   override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = dataType match {
     case _: DecimalType =>
       defineCodeGen(ctx, ev, (eval1, eval2) => s"$eval1.$decimalMethod($eval2)")
     case CalendarIntervalType =>
       defineCodeGen(ctx, ev, (eval1, eval2) => s"$eval1.$calendarIntervalMethod($eval2)")
+    // In the following cases, overflow can happen, so we need to check the result is valid.
+    // Otherwise we throw an ArithmeticException
     // byte and short are casted into int when add, minus, times or divide
     case ByteType | ShortType =>
-      defineCodeGen(ctx, ev,
-        (eval1, eval2) => s"(${CodeGenerator.javaType(dataType)})($eval1 $symbol $eval2)")
+      nullSafeCodeGen(ctx, ev, (eval1, eval2) => {
+        val overflowCheck = if (checkOverflow) {
+          checkOverflowCode(ev.value, eval1, eval2)
+        } else {
+          ""
+        }
+        s"""
+           |${ev.value} = (${CodeGenerator.javaType(dataType)})($eval1 $symbol $eval2);
+           |$overflowCheck
+         """.stripMargin
+      })
     case _ =>
-      defineCodeGen(ctx, ev, (eval1, eval2) => s"$eval1 $symbol $eval2")
+      nullSafeCodeGen(ctx, ev, (eval1, eval2) => {
+        val overflowCheck = if (checkOverflow) {
+          checkOverflowCode(ev.value, eval1, eval2)
+        } else {
+          ""
+        }
+        s"""
+           |${ev.value} = $eval1 $symbol $eval2;
+           |$overflowCheck
+         """.stripMargin
+      })
   }
 }
 
@@ -170,9 +196,27 @@ case class Add(left: Expression, right: Expression) extends BinaryArithmetic {
     if (dataType.isInstanceOf[CalendarIntervalType]) {
       input1.asInstanceOf[CalendarInterval].add(input2.asInstanceOf[CalendarInterval])
     } else {
-      numeric.plus(input1, input2)
+      val result = numeric.plus(input1, input2)
+      if (checkOverflow) {
+        val resSignum = numeric.signum(result)
+        val input1Signum = numeric.signum(input1)
+        val input2Signum = numeric.signum(input2)
+        if (resSignum != -1 && input1Signum == -1 && input2Signum == -1
+          || resSignum != 1 && input1Signum == 1 && input2Signum == 1) {
+          throw new ArithmeticException(s"$input1 + $input2 caused overflow.")
+        }
+      }
+      result
     }
   }
+
+  override def checkOverflowCode(result: String, op1: String, op2: String): String = {
+    s"""
+       |if ($result >= 0 && $op1 < 0 && $op2 < 0 || $result <= 0 && $op1 > 0 && $op2 > 0) {
+       |  throw new ArithmeticException($op1 + " + " + $op2 + " caused overflow.");
+       |}
+     """.stripMargin
+  }
 }
 
 @ExpressionDescription(
@@ -198,9 +242,27 @@ case class Subtract(left: Expression, right: Expression) extends BinaryArithmeti
     if (dataType.isInstanceOf[CalendarIntervalType]) {
       input1.asInstanceOf[CalendarInterval].subtract(input2.asInstanceOf[CalendarInterval])
     } else {
-      numeric.minus(input1, input2)
+      val result = numeric.minus(input1, input2)
+      if (checkOverflow) {
+        val resSignum = numeric.signum(result)
+        val input1Signum = numeric.signum(input1)
+        val input2Signum = numeric.signum(input2)
+        if (resSignum != 1 && input1Signum == 1 && input2Signum == -1
+          || resSignum != -1 && input1Signum == -1 && input2Signum == 1) {
+          throw new ArithmeticException(s"$input1 - $input2 caused overflow.")
+        }
+      }
+      result
     }
   }
+
+  override def checkOverflowCode(result: String, op1: String, op2: String): String = {
+    s"""
+       |if ($result <= 0 && $op1 > 0 && $op2 < 0 || $result >= 0 && $op1 < 0 && $op2 > 0) {
+       |  throw new ArithmeticException($op1 + " - " + $op2 + " caused overflow.");
+       |}
+     """.stripMargin
+  }
 }
 
 @ExpressionDescription(
@@ -219,7 +281,29 @@ case class Multiply(left: Expression, right: Expression) extends BinaryArithmeti
 
   private lazy val numeric = TypeUtils.getNumeric(dataType)
 
-  protected override def nullSafeEval(input1: Any, input2: Any): Any = numeric.times(input1, input2)
+  protected override def nullSafeEval(input1: Any, input2: Any): Any = {
+    val result = numeric.times(input1, input2)
+    if (checkOverflow) {
+      if (numeric.signum(result) != numeric.signum(input1) * numeric.signum(input2) &&
+        !(result.isInstanceOf[Double] && !result.asInstanceOf[Double].isNaN) &&
+        !(result.isInstanceOf[Float] && !result.asInstanceOf[Float].isNaN)) {
+        throw new ArithmeticException(s"$input1 * $input2 caused overflow.")
+      }
+    }
+    result
+  }
+
+  override def checkOverflowCode(result: String, op1: String, op2: String): String = {
+    val isNaNCheck = dataType match {
+      case DoubleType | FloatType => s" && !java.lang.Double.isNaN($result)"
+      case _ => ""
+    }
+    s"""
+       |if (Math.signum($result) != Math.signum($op1) * Math.signum($op2)$isNaNCheck) {
+       |  throw new ArithmeticException($op1 + " * " + $op2 + " caused overflow.");
+       |}
+     """.stripMargin
+  }
 }
 
 // Common base trait for Divide and Remainder, since these two classes are almost identical

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/bitwiseExpressions.scala

@@ -51,6 +51,8 @@ case class BitwiseAnd(left: Expression, right: Expression) extends BinaryArithme
   }
 
   protected override def nullSafeEval(input1: Any, input2: Any): Any = and(input1, input2)
+
+  override def checkOverflowCode(result: String, op1: String, op2: String): String = ""
 }
 
 /**
@@ -83,6 +85,8 @@ case class BitwiseOr(left: Expression, right: Expression) extends BinaryArithmet
   }
 
   protected override def nullSafeEval(input1: Any, input2: Any): Any = or(input1, input2)
+
+  override def checkOverflowCode(result: String, op1: String, op2: String): String = ""
 }
 
 /**
@@ -115,6 +119,8 @@ case class BitwiseXor(left: Expression, right: Expression) extends BinaryArithme
   }
 
   protected override def nullSafeEval(input1: Any, input2: Any): Any = xor(input1, input2)
+
+  override def checkOverflowCode(result: String, op1: String, op2: String): String = ""
 }
 
 /**

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

@@ -1764,6 +1764,14 @@ object SQLConf {
     .booleanConf
     .createWithDefault(false)
 
+  val ARITHMETIC_OPERATION_OVERFLOW_CHECK = buildConf("spark.sql.arithmetic.checkOverflow")
+    .doc("If it is set to true, all arithmetic operations on non-decimal fields throw an " +
+      "exception if an overflow occurs. If it is false (default), in case of overflow a wrong " +
+      "result is returned.")
+    .internal()
+    .booleanConf
+    .createWithDefault(false)
+
   val LEGACY_HAVING_WITHOUT_GROUP_BY_AS_WHERE =
     buildConf("spark.sql.legacy.parser.havingWithoutGroupByAsWhere")
       .internal()

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

@@ -22,6 +22,7 @@ import java.math.{BigInteger, MathContext, RoundingMode}
 
 import org.apache.spark.annotation.Unstable
 import org.apache.spark.sql.AnalysisException
+import org.apache.spark.sql.internal.SQLConf
 
 /**
  * A mutable implementation of BigDecimal that can hold a Long if values are small enough.
@@ -228,6 +229,12 @@ final class Decimal extends Ordered[Decimal] with Serializable {
     if (decimalVal.eq(null)) {
       longVal / POW_10(_scale)
     } else {
+      if (SQLConf.get.getConf(SQLConf.ARITHMETIC_OPERATION_OVERFLOW_CHECK)) {
+        // This will throw an exception if overflow occurs
+        if (decimalVal.compare(LONG_MIN_BIG_DEC) < 0 || decimalVal.compare(LONG_MAX_BIG_DEC) > 0) {
+          throw new ArithmeticException("Overflow")
+        }
+      }
       decimalVal.longValue()
     }
   }
@@ -456,6 +463,9 @@ object Decimal {
   private val LONG_MAX_BIG_INT = BigInteger.valueOf(JLong.MAX_VALUE)
   private val LONG_MIN_BIG_INT = BigInteger.valueOf(JLong.MIN_VALUE)
 
+  private val LONG_MAX_BIG_DEC = BigDecimal.valueOf(JLong.MAX_VALUE)
+  private val LONG_MIN_BIG_DEC = BigDecimal.valueOf(JLong.MIN_VALUE)
+
   def apply(value: Double): Decimal = new Decimal().set(value)
 
   def apply(value: Long): Decimal = new Decimal().set(value)

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

@@ -59,8 +59,12 @@ class ArithmeticExpressionSuite extends SparkFunSuite with ExpressionEvalHelper
     checkEvaluation(Add(positiveIntLit, negativeIntLit), -1)
     checkEvaluation(Add(positiveLongLit, negativeLongLit), -1L)
 
-    DataTypeTestUtils.numericAndInterval.foreach { tpe =>
-      checkConsistencyBetweenInterpretedAndCodegen(Add, tpe, tpe)
+    Seq("true", "false").foreach { checkOverflow =>
+      withSQLConf(SQLConf.ARITHMETIC_OPERATION_OVERFLOW_CHECK.key -> checkOverflow) {
+        DataTypeTestUtils.numericAndInterval.foreach { tpe =>
+          checkConsistencyBetweenInterpretedAndCodegenAllowingException(Add, tpe, tpe)
+        }
+      }
     }
   }
 
@@ -100,8 +104,12 @@ class ArithmeticExpressionSuite extends SparkFunSuite with ExpressionEvalHelper
     checkEvaluation(Subtract(positiveIntLit, negativeIntLit), positiveInt - negativeInt)
     checkEvaluation(Subtract(positiveLongLit, negativeLongLit), positiveLong - negativeLong)
 
-    DataTypeTestUtils.numericAndInterval.foreach { tpe =>
-      checkConsistencyBetweenInterpretedAndCodegen(Subtract, tpe, tpe)
+    Seq("true", "false").foreach { checkOverflow =>
+      withSQLConf(SQLConf.ARITHMETIC_OPERATION_OVERFLOW_CHECK.key -> checkOverflow) {
+        DataTypeTestUtils.numericAndInterval.foreach { tpe =>
+          checkConsistencyBetweenInterpretedAndCodegenAllowingException(Subtract, tpe, tpe)
+        }
+      }
     }
   }
 
@@ -118,8 +126,12 @@ class ArithmeticExpressionSuite extends SparkFunSuite with ExpressionEvalHelper
     checkEvaluation(Multiply(positiveIntLit, negativeIntLit), positiveInt * negativeInt)
     checkEvaluation(Multiply(positiveLongLit, negativeLongLit), positiveLong * negativeLong)
 
-    DataTypeTestUtils.numericTypeWithoutDecimal.foreach { tpe =>
-      checkConsistencyBetweenInterpretedAndCodegen(Multiply, tpe, tpe)
+    Seq("true", "false").foreach { checkOverflow =>
+      withSQLConf(SQLConf.ARITHMETIC_OPERATION_OVERFLOW_CHECK.key -> checkOverflow) {
+        DataTypeTestUtils.numericTypeWithoutDecimal.foreach { tpe =>
+          checkConsistencyBetweenInterpretedAndCodegenAllowingException(Multiply, tpe, tpe)
+        }
+      }
     }
   }
 
@@ -376,4 +388,28 @@ class ArithmeticExpressionSuite extends SparkFunSuite with ExpressionEvalHelper
     Greatest(Seq(Literal(1), Literal(1))).genCode(ctx2)
     assert(ctx2.inlinedMutableStates.size == 1)
   }
+
+  test("SPARK-24598: overflow on BigInt returns wrong result") {
+    val maxLongLiteral = Literal(Long.MaxValue)
+    val minLongLiteral = Literal(Long.MinValue)
+    val e1 = Add(maxLongLiteral, Literal(1L))
+    val e2 = Subtract(maxLongLiteral, Literal(-1L))
+    val e3 = Multiply(maxLongLiteral, Literal(2L))
+    val e4 = Add(minLongLiteral, minLongLiteral)
+    val e5 = Subtract(minLongLiteral, maxLongLiteral)
+    val e6 = Multiply(minLongLiteral, minLongLiteral)
+    withSQLConf(SQLConf.ARITHMETIC_OPERATION_OVERFLOW_CHECK.key -> "true") {
+      Seq(e1, e2, e3, e4, e5, e6).foreach { e =>
+        checkExceptionInExpression[ArithmeticException](e, "caused overflow")
+      }
+    }
+    withSQLConf(SQLConf.ARITHMETIC_OPERATION_OVERFLOW_CHECK.key -> "false") {
+      checkEvaluation(e1, Long.MinValue)
+      checkEvaluation(e2, Long.MinValue)
+      checkEvaluation(e3, -2L)
+      checkEvaluation(e4, 0L)
+      checkEvaluation(e5, 1L)
+      checkEvaluation(e6, 0L)
+    }
+  }
 }

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

@@ -29,6 +29,7 @@ import org.apache.spark.sql.catalyst.expressions.codegen.CodegenContext
 import org.apache.spark.sql.catalyst.util.DateTimeTestUtils._
 import org.apache.spark.sql.catalyst.util.DateTimeUtils
 import org.apache.spark.sql.catalyst.util.DateTimeUtils._
+import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types._
 import org.apache.spark.unsafe.types.UTF8String
 
@@ -956,6 +957,19 @@ class CastSuite extends SparkFunSuite with ExpressionEvalHelper {
     checkEvaluation(ret6, "[1, [1 -> a, 2 -> b, 3 -> c]]")
   }
 
+  test("SPARK-24598: Cast to long should fail on overflow") {
+    val overflowCast = cast(Literal.create(Decimal(Long.MaxValue) + Decimal(1)), LongType)
+    val nonOverflowCast = cast(Literal.create(Decimal(Long.MaxValue)), LongType)
+    withSQLConf(SQLConf.ARITHMETIC_OPERATION_OVERFLOW_CHECK.key -> "true") {
+      checkExceptionInExpression[ArithmeticException](overflowCast, "Overflow")
+      checkEvaluation(nonOverflowCast, Long.MaxValue)
+    }
+    withSQLConf(SQLConf.ARITHMETIC_OPERATION_OVERFLOW_CHECK.key -> "false") {
+      checkEvaluation(overflowCast, Long.MinValue)
+      checkEvaluation(nonOverflowCast, Long.MaxValue)
+    }
+  }
+
   test("up-cast") {
     def isCastSafe(from: NumericType, to: NumericType): Boolean = (from, to) match {
       case (_, dt: DecimalType) => dt.isWiderThan(from)