From 447be100e578760ece3918b58407c74d2b7fcd64 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 6 Sep 2026 14:21:11 +0800 Subject: [PATCH 1/2] perf: avoid repeated decimal promotion in expression serialization --- .../apache/comet/serde/QueryPlanSerde.scala | 13 +++- .../scala/org/apache/comet/serde/arrays.scala | 58 ++++++++--------- .../org/apache/comet/serde/bitwise.scala | 6 +- .../comet/CometDecimalPromotionSuite.scala | 63 +++++++++++++++++-- 4 files changed, 102 insertions(+), 38 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 70646ce2e4..fbd8bfe998 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -740,6 +740,8 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { } } + // Aggregate functions bypass exprToProto: their arguments and filters are independent roots + // and must enter through exprToProto to receive decimal promotion. def aggExprToProto( aggExpr: AggregateExpression, inputs: Seq[Attribute], @@ -844,7 +846,9 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { * expression. * * This method performs a transformation on the plan to handle decimal promotion and then calls - * into the recursive method [[exprToProtoInternal]]. + * into the recursive method [[exprToProtoInternal]]. Use this entry point for independent roots + * (including aggregate arguments and filters) and synthesized trees needing decimal promotion. + * Serdes must use [[exprToProtoInternal]] for children of the already-promoted tree. * * @param expr * The input expression @@ -913,6 +917,11 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { * Convert a Spark expression to a protocol-buffer representation of a native Comet/DataFusion * expression. * + * The caller owns decimal promotion: this method serializes children of an already-promoted + * root without traversing them again. Literals and wrappers that introduce no decimal + * arithmetic can also use this path. Newly synthesized arithmetic must enter through + * [[exprToProto]]. + * * @param expr * The input expression * @param inputs @@ -1062,7 +1071,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { binding: Boolean, f: (ExprOuterClass.Expr.Builder, ExprOuterClass.UnaryExpr) => ExprOuterClass.Expr.Builder) : Option[ExprOuterClass.Expr] = { - val childExpr = exprToProtoInternal(child, inputs, binding) // TODO review + val childExpr = exprToProtoInternal(child, inputs, binding) if (childExpr.isDefined) { // create the generic UnaryExpr message val inner = ExprOuterClass.UnaryExpr diff --git a/spark/src/main/scala/org/apache/comet/serde/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala index 27a742ce58..18622ae094 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -44,8 +44,8 @@ object CometArrayRemove expr: ArrayRemove, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - val arrayExprProto = exprToProto(expr.left, inputs, binding) - val keyExprProto = exprToProto(expr.right, inputs, binding) + val arrayExprProto = exprToProtoInternal(expr.left, inputs, binding) + val keyExprProto = exprToProtoInternal(expr.right, inputs, binding) scalarFunctionExprToProto("array_remove_all", arrayExprProto, keyExprProto) } @@ -60,8 +60,8 @@ object CometArrayAppend extends CometExpressionSerde[ArrayAppend] { val child = expr.children.head val elementType = child.dataType.asInstanceOf[ArrayType].elementType - val arrayExprProto = exprToProto(expr.children.head, inputs, binding) - val keyExprProto = exprToProto(expr.children(1), inputs, binding) + val arrayExprProto = exprToProtoInternal(expr.children.head, inputs, binding) + val keyExprProto = exprToProtoInternal(expr.children(1), inputs, binding) // DataFusion's array_append always returns a list with nullable elements, // so we must promise ArrayType(elementType, containsNull = true) here even if @@ -128,8 +128,8 @@ object CometArrayContains expr: ArrayContains, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - val arrayExprProto = exprToProto(expr.children.head, inputs, binding) - val keyExprProto = exprToProto(expr.children(1), inputs, binding) + val arrayExprProto = exprToProtoInternal(expr.children.head, inputs, binding) + val keyExprProto = exprToProtoInternal(expr.children(1), inputs, binding) scalarFunctionExprToProto("array_contains", arrayExprProto, keyExprProto) } @@ -228,8 +228,8 @@ object CometArrayIntersect expr: ArrayIntersect, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - val leftArrayExprProto = exprToProto(expr.children.head, inputs, binding) - val rightArrayExprProto = exprToProto(expr.children(1), inputs, binding) + val leftArrayExprProto = exprToProtoInternal(expr.children.head, inputs, binding) + val rightArrayExprProto = exprToProtoInternal(expr.children(1), inputs, binding) val arraysIntersectScalarExpr = scalarFunctionExprToProto("array_intersect", leftArrayExprProto, rightArrayExprProto) @@ -242,7 +242,7 @@ object CometArrayMax extends CometExpressionSerde[ArrayMax] { expr: ArrayMax, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - val arrayExprProto = exprToProto(expr.children.head, inputs, binding) + val arrayExprProto = exprToProtoInternal(expr.children.head, inputs, binding) val arrayMaxScalarExpr = scalarFunctionExprToProto("array_max", arrayExprProto) @@ -255,7 +255,7 @@ object CometArrayMin extends CometExpressionSerde[ArrayMin] { expr: ArrayMin, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - val arrayExprProto = exprToProto(expr.children.head, inputs, binding) + val arrayExprProto = exprToProtoInternal(expr.children.head, inputs, binding) val arrayMinScalarExpr = scalarFunctionExprToProto("array_min", arrayExprProto) arrayMinScalarExpr @@ -267,8 +267,8 @@ object CometArraysOverlap extends CometExpressionSerde[ArraysOverlap] { expr: ArraysOverlap, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - val leftArrayExprProto = exprToProto(expr.left, inputs, binding) - val rightArrayExprProto = exprToProto(expr.right, inputs, binding) + val leftArrayExprProto = exprToProtoInternal(expr.left, inputs, binding) + val rightArrayExprProto = exprToProtoInternal(expr.right, inputs, binding) val arraysOverlapScalarExpr = scalarFunctionExprToProtoWithReturnType( "spark_arrays_overlap", @@ -287,7 +287,7 @@ object CometArrayCompact extends CometExpressionSerde[Expression] { inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { val child = expr.children.head - val arrayExprProto = exprToProto(child, inputs, binding) + val arrayExprProto = exprToProtoInternal(child, inputs, binding) val arrayCompactScalarExpr = scalarFunctionExprToProto("array_compact", arrayExprProto) arrayCompactScalarExpr @@ -343,8 +343,8 @@ object CometArrayExcept return None case None => } - val leftArrayExprProto = exprToProto(expr.left, inputs, binding) - val rightArrayExprProto = exprToProto(expr.right, inputs, binding) + val leftArrayExprProto = exprToProtoInternal(expr.left, inputs, binding) + val rightArrayExprProto = exprToProtoInternal(expr.right, inputs, binding) val arrayExceptScalarExpr = scalarFunctionExprToProto("array_except", leftArrayExprProto, rightArrayExprProto) @@ -400,8 +400,8 @@ object CometArrayJoin expr: ArrayJoin, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - val arrayExprProto = exprToProto(expr.array, inputs, binding) - val delimiterExprProto = exprToProto(expr.delimiter, inputs, binding) + val arrayExprProto = exprToProtoInternal(expr.array, inputs, binding) + val delimiterExprProto = exprToProtoInternal(expr.delimiter, inputs, binding) val joined = expr.nullReplacement match { case Some(nullReplacementExpr) => @@ -409,7 +409,7 @@ object CometArrayJoin "array_to_string", arrayExprProto, delimiterExprProto, - exprToProto(nullReplacementExpr, inputs, binding)) + exprToProtoInternal(nullReplacementExpr, inputs, binding)) case None => scalarFunctionExprToProto("array_to_string", arrayExprProto, delimiterExprProto) } @@ -498,7 +498,7 @@ object CometSlice extends CometExpressionSerde[Slice] { inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { val elementType = expr.x.dataType.asInstanceOf[ArrayType].elementType - val arrayExprProto = exprToProto(expr.x, inputs, binding) + val arrayExprProto = exprToProtoInternal(expr.x, inputs, binding) val startExprProto = exprToProto(Cast(expr.start, LongType), inputs, binding) val lengthExprProto = exprToProto(Cast(expr.length, LongType), inputs, binding) // DataFusion list types always have nullable inner elements, so promise @@ -521,8 +521,8 @@ object CometArrayUnion extends CometExpressionSerde[ArrayUnion] { expr: ArrayUnion, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - val leftArrayExprProto = exprToProto(expr.children.head, inputs, binding) - val rightArrayExprProto = exprToProto(expr.children(1), inputs, binding) + val leftArrayExprProto = exprToProtoInternal(expr.children.head, inputs, binding) + val rightArrayExprProto = exprToProtoInternal(expr.children(1), inputs, binding) val arraysUnionScalarExpr = scalarFunctionExprToProto("array_union", leftArrayExprProto, rightArrayExprProto) @@ -629,7 +629,7 @@ object CometArrayReverse extends CometExpressionSerde[Reverse] with ArraysBase { withFallbackReason(expr, s"child data type not supported: ${expr.child.dataType}") return None } - val reverseExprProto = exprToProto(expr.child, inputs, binding) + val reverseExprProto = exprToProtoInternal(expr.child, inputs, binding) val reverseScalarExpr = scalarFunctionExprToProto("array_reverse", reverseExprProto) reverseScalarExpr } @@ -735,7 +735,7 @@ object CometFlatten extends CometExpressionSerde[Flatten] with ArraysBase { expr: Flatten, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - val flattenExprProto = exprToProto(expr.child, inputs, binding) + val flattenExprProto = exprToProtoInternal(expr.child, inputs, binding) val flattenScalarExpr = scalarFunctionExprToProto("flatten", flattenExprProto) flattenScalarExpr } @@ -780,7 +780,7 @@ object CometSize extends CometExpressionSerde[Size] { expr: Size, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - val arrayExprProto = exprToProto(expr.child, inputs, binding) + val arrayExprProto = exprToProtoInternal(expr.child, inputs, binding) for { isNotNullExprProto <- createIsNotNullExprProto(expr, inputs, binding) sizeScalarExprProto <- scalarFunctionExprToProto("size", arrayExprProto) @@ -833,8 +833,8 @@ object CometArrayPosition extends CometExpressionSerde[ArrayPosition] with Array expr: ArrayPosition, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - val arrayExprProto = exprToProto(expr.left, inputs, binding) - val elementExprProto = exprToProto(expr.right, inputs, binding) + val arrayExprProto = exprToProtoInternal(expr.left, inputs, binding) + val elementExprProto = exprToProtoInternal(expr.right, inputs, binding) // Use spark_array_position which returns Int64 and 0 when not found // (matching Spark's behavior) @@ -999,12 +999,12 @@ object CometSequence extends CometExpressionSerde[Sequence] with CodegenDispatch expr: Sequence, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - val startExprProto = exprToProto(expr.start, inputs, binding) - val stopExprProto = exprToProto(expr.stop, inputs, binding) + val startExprProto = exprToProtoInternal(expr.start, inputs, binding) + val stopExprProto = exprToProtoInternal(expr.stop, inputs, binding) // With no step argument the native kernel computes Spark's per-row default, // `start <= stop ? 1 : -1`, which cannot be expressed as a plan-time literal. val argProtos = Seq(startExprProto, stopExprProto) ++ - expr.stepOpt.map(exprToProto(_, inputs, binding)) + expr.stepOpt.map(exprToProtoInternal(_, inputs, binding)) scalarFunctionExprToProtoWithReturnType( "spark_sequence", expr.dataType, diff --git a/spark/src/main/scala/org/apache/comet/serde/bitwise.scala b/spark/src/main/scala/org/apache/comet/serde/bitwise.scala index 115bd80422..ec06c71253 100644 --- a/spark/src/main/scala/org/apache/comet/serde/bitwise.scala +++ b/spark/src/main/scala/org/apache/comet/serde/bitwise.scala @@ -50,7 +50,7 @@ object CometBitwiseNot extends CometExpressionSerde[BitwiseNot] { expr: BitwiseNot, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - val childProto = exprToProto(expr.child, inputs, binding) + val childProto = exprToProtoInternal(expr.child, inputs, binding) val bitNotScalarExpr = scalarFunctionExprToProto("bitwise_not", childProto) bitNotScalarExpr @@ -144,8 +144,8 @@ object CometBitwiseGet extends CometExpressionSerde[BitwiseGet] { expr: BitwiseGet, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - val argProto = exprToProto(expr.left, inputs, binding) - val posProto = exprToProto(expr.right, inputs, binding) + val argProto = exprToProtoInternal(expr.left, inputs, binding) + val posProto = exprToProtoInternal(expr.right, inputs, binding) val bitGetScalarExpr = scalarFunctionExprToProtoWithReturnType("bit_get", ByteType, false, argProto, posProto) bitGetScalarExpr 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 346e7bef27..0dbfc14750 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 @@ -20,7 +20,8 @@ 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.catalyst.expressions.{Add, ArrayContains, AttributeReference, BitwiseNot, Cast, CreateArray, Divide, EvalMode, Multiply, NamedExpression} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Partial, Sum} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DecimalType, IntegerType} @@ -62,9 +63,7 @@ class CometDecimalPromotionSuite extends CometTestBase { DecimalPrecision.promote(promoted) == promoted, s"$name promotion is not idempotent: $promoted") - // This proto-shape check relies on CometArrayContains re-entering exprToProto for its - // children. If https://github.com/apache/datafusion-comet/issues/5248 changes that, - // re-point it to another recursively serializing serde. + // Recursive child serialization must retain exactly one equivalent overflow wrapper. val arithmeticProto = QueryPlanSerde .exprToProto(expression, plan.children.head.output) .get @@ -95,6 +94,62 @@ class CometDecimalPromotionSuite extends CometTestBase { s"try_divide($left, $right)") } + test("issue #5248: nested decimal children and aggregate roots retain overflow wrappers") { + val left = AttributeReference("left", DecimalType(10, 0))() + val right = AttributeReference("right", DecimalType(10, 0))() + val inputs = Seq(left, right) + + Seq(EvalMode.LEGACY, EvalMode.ANSI, EvalMode.TRY).foreach { mode => + val multiply = Multiply(left, right, mode) + val arithmetic = Add(multiply, right, mode) + val contains = ArrayContains(CreateArray(Seq(arithmetic)), arithmetic) + + def check(proto: ExprOuterClass.Expr): Unit = { + assert(proto.hasCheckOverflow, s"$mode: $proto") + val outer = proto.getCheckOverflow + assert(outer.getDatatype === QueryPlanSerde.serializeDataType(arithmetic.dataType).get) + assert(outer.getFailOnError === (mode == EvalMode.ANSI)) + assert(outer.getChild.hasAdd, s"Duplicate outer CheckOverflow: $proto") + val inner = outer.getChild.getAdd.getLeft + assert(inner.hasCheckOverflow, s"Missing nested CheckOverflow: $proto") + assert( + inner.getCheckOverflow.getDatatype === + QueryPlanSerde.serializeDataType(multiply.dataType).get) + assert(inner.getCheckOverflow.getFailOnError === (mode == EvalMode.ANSI)) + assert( + inner.getCheckOverflow.getChild.hasMultiply, + s"Duplicate inner CheckOverflow: $proto") + } + + // Exercise both array child paths, including CreateArray's recursive serialization. + Seq(true, false).foreach { binding => + val proto = QueryPlanSerde.exprToProto(contains, inputs, binding).get.getScalarFunc + check(proto.getArgs(0).getScalarFunc.getArgs(0)) + check(proto.getArgs(1)) + val bitwise = BitwiseNot(Cast(arithmetic, IntegerType)) + check( + QueryPlanSerde + .exprToProto(bitwise, inputs, binding) + .get + .getScalarFunc + .getArgs(0) + .getCast + .getChild) + } + + // Aggregate serialization does not promote the aggregate tree before visiting its inputs. + val aggregate = AggregateExpression( + Sum(arithmetic), + Partial, + false, + Some(contains), + NamedExpression.newExprId) + val proto = QueryPlanSerde.aggExprToProto(aggregate, inputs, true, SQLConf.get).get + check(proto.getSum.getChild) + check(proto.getFilter.getScalarFunc.getArgs(1)) + } + } + 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))() From d3c44512fa2379964e5d7beba831b396e2bf4545 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Mon, 7 Sep 2026 00:49:24 +0800 Subject: [PATCH 2/2] perf: finish decimal promotion audit and verify overflow results --- .../apache/comet/serde/QueryPlanSerde.scala | 6 +- .../scala/org/apache/comet/serde/arrays.scala | 19 ++-- .../comet/CometDecimalPromotionSuite.scala | 90 +++++++++++++++++-- 3 files changed, 99 insertions(+), 16 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 fbd8bfe998..6974f4c7f5 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -740,8 +740,10 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { } } - // Aggregate functions bypass exprToProto: their arguments and filters are independent roots - // and must enter through exprToProto to receive decimal promotion. + /** + * This method does not promote the aggregate tree. Its arguments and filters are independent + * roots and must be serialized through [[exprToProto]] to receive decimal promotion. + */ def aggExprToProto( aggExpr: AggregateExpression, inputs: Seq[Attribute], diff --git a/spark/src/main/scala/org/apache/comet/serde/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala index 18622ae094..8da696da46 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -81,7 +81,8 @@ object CometArrayAppend extends CometExpressionSerde[ArrayAppend] { binding, (builder, unaryExpr) => builder.setIsNotNull(unaryExpr)) - val nullLiteralProto = exprToProto(Literal(null, elementType), Seq.empty) + val nullLiteralProto = + exprToProtoInternal(Literal(null, elementType), Seq.empty, binding = true) if (arrayAppendScalarExpr.isDefined && isNotNullExpr.isDefined && nullLiteralProto.isDefined) { val caseWhenExpr = ExprOuterClass.CaseWhen @@ -420,8 +421,8 @@ object CometArrayJoin case Some(nullReplacementExpr) => for { innerProto <- joined - replacementIsNull <- exprToProto(IsNull(nullReplacementExpr), inputs, binding) - nullLiteral <- exprToProto(Literal(null, expr.dataType), inputs, binding) + replacementIsNull <- exprToProtoInternal(IsNull(nullReplacementExpr), inputs, binding) + nullLiteral <- exprToProtoInternal(Literal(null, expr.dataType), inputs, binding) } yield ExprOuterClass.Expr .newBuilder() .setIf( @@ -499,8 +500,8 @@ object CometSlice extends CometExpressionSerde[Slice] { binding: Boolean): Option[ExprOuterClass.Expr] = { val elementType = expr.x.dataType.asInstanceOf[ArrayType].elementType val arrayExprProto = exprToProtoInternal(expr.x, inputs, binding) - val startExprProto = exprToProto(Cast(expr.start, LongType), inputs, binding) - val lengthExprProto = exprToProto(Cast(expr.length, LongType), inputs, binding) + val startExprProto = exprToProtoInternal(Cast(expr.start, LongType), inputs, binding) + val lengthExprProto = exprToProtoInternal(Cast(expr.length, LongType), inputs, binding) // DataFusion list types always have nullable inner elements, so promise // ArrayType(elementType, containsNull = true) here even if Spark's // expr.dataType reports containsNull = false (e.g. for array(1, 2, 3)). @@ -697,7 +698,8 @@ object CometElementAt extends CometExpressionSerde[ElementAt] { inputs, binding, (builder, unaryExpr) => builder.setIsNotNull(unaryExpr)) - val nullLiteralProto = exprToProto(Literal(null, expr.dataType), Seq.empty) + val nullLiteralProto = + exprToProtoInternal(Literal(null, expr.dataType), Seq.empty, binding = true) for { base <- baseExpr notNull <- isNotNullExpr @@ -813,7 +815,7 @@ object CometSize extends CometExpressionSerde[Size] { private def createLiteralExprProto(legacySizeOfNull: Boolean): Option[ExprOuterClass.Expr] = { val value = if (legacySizeOfNull) -1 else null - exprToProto(Literal(value, IntegerType), Seq.empty) + exprToProtoInternal(Literal(value, IntegerType), Seq.empty, binding = true) } } @@ -884,7 +886,8 @@ object CometArraysZip extends CometExpressionSerde[ArraysZip] { // mimic Spark's ArraysZip behavior: returns NULL if any argument is NULL val combinedNullCheck = expr.children.map(child => IsNotNull(child)).reduce(And) val isNotNullExpr = exprToProtoInternal(combinedNullCheck, inputs, binding) - val nullLiteralProto = exprToProto(Literal(null, expr.dataType), Seq.empty) + val nullLiteralProto = + exprToProtoInternal(Literal(null, expr.dataType), Seq.empty, binding = true) if (exprChildren.forall( _.isDefined) && isNotNullExpr.isDefined && nullLiteralProto.isDefined) { 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 0dbfc14750..8181628242 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 @@ -25,6 +25,7 @@ import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DecimalType, IntegerType} +import org.apache.comet.CometConf import org.apache.comet.serde.{CometDivide, ExprOuterClass, QueryPlanSerde, Unsupported} class CometDecimalPromotionSuite extends CometTestBase { @@ -63,9 +64,10 @@ class CometDecimalPromotionSuite extends CometTestBase { DecimalPrecision.promote(promoted) == promoted, s"$name promotion is not idempotent: $promoted") - // Recursive child serialization must retain exactly one equivalent overflow wrapper. + // Deliberately re-enter the public serializer with an already-promoted tree. Recursive + // serdes no longer do this, but re-promotion must still preserve the protobuf shape. val arithmeticProto = QueryPlanSerde - .exprToProto(expression, plan.children.head.output) + .exprToProto(promoted, plan.children.head.output) .get .getScalarFunc .getArgs(1) @@ -104,7 +106,7 @@ class CometDecimalPromotionSuite extends CometTestBase { val arithmetic = Add(multiply, right, mode) val contains = ArrayContains(CreateArray(Seq(arithmetic)), arithmetic) - def check(proto: ExprOuterClass.Expr): Unit = { + def check(proto: ExprOuterClass.Expr, binding: Boolean = true): Unit = { assert(proto.hasCheckOverflow, s"$mode: $proto") val outer = proto.getCheckOverflow assert(outer.getDatatype === QueryPlanSerde.serializeDataType(arithmetic.dataType).get) @@ -119,13 +121,19 @@ class CometDecimalPromotionSuite extends CometTestBase { assert( inner.getCheckOverflow.getChild.hasMultiply, s"Duplicate inner CheckOverflow: $proto") + val reference = inner.getCheckOverflow.getChild.getMultiply.getLeft + if (binding) { + assert(reference.hasBound && reference.getBound.getIndex == 0) + } else { + assert(reference.hasUnbound && reference.getUnbound.getName == "left") + } } // Exercise both array child paths, including CreateArray's recursive serialization. Seq(true, false).foreach { binding => val proto = QueryPlanSerde.exprToProto(contains, inputs, binding).get.getScalarFunc - check(proto.getArgs(0).getScalarFunc.getArgs(0)) - check(proto.getArgs(1)) + check(proto.getArgs(0).getScalarFunc.getArgs(0), binding) + check(proto.getArgs(1), binding) val bitwise = BitwiseNot(Cast(arithmetic, IntegerType)) check( QueryPlanSerde @@ -134,7 +142,8 @@ class CometDecimalPromotionSuite extends CometTestBase { .getScalarFunc .getArgs(0) .getCast - .getChild) + .getChild, + binding) } // Aggregate serialization does not promote the aggregate tree before visiting its inputs. @@ -150,6 +159,75 @@ class CometDecimalPromotionSuite extends CometTestBase { } } + Seq(false, true).foreach { ansi => + test(s"issue #5248: decimal overflow values under recursive serdes, ANSI=$ansi") { + withSQLConf( + SQLConf.ANSI_ENABLED.key -> ansi.toString, + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false", + CometConf.getExprAllowIncompatConfigKey("ArrayIntersect") -> "true", + CometConf.getExprAllowIncompatConfigKey("ArrayExcept") -> "true", + CometConf.getExprAllowIncompatConfigKey("ArrayJoin") -> "true") { + withTempPath { path => + // Read actual decimal columns from Parquet so constant folding cannot hide promotion. + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sql(s"""SELECT CAST(a AS DECIMAL(38,0)) a, CAST(b AS DECIMAL(38,0)) b, + |CAST(c AS DECIMAL(38,6)) c, CAST(d AS DECIMAL(38,6)) d + |FROM VALUES + |('${"9" * 38}', '2', '${"9" * 32}.999999', '0.000001'), + |('4', '2', '4', '2'), (NULL, NULL, NULL, NULL) AS t(a,b,c,d) + |""".stripMargin).write.parquet(path.toString) + } + withParquetTable(path.toString, "decimal_overflow") { + val expressions = Seq( + "array_remove(array($e), $e)", + "array_append(array($e), $e)", + "array_contains(array($e), $e)", + "array_intersect(array($e), array($e))", + "array_max(array($e))", + "array_min(array($e))", + "arrays_overlap(array($e), array($e))", + "array_compact(array($e))", + "array_except(array($e), array($e))", + "array_join(array(CAST($e AS STRING)), ',')", + // Spark's ArrayJoin codegen needs a nullable array or delimiter to clear isNull + // when the replacement is nullable. Use a column-based delimiter for this case. + "array_join(array('x', NULL), CAST(a AS STRING), CAST($e AS STRING))", + "slice(array($e), 1, 1)", + "slice(array(1, 2), CAST($e AS INT), 2)", + "slice(array(1, 2), 1, CAST($e AS INT))", + "array_union(array($e), array($e))", + "reverse(array($e))", + "flatten(array(array($e)))", + "size(array($e))", + "array_position(array($e), $e)", + "~CAST($e AS BIGINT)", + "bit_get(CAST($e AS BIGINT), 1)", + "element_at(array($e), 1)", + "arrays_zip(array($e), array($e))") + // Decimal division's overflow sentinel needs CheckOverflow to become NULL in LEGACY. + for (arithmetic <- Seq("a * b", "c / d"); expression <- expressions) { + val query = s"SELECT ${expression.replace("$e", arithmetic)} FROM decimal_overflow" + withClue(query) { + if (ansi) { + val df = sql(query) + assert(df.queryExecution.executedPlan.collect { case _: CometProjectExec => + true + }.nonEmpty) + val (sparkError, cometError) = checkSparkAnswerMaybeThrows(df) + assert(sparkError.isDefined == cometError.isDefined) + } else { + checkSparkAnswerAndOperator( + sql(query), + includeClasses = Seq(classOf[CometNativeScanExec])) + } + } + } + } + } + } + } + } + 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))()