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 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 8023231b35..5535f1540a 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 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) + } + } override def convert( expr: Divide, @@ -263,7 +271,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 +280,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/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..ff70e1ea41 --- /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 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) + + 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 37189714ad..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 @@ -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(