diff --git a/spark/pom.xml b/spark/pom.xml index fbe059c739..4b5d539e16 100644 --- a/spark/pom.xml +++ b/spark/pom.xml @@ -536,6 +536,7 @@ under the License. src/test/${shims.majorVerSrc} + src/test/${shims.minorPlusVerSrc} src/test/${shims.minorVerSrc} diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index 5df16eab67..4ae8144648 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -785,7 +785,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { inputs: Seq[Attribute], binding: Boolean = true): Option[Expr] = { - val newExpr = DecimalPrecision.promote(expr, !SQLConf.get.ansiEnabled) + val newExpr = DecimalPrecision.promote(expr) exprToProtoInternal(newExpr, inputs, binding) } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/DecimalPrecision.scala b/spark/src/main/scala/org/apache/spark/sql/comet/DecimalPrecision.scala index 6309a4304a..2d5247df7e 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/DecimalPrecision.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/DecimalPrecision.scala @@ -33,9 +33,13 @@ import org.apache.spark.sql.types.DecimalType * (SPARK-53968) it preserves the per-expression `allowDecimalPrecisionLoss` captured at view * creation time. Recomputing from the live `SQLConf` would re-label a stored DEC(38, 17) result * as DEC(38, 18) (or vice versa) and shift values by 10x (issue #4124). + * + * The `CheckOverflow` behavior also follows each expression's captured `evalMode`, since TRY + * arithmetic can be nested inside ANSI arithmetic (or vice versa) independently of the live + * session ANSI setting. */ object DecimalPrecision { - def promote(expr: Expression, nullOnOverflow: Boolean): Expression = { + def promote(expr: Expression): Expression = { expr.transformUp { // This means the binary expression is already optimized with the rule in Spark. This can // happen if the Spark version is < 3.4 @@ -43,23 +47,23 @@ object DecimalPrecision { case add @ Add(DecimalExpression(_, _), DecimalExpression(_, _), _) if add.dataType.isInstanceOf[DecimalType] => - CheckOverflow(add, add.dataType.asInstanceOf[DecimalType], nullOnOverflow) + CheckOverflow(add, add.dataType.asInstanceOf[DecimalType], add.evalMode != EvalMode.ANSI) case sub @ Subtract(DecimalExpression(_, _), DecimalExpression(_, _), _) if sub.dataType.isInstanceOf[DecimalType] => - CheckOverflow(sub, sub.dataType.asInstanceOf[DecimalType], nullOnOverflow) + CheckOverflow(sub, sub.dataType.asInstanceOf[DecimalType], sub.evalMode != EvalMode.ANSI) case mul @ Multiply(DecimalExpression(_, _), DecimalExpression(_, _), _) if mul.dataType.isInstanceOf[DecimalType] => - CheckOverflow(mul, mul.dataType.asInstanceOf[DecimalType], nullOnOverflow) + CheckOverflow(mul, mul.dataType.asInstanceOf[DecimalType], mul.evalMode != EvalMode.ANSI) case div @ Divide(DecimalExpression(_, _), DecimalExpression(_, _), _) if div.dataType.isInstanceOf[DecimalType] => - CheckOverflow(div, div.dataType.asInstanceOf[DecimalType], nullOnOverflow) + CheckOverflow(div, div.dataType.asInstanceOf[DecimalType], div.evalMode != EvalMode.ANSI) case rem @ Remainder(DecimalExpression(_, _), DecimalExpression(_, _), _) if rem.dataType.isInstanceOf[DecimalType] => - CheckOverflow(rem, rem.dataType.asInstanceOf[DecimalType], nullOnOverflow) + CheckOverflow(rem, rem.dataType.asInstanceOf[DecimalType], rem.evalMode != EvalMode.ANSI) case e => e } diff --git a/spark/src/test/spark-4.1/org/apache/spark/sql/comet/CometDecimalArithmeticViewSuite.scala b/spark/src/test/spark-4.1+/org/apache/spark/sql/comet/CometDecimalArithmeticViewSuite.scala similarity index 54% rename from spark/src/test/spark-4.1/org/apache/spark/sql/comet/CometDecimalArithmeticViewSuite.scala rename to spark/src/test/spark-4.1+/org/apache/spark/sql/comet/CometDecimalArithmeticViewSuite.scala index 0a16f8e55c..37189714ad 100644 --- a/spark/src/test/spark-4.1/org/apache/spark/sql/comet/CometDecimalArithmeticViewSuite.scala +++ b/spark/src/test/spark-4.1+/org/apache/spark/sql/comet/CometDecimalArithmeticViewSuite.scala @@ -20,10 +20,12 @@ package org.apache.spark.sql.comet import org.apache.spark.sql.CometTestBase -import org.apache.spark.sql.catalyst.expressions.{Add, AttributeReference, CheckOverflow, EvalMode, NumericEvalContext} +import org.apache.spark.sql.catalyst.expressions.{Add, AttributeReference, BinaryArithmetic, CheckOverflow, Divide, EvalMode, Expression, Multiply, NumericEvalContext, Remainder, Subtract} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.DecimalType +import org.apache.comet.serde.{ExprOuterClass, QueryPlanSerde} + class CometDecimalArithmeticViewSuite extends CometTestBase { // Spark 4.1+ (SPARK-53968) stores `spark.sql.decimalOperations.allowPrecisionLoss` per @@ -50,7 +52,7 @@ class CometDecimalArithmeticViewSuite extends CometTestBase { Seq((true, storedFalse), (false, storedTrue)).foreach { case (currentConf, add) => withSQLConf(SQLConf.DECIMAL_OPERATIONS_ALLOW_PREC_LOSS.key -> currentConf.toString) { val promoted = org.apache.spark.sql.comet.DecimalPrecision - .promote(add, nullOnOverflow = true) + .promote(add) promoted match { case CheckOverflow(_, dt, _) => assert( @@ -63,4 +65,58 @@ class CometDecimalArithmeticViewSuite extends CometTestBase { } } } + + test("issue #5075: DecimalPrecision.promote honours per-expression eval mode") { + val left = AttributeReference("a", DecimalType(10, 0))() + val right = AttributeReference("b", DecimalType(10, 0))() + val third = AttributeReference("c", DecimalType(10, 0))() + + val operations: Seq[( + String, + (Expression, Expression, NumericEvalContext) => BinaryArithmetic, + ExprOuterClass.Expr => ExprOuterClass.MathExpr)] = Seq( + ("add", Add.apply, _.getAdd), + ("subtract", Subtract.apply, _.getSubtract), + ("multiply", Multiply.apply, _.getMultiply), + ("divide", Divide.apply, _.getDivide), + ("remainder", Remainder.apply, _.getRemainder)) + val tryContext = NumericEvalContext(EvalMode.TRY, allowDecimalPrecisionLoss = true) + val ansiContext = NumericEvalContext(EvalMode.ANSI, allowDecimalPrecisionLoss = true) + val expressionTrees = operations.map { case (name, operation, getMathExpr) => + (name, operation(operation(left, right, tryContext), third, ansiContext), getMathExpr) + } + + Seq(false, true).foreach { sessionAnsiEnabled => + withSQLConf(SQLConf.ANSI_ENABLED.key -> sessionAnsiEnabled.toString) { + expressionTrees.foreach { case (name, expression, getMathExpr) => + val proto = QueryPlanSerde.exprToProto(expression, Seq(left, right, third)).get + assert(proto.hasCheckOverflow, s"$name under session ANSI=$sessionAnsiEnabled") + val ansiOverflow = proto.getCheckOverflow + assert(ansiOverflow.getFailOnError, s"$name under session ANSI=$sessionAnsiEnabled") + + // Decimal division already adds its own CheckOverflow inside the one added by + // DecimalPrecision, so peel that wrapper before inspecting the Divide proto. + val mathExprProto = + if (name == "divide") { + assert( + ansiOverflow.getChild.hasCheckOverflow, + s"$name under session ANSI=$sessionAnsiEnabled") + val divideOverflow = ansiOverflow.getChild.getCheckOverflow + assert( + divideOverflow.getFailOnError, + s"$name under session ANSI=$sessionAnsiEnabled") + divideOverflow.getChild + } else { + ansiOverflow.getChild + } + + val tryExprProto = getMathExpr(mathExprProto).getLeft + assert(tryExprProto.hasCheckOverflow, s"$name under session ANSI=$sessionAnsiEnabled") + assert( + !tryExprProto.getCheckOverflow.getFailOnError, + s"$name under session ANSI=$sessionAnsiEnabled") + } + } + } + } }