diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 2d01ecc104..e7071b94a7 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.DecimalPrecisionSuite 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 5f60c07451..15c5deb363 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.DecimalPrecisionSuite 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/spark/sql/comet/DecimalPrecision.scala b/spark/src/main/scala/org/apache/spark/sql/comet/DecimalPrecision.scala index 2d5247df7e..42a30db330 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 @@ -40,6 +40,24 @@ import org.apache.spark.sql.types.DecimalType */ object DecimalPrecision { def promote(expr: Expression): Expression = { + // `transformUp` walks and rebuilds every node even when no case matches, and the serde calls + // this once per expression it converts, so skip it when there is nothing to rewrite. + if (containsDecimalArithmetic(expr)) rewrite(expr) else expr + } + + /** + * Whether [[promote]] could rewrite anything in this tree. Deliberately a superset of the rule: + * it ignores the operands, so any arithmetic node the rule handles that produces a decimal + * qualifies. Wrapping a child in `CheckOverflow` does not change that child's data type, so a + * tree without such a node cannot grow one part way through the rewrite. + */ + private def containsDecimalArithmetic(expr: Expression): Boolean = expr.exists { + case e @ (_: Add | _: Subtract | _: Multiply | _: Divide | _: Remainder) => + e.dataType.isInstanceOf[DecimalType] + case _ => false + } + + private def rewrite(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 diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/DecimalPrecisionSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/DecimalPrecisionSuite.scala new file mode 100644 index 0000000000..be4c7db5de --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/DecimalPrecisionSuite.scala @@ -0,0 +1,83 @@ +/* + * 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.SparkFunSuite +import org.apache.spark.sql.catalyst.expressions.{Add, Alias, AttributeReference, CaseWhen, Cast, CheckOverflow, Coalesce, Divide, Expression, GreaterThan, IsNotNull, Literal, Multiply, Remainder, Subtract} +import org.apache.spark.sql.types.{DecimalType, IntegerType, LongType, StringType} + +class DecimalPrecisionSuite extends SparkFunSuite { + + private val decA = AttributeReference("a", DecimalType(10, 2))() + private val decB = AttributeReference("b", DecimalType(10, 2))() + private val intA = AttributeReference("c", IntegerType)() + private val intB = AttributeReference("d", IntegerType)() + + private def assertUnchanged(expr: Expression): Unit = { + // Reference equality, not `===`: the point of the guard is that no node is rebuilt. + val promoted = DecimalPrecision.promote(expr) + assert(promoted.eq(expr), s"expected promote to leave $expr untouched, got: $promoted") + } + + test("promote leaves trees without decimal arithmetic untouched") { + Seq( + decA, + Literal(1), + Add(intA, intB), + Multiply(intA, Cast(decA, IntegerType)), + GreaterThan(decA, decB), + Cast(Add(intA, intB), DecimalType(20, 4)), + Coalesce(Seq(decA, decB)), + CaseWhen(Seq((IsNotNull(decA), decA)), Some(decB)), + Alias(Add(Cast(decA, LongType), Cast(decB, LongType)), "x")()).foreach(assertUnchanged) + } + + test("promote wraps decimal arithmetic wherever it appears in the tree") { + Seq[Expression]( + Add(decA, decB), + Subtract(decA, decB), + Multiply(decA, decB), + Divide(decA, decB), + Remainder(decA, decB)).foreach { arithmetic => + DecimalPrecision.promote(arithmetic) match { + case CheckOverflow(child, dt, _) => + assert(child.eq(arithmetic)) + assert(dt === arithmetic.dataType) + case other => fail(s"expected $arithmetic to be wrapped in CheckOverflow, got: $other") + } + + // The guard has to inspect the whole tree, not just the root: a decimal arithmetic node + // buried under expressions the rule does not rewrite must still be promoted. + val buried = Cast(Coalesce(Seq(decA, arithmetic)), StringType) + val promoted = DecimalPrecision.promote(buried) + assert(!promoted.eq(buried)) + assert(promoted.collectFirst { case c: CheckOverflow => c }.isDefined) + } + } + + test("promote rewrites only the decimal arithmetic in a mixed tree") { + val intAdd = Add(intA, intB) + val decAdd = Add(decA, decB) + val promoted = DecimalPrecision.promote(GreaterThan(Cast(intAdd, DecimalType(20, 2)), decAdd)) + val overflows = promoted.collect { case c: CheckOverflow => c } + assert(overflows.length === 1) + assert(overflows.head.child.eq(decAdd)) + } +}