From 77397316baaf72f085bd70ccf9300e08fb1ccc48 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Fri, 31 Jul 2026 23:04:51 +0800 Subject: [PATCH 1/2] fix: use per-expression eval mode for decimal promotion --- .../apache/comet/serde/QueryPlanSerde.scala | 2 +- .../spark/sql/comet/DecimalPrecision.scala | 12 ++++----- .../CometDecimalArithmeticViewSuite.scala | 25 ++++++++++++++++++- 3 files changed, 31 insertions(+), 8 deletions(-) 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 6f8306862a..59156f4a3c 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -784,7 +784,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..cd2ce24155 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 @@ -35,7 +35,7 @@ import org.apache.spark.sql.types.DecimalType * as DEC(38, 18) (or vice versa) and shift values by 10x (issue #4124). */ 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 +43,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 index 0a16f8e55c..977778a653 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 @@ -24,6 +24,8 @@ import org.apache.spark.sql.catalyst.expressions.{Add, AttributeReference, Check import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.DecimalType +import org.apache.comet.serde.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,25 @@ 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 tryAdd = + Add(left, right, NumericEvalContext(EvalMode.TRY, allowDecimalPrecisionLoss = true)) + val ansiAdd = + Add(tryAdd, third, NumericEvalContext(EvalMode.ANSI, allowDecimalPrecisionLoss = true)) + + withSQLConf(SQLConf.ANSI_ENABLED.key -> "false") { + val proto = QueryPlanSerde.exprToProto(ansiAdd, Seq(left, right, third)).get + assert(proto.hasCheckOverflow) + val ansiOverflow = proto.getCheckOverflow + assert(ansiOverflow.getFailOnError) + + val tryAddProto = ansiOverflow.getChild.getAdd.getLeft + assert(tryAddProto.hasCheckOverflow) + assert(!tryAddProto.getCheckOverflow.getFailOnError) + } + } } From f6c433121b118fb3c904839e5e33ce4db9bb90ef Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sat, 1 Aug 2026 10:35:51 +0800 Subject: [PATCH 2/2] andy's review --- spark/pom.xml | 1 + .../spark/sql/comet/DecimalPrecision.scala | 4 ++ .../CometDecimalArithmeticViewSuite.scala | 61 ++++++++++++++----- 3 files changed, 52 insertions(+), 14 deletions(-) rename spark/src/test/{spark-4.1 => spark-4.1+}/org/apache/spark/sql/comet/CometDecimalArithmeticViewSuite.scala (59%) 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/spark/sql/comet/DecimalPrecision.scala b/spark/src/main/scala/org/apache/spark/sql/comet/DecimalPrecision.scala index cd2ce24155..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,6 +33,10 @@ 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): Expression = { 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 59% 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 977778a653..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,11 +20,11 @@ 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.QueryPlanSerde +import org.apache.comet.serde.{ExprOuterClass, QueryPlanSerde} class CometDecimalArithmeticViewSuite extends CometTestBase { @@ -70,20 +70,53 @@ class CometDecimalArithmeticViewSuite extends CometTestBase { val left = AttributeReference("a", DecimalType(10, 0))() val right = AttributeReference("b", DecimalType(10, 0))() val third = AttributeReference("c", DecimalType(10, 0))() - val tryAdd = - Add(left, right, NumericEvalContext(EvalMode.TRY, allowDecimalPrecisionLoss = true)) - val ansiAdd = - Add(tryAdd, third, NumericEvalContext(EvalMode.ANSI, allowDecimalPrecisionLoss = true)) - withSQLConf(SQLConf.ANSI_ENABLED.key -> "false") { - val proto = QueryPlanSerde.exprToProto(ansiAdd, Seq(left, right, third)).get - assert(proto.hasCheckOverflow) - val ansiOverflow = proto.getCheckOverflow - assert(ansiOverflow.getFailOnError) + 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 tryAddProto = ansiOverflow.getChild.getAdd.getLeft - assert(tryAddProto.hasCheckOverflow) - assert(!tryAddProto.getCheckOverflow.getFailOnError) + val tryExprProto = getMathExpr(mathExprProto).getLeft + assert(tryExprProto.hasCheckOverflow, s"$name under session ANSI=$sessionAnsiEnabled") + assert( + !tryExprProto.getCheckOverflow.getFailOnError, + s"$name under session ANSI=$sessionAnsiEnabled") + } + } } } }