From afbebbf2311687399a15603307629f75a3954b1f Mon Sep 17 00:00:00 2001 From: peterxcli Date: Mon, 3 Aug 2026 06:45:19 +0800 Subject: [PATCH 1/5] fix: avoid duplicate CheckOverflow evaluation for decimal division --- .../org/apache/comet/serde/arithmetic.scala | 17 +------ .../spark/sql/comet/DecimalPrecision.scala | 11 ++++ .../CometDecimalArithmeticViewSuite.scala | 51 +++++++++++++------ 3 files changed, 47 insertions(+), 32 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala b/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala index 614fa4a5c0..7b1d878bd3 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala @@ -263,7 +263,7 @@ object CometDivide extends CometExpressionSerde[Divide] with MathBase { // For now, use NullIf to swap zeros with nulls. val rightExpr = if (expr.evalMode != EvalMode.ANSI) nullIfWhenPrimitive(expr.right) else expr.right - val divideExpr = createMathExpression( + createMathExpression( expr, expr.left, rightExpr, @@ -272,21 +272,6 @@ object CometDivide extends CometExpressionSerde[Divide] with MathBase { expr.dataType, expr.evalMode, (builder, mathExpr) => builder.setDivide(mathExpr)) - - // For decimal division Spark applies CheckOverflow after dividing: in ANSI mode overflow - // throws NUMERIC_VALUE_OUT_OF_RANGE; in legacy/try mode it returns null. The Rust - // spark_decimal_div_internal uses i128::MAX as a sentinel for overflow, so without this - // wrapper an ANSI overflow would silently return a garbage value instead of throwing. - if (divideExpr.isDefined && expr.dataType.isInstanceOf[DecimalType] && - serializeDataType(expr.dataType).isDefined) { - val builder = ExprOuterClass.CheckOverflow.newBuilder() - builder.setChild(divideExpr.get) - builder.setFailOnError(expr.evalMode == EvalMode.ANSI) - builder.setDatatype(serializeDataType(expr.dataType).get) - Some(ExprOuterClass.Expr.newBuilder().setCheckOverflow(builder).build()) - } else { - divideExpr - } } } 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 2d5247df7e..7298a1eae3 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 @@ -45,6 +45,15 @@ object DecimalPrecision { // happen if the Spark version is < 3.4 case e: BinaryArithmetic if e.left.prettyName == "promote_precision" => e + // Recursive exprToProto calls can re-promote every decimal binary operator below. Collapse + // only equivalent wrappers so the shared promotion rule remains idempotent. + case outer @ CheckOverflow( + inner @ CheckOverflow(child: BinaryArithmetic, _, _), + dataType, + nullOnOverflow) + if inner.dataType == dataType && inner.nullOnOverflow == nullOnOverflow => + outer.copy(child = child) + case add @ Add(DecimalExpression(_, _), DecimalExpression(_, _), _) if add.dataType.isInstanceOf[DecimalType] => CheckOverflow(add, add.dataType.asInstanceOf[DecimalType], add.evalMode != EvalMode.ANSI) @@ -57,6 +66,8 @@ object DecimalPrecision { if mul.dataType.isInstanceOf[DecimalType] => CheckOverflow(mul, mul.dataType.asInstanceOf[DecimalType], mul.evalMode != EvalMode.ANSI) + // Native decimal division uses i128::MAX as an overflow sentinel, so this wrapper must + // convert it to null or an ANSI error according to the expression's eval mode. case div @ Divide(DecimalExpression(_, _), DecimalExpression(_, _), _) if div.dataType.isInstanceOf[DecimalType] => CheckOverflow(div, div.dataType.asInstanceOf[DecimalType], div.evalMode != EvalMode.ANSI) 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 37189714ad..1b01fa27bb 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 @@ -94,22 +94,7 @@ class CometDecimalArithmeticViewSuite extends CometTestBase { 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 mathExprProto = ansiOverflow.getChild val tryExprProto = getMathExpr(mathExprProto).getLeft assert(tryExprProto.hasCheckOverflow, s"$name under session ANSI=$sessionAnsiEnabled") assert( @@ -119,4 +104,38 @@ class CometDecimalArithmeticViewSuite extends CometTestBase { } } } + + test("issue #5190: recursive serialization does not duplicate decimal CheckOverflow") { + val left = "CAST(id AS DECIMAL(10, 0))" + val right = "CAST(id + 1 AS DECIMAL(10, 0))" + val operations: Seq[(String, Boolean, String, ExprOuterClass.Expr => Boolean, Boolean)] = Seq( + ("add", false, s"$left + $right", _.hasAdd, false), + ("subtract", false, s"$left - $right", _.hasSubtract, false), + ("multiply", false, s"$left * $right", _.hasMultiply, false), + ("remainder", false, s"$left % $right", _.hasRemainder, false), + ("divide LEGACY", false, s"$left / $right", _.hasDivide, false), + ("divide TRY", true, s"try_divide($left, $right)", _.hasDivide, false), + ("divide ANSI", true, s"$left / $right", _.hasDivide, true)) + + operations.foreach { case (name, ansiEnabled, arithmetic, hasOperation, failOnError) => + withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) { + val plan = spark + .sql(s"SELECT array_contains(array($arithmetic), $arithmetic) FROM range(1, 4)") + .queryExecution + .optimizedPlan + val arithmeticProto = QueryPlanSerde + .exprToProto(plan.expressions.head, plan.children.head.output) + .get + .getScalarFunc + .getArgs(1) + + assert(arithmeticProto.hasCheckOverflow, s"$name: $arithmeticProto") + val overflow = arithmeticProto.getCheckOverflow + assert(overflow.getFailOnError === failOnError, s"$name: $arithmeticProto") + assert( + hasOperation(overflow.getChild), + s"$name has duplicate CheckOverflow: $arithmeticProto") + } + } + } } From c61af2ed37fb10554bc3ac67e07f30619092eebb Mon Sep 17 00:00:00 2001 From: peterxcli Date: Tue, 4 Aug 2026 19:57:06 +0800 Subject: [PATCH 2/5] make test run on all spark version with better parameterize name --- .../org/apache/comet/serde/arithmetic.scala | 12 +- .../comet/CometDecimalPromotionSuite.scala | 104 ++++++++++++++++++ .../CometDecimalArithmeticViewSuite.scala | 34 ------ 3 files changed, 114 insertions(+), 36 deletions(-) create mode 100644 spark/src/test/scala/org/apache/spark/sql/comet/CometDecimalPromotionSuite.scala diff --git a/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala b/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala index 7b1d878bd3..55b177298d 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala @@ -251,8 +251,16 @@ object CometMultiply extends CometExpressionSerde[Multiply] with MathBase { object CometDivide extends CometExpressionSerde[Divide] with MathBase { - override def getSupportLevel(expr: Divide): SupportLevel = - mathDataTypeSupportLevel(expr.left.dataType) + override def getSupportLevel(expr: Divide): SupportLevel = { + if (expr.dataType.isInstanceOf[DecimalType] && + (!expr.left.dataType.isInstanceOf[DecimalType] || + !expr.right.dataType.isInstanceOf[DecimalType])) { + // This is just a sanity check. Spark should not allow this case to be created cause type coercion. + Unsupported(Some("Decimal division with a decimal result requires decimal operands")) + } else { + mathDataTypeSupportLevel(expr.left.dataType) + } + } override def convert( expr: Divide, diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometDecimalPromotionSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometDecimalPromotionSuite.scala new file mode 100644 index 0000000000..eaca6d4030 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometDecimalPromotionSuite.scala @@ -0,0 +1,104 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.comet + +import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.catalyst.expressions.{ArrayContains, AttributeReference, Divide, EvalMode} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{DecimalType, IntegerType} + +import org.apache.comet.serde.{CometDivide, ExprOuterClass, QueryPlanSerde, Unsupported} + +class CometDecimalPromotionSuite extends CometTestBase { + + private case class Operation( + name: String, + symbol: String, + hasProto: ExprOuterClass.Expr => Boolean) + + private case class TestMode(name: String, ansiEnabled: Boolean, failOnError: Boolean) + + test("issue #5190: decimal promotion is idempotent during recursive serialization") { + val left = "CAST(id AS DECIMAL(10, 0))" + val right = "CAST(id + 1 AS DECIMAL(10, 0))" + val divide = Operation(name = "divide", symbol = "/", hasProto = _.hasDivide) + val operations = Seq( + Operation(name = "add", symbol = "+", hasProto = _.hasAdd), + Operation(name = "subtract", symbol = "-", hasProto = _.hasSubtract), + Operation(name = "multiply", symbol = "*", hasProto = _.hasMultiply), + divide, + Operation(name = "remainder", symbol = "%", hasProto = _.hasRemainder)) + val legacy = TestMode(name = "LEGACY", ansiEnabled = false, failOnError = false) + val ansi = TestMode(name = "ANSI", ansiEnabled = true, failOnError = true) + + def check(operation: Operation, mode: TestMode, arithmetic: String): Unit = { + val name = s"${operation.name} ${mode.name}" + withSQLConf(SQLConf.ANSI_ENABLED.key -> mode.ansiEnabled.toString) { + val plan = spark + .sql(s"SELECT array_contains(array($arithmetic), $arithmetic) FROM range(1, 4)") + .queryExecution + .optimizedPlan + val expression = plan.expressions.head + val arrayContains = expression.collectFirst { case e: ArrayContains => e }.get + val promoted = DecimalPrecision.promote(expression) + assert( + DecimalPrecision.promote(promoted) == promoted, + s"$name promotion is not idempotent: $promoted") + + val arithmeticProto = QueryPlanSerde + .exprToProto(expression, plan.children.head.output) + .get + .getScalarFunc + .getArgs(1) + assert(arithmeticProto.hasCheckOverflow, s"$name: $arithmeticProto") + val overflow = arithmeticProto.getCheckOverflow + assert( + overflow.getDatatype === QueryPlanSerde + .serializeDataType(arrayContains.right.dataType) + .get, + s"$name has the wrong CheckOverflow datatype: $arithmeticProto") + assert(overflow.getFailOnError === mode.failOnError, s"$name: $arithmeticProto") + assert( + operation.hasProto(overflow.getChild), + s"$name has duplicate CheckOverflow: $arithmeticProto") + } + } + + operations.foreach { operation => + Seq(legacy, ansi).foreach { mode => + check(operation, mode, s"$left ${operation.symbol} $right") + } + } + check( + divide, + TestMode(name = "TRY", ansiEnabled = true, failOnError = false), + s"try_divide($left, $right)") + } + + test("decimal Divide with a non-decimal operand is unsupported") { + // This is just a sanity check. Spark should not allow this case to be created cause type coercion + val decimal = AttributeReference("decimal", DecimalType(10, 0))() + val integer = AttributeReference("integer", IntegerType)() + val divide = Divide(decimal, integer, EvalMode.LEGACY) + + assert(divide.dataType === DecimalType(10, 0)) + assert(CometDivide.getSupportLevel(divide).isInstanceOf[Unsupported]) + } +} 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 1b01fa27bb..0a214cf91f 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 @@ -104,38 +104,4 @@ class CometDecimalArithmeticViewSuite extends CometTestBase { } } } - - test("issue #5190: recursive serialization does not duplicate decimal CheckOverflow") { - val left = "CAST(id AS DECIMAL(10, 0))" - val right = "CAST(id + 1 AS DECIMAL(10, 0))" - val operations: Seq[(String, Boolean, String, ExprOuterClass.Expr => Boolean, Boolean)] = Seq( - ("add", false, s"$left + $right", _.hasAdd, false), - ("subtract", false, s"$left - $right", _.hasSubtract, false), - ("multiply", false, s"$left * $right", _.hasMultiply, false), - ("remainder", false, s"$left % $right", _.hasRemainder, false), - ("divide LEGACY", false, s"$left / $right", _.hasDivide, false), - ("divide TRY", true, s"try_divide($left, $right)", _.hasDivide, false), - ("divide ANSI", true, s"$left / $right", _.hasDivide, true)) - - operations.foreach { case (name, ansiEnabled, arithmetic, hasOperation, failOnError) => - withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) { - val plan = spark - .sql(s"SELECT array_contains(array($arithmetic), $arithmetic) FROM range(1, 4)") - .queryExecution - .optimizedPlan - val arithmeticProto = QueryPlanSerde - .exprToProto(plan.expressions.head, plan.children.head.output) - .get - .getScalarFunc - .getArgs(1) - - assert(arithmeticProto.hasCheckOverflow, s"$name: $arithmeticProto") - val overflow = arithmeticProto.getCheckOverflow - assert(overflow.getFailOnError === failOnError, s"$name: $arithmeticProto") - assert( - hasOperation(overflow.getChild), - s"$name has duplicate CheckOverflow: $arithmeticProto") - } - } - } } From 3f820f79f3f6dc76e6351e8d75d966c8e6ef0b44 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Tue, 4 Aug 2026 20:04:57 +0800 Subject: [PATCH 3/5] ci: register decimal promotion suite --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + 2 files changed, 2 insertions(+) diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 0a71607ead..3a877613ea 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -364,6 +364,7 @@ jobs: org.apache.spark.sql.comet.CometShuffleFallbackStickinessSuite org.apache.spark.sql.comet.PlanDataInjectorSuite org.apache.spark.sql.comet.CometDecimalArithmeticViewSuite + org.apache.spark.sql.comet.CometDecimalPromotionSuite org.apache.spark.sql.comet.CometScanWithPlanDataSuite org.apache.spark.sql.comet.util.UtilsSuite org.apache.comet.vector.NativeUtilSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 2335588b70..213ca1a4e5 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -180,6 +180,7 @@ jobs: org.apache.spark.sql.comet.CometShuffleFallbackStickinessSuite org.apache.spark.sql.comet.PlanDataInjectorSuite org.apache.spark.sql.comet.CometDecimalArithmeticViewSuite + org.apache.spark.sql.comet.CometDecimalPromotionSuite org.apache.spark.sql.comet.CometScanWithPlanDataSuite org.apache.spark.sql.comet.util.UtilsSuite org.apache.comet.vector.NativeUtilSuite From 1a38b7019387de93d9146bb3a950d2f50e52d348 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Wed, 5 Aug 2026 00:41:46 +0800 Subject: [PATCH 4/5] style: fix decimal division comment length --- spark/src/main/scala/org/apache/comet/serde/arithmetic.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala b/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala index 2aa81d8c8c..5535f1540a 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala @@ -255,7 +255,7 @@ object CometDivide extends CometExpressionSerde[Divide] with MathBase { if (expr.dataType.isInstanceOf[DecimalType] && (!expr.left.dataType.isInstanceOf[DecimalType] || !expr.right.dataType.isInstanceOf[DecimalType])) { - // This is just a sanity check. Spark should not allow this case to be created cause type coercion. + // This is only a sanity check; Spark's type coercion should prevent this case. Unsupported(Some("Decimal division with a decimal result requires decimal operands")) } else { mathDataTypeSupportLevel(expr.left.dataType) From e27a9f645a0015802ed62ad4d4f77c9654ac1bf8 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Wed, 5 Aug 2026 00:44:28 +0800 Subject: [PATCH 5/5] scala style --- .../org/apache/spark/sql/comet/CometDecimalPromotionSuite.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometDecimalPromotionSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometDecimalPromotionSuite.scala index eaca6d4030..ff70e1ea41 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/CometDecimalPromotionSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometDecimalPromotionSuite.scala @@ -93,7 +93,7 @@ class CometDecimalPromotionSuite extends CometTestBase { } test("decimal Divide with a non-decimal operand is unsupported") { - // This is just a sanity check. Spark should not allow this case to be created cause type coercion + // This is only a sanity check; Spark's type coercion should prevent this case. val decimal = AttributeReference("decimal", DecimalType(10, 0))() val integer = AttributeReference("integer", IntegerType)() val divide = Divide(decimal, integer, EvalMode.LEGACY)