0lai0 commented on code in PR #5050:
URL: https://github.com/apache/datafusion-comet/pull/5050#discussion_r4307675772
##########
spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala:
##########
@@ -1486,6 +1487,111 @@ class CometExpressionSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
}
}
+ test("division on negative-scale decimal reports unsupported via
getSupportLevel") {
+ // Issue #5013: `spark_decimal_div` computes its scale-alignment exponents
in `u32`, so a
+ // negative scale wraps to roughly 4.29e9 and the BigInt branch attempts
an unbounded
+ // allocation. Only division is affected. The other arithmetic operators
scale-align by
+ // `10^(max_scale - s)`, which is never negative, so they stay native and
are pinned by
+ // "safe ops on negative-scale decimal run natively" below. Constructing
the expressions
+ // requires allowNegativeScaleOfDecimal so DecimalType(_, s<0) passes
Spark's own check.
+ withSQLConf("spark.sql.legacy.allowNegativeScaleOfDecimal" -> "true") {
+ Seq(DecimalType(10, -1), DecimalType(20, -5)).foreach { negScale =>
+ val posDec = DecimalType(10, 0)
+ val negAttr = AttributeReference("n", negScale)()
+ val posAttr = AttributeReference("p", posDec)()
+ val expected =
Unsupported(Some(CometDivide.negScaleDecimalArithmeticReason))
+
+ assert(CometDivide.getSupportLevel(Divide(negAttr, posAttr)) ==
expected)
+ assert(CometIntegralDivide.getSupportLevel(IntegralDivide(negAttr,
posAttr)) == expected)
+ // Guard also fires when negative scale is only on the right operand.
+ assert(CometDivide.getSupportLevel(Divide(posAttr, negAttr)) ==
expected)
+
+ // The unguarded operators must not report unsupported for the same
operands.
+ assert(CometAdd.getSupportLevel(Add(negAttr, posAttr)) == Compatible())
+ assert(CometSubtract.getSupportLevel(Subtract(negAttr, posAttr)) ==
Compatible())
+ assert(CometRemainder.getSupportLevel(Remainder(negAttr, posAttr)) ==
Compatible())
+ assert(CometMultiply.getSupportLevel(Multiply(negAttr, posAttr)) ==
Compatible())
+ }
+ }
+ }
+
+ test("division on negative-scale decimal falls back and returns correct
results") {
+ // End-to-end proof for issue #5013 that the previously-panicking queries
now complete.
+ // `checkSparkAnswerAndFallbackReason` rather than `checkSparkAnswer`:
these five serdes do
+ // not mix in `CodegenDispatchFallback`, so `Unsupported` falls the whole
projection back to
+ // Spark, and asserting the reason is what proves the guard fired. A plain
answer comparison
+ // would pass whether or not it did.
+ // ConvertToLocalRelation is excluded so the arithmetic actually runs on
the plan rather
+ // than being folded at plan time (#4789). `DecimalType(_, s<0)` must be
constructed under
+ // allowNegativeScaleOfDecimal=true because the case class initializer
reads the flag, so
+ // wrap everything (including the type value) in one block.
+ withSQLConf(
+ "spark.sql.legacy.allowNegativeScaleOfDecimal" -> "true",
+ CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true",
+ SQLConf.OPTIMIZER_EXCLUDED_RULES.key ->
+ "org.apache.spark.sql.catalyst.optimizer.ConvertToLocalRelation") {
+ // Cover a narrow and a wide negative-scale type: at precision 20 the
combined operand
+ // width crosses DECIMAL128_MAX_PRECISION, which routes the operation to
+ // `WideDecimalBinaryExpr` / `decimal_div`'s BigInt path rather than the
narrow arrow
+ // kernels. `unit` is the smallest magnitude representable at that scale.
+ val negScaleCases = Seq((DecimalType(10, -1), 10), (DecimalType(20, -5),
100000))
+ negScaleCases.foreach { case (negScaleType, unit) =>
+ // Build the neg-scale column via string -> Decimal(neg), a known-safe
cast, so the
+ // source construction does not depend on any of the guards under test.
+ val strs = Seq(10 * unit, 20 * unit, 30 * unit).map(_.toString)
+ val v = strs.toDF("s").select(col("s").cast(negScaleType).as("v"))
+ Seq("v / cast(3 as decimal(10, 0))", "v div cast(3 as decimal(10,
0))").foreach { e =>
+ checkSparkAnswerAndFallbackReason(
+ v.selectExpr(e),
+ CometDivide.negScaleDecimalArithmeticReason)
+ }
+ }
+ }
+ }
+
+ test("safe ops on negative-scale decimal run natively (regression pin)") {
+ // Multiply and UnaryMinus on negative-scale decimals don't scale-align an
integer and
+ // ran natively at the time #5013 was fixed. The arithmetic guard
intentionally does
+ // not cover them. If a future change ever adds scale-alignment to
Multiply or a
+ // decimal codepath is added to UnaryMinus that does the multiply-by-10^n
trick,
+ // these assertions will start returning a Comet fallback plan and this
test will
+ // catch it before it reaches production as a silent-wrap or debug panic.
+ withSQLConf(
+ "spark.sql.legacy.allowNegativeScaleOfDecimal" -> "true",
+ CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true",
+ SQLConf.OPTIMIZER_EXCLUDED_RULES.key ->
+ "org.apache.spark.sql.catalyst.optimizer.ConvertToLocalRelation") {
+ // The wide case matters most for `v * v`: Decimal(20,-5) *
Decimal(20,-5) has
+ // p1 + p2 >= DECIMAL128_MAX_PRECISION, so it takes
`WideDecimalBinaryExpr`'s multiply
+ // branch (which does compute a scale adjustment against the output
precision/scale)
+ // and its result type is Decimal(38,-10) -- a negative *output* scale,
not just
+ // negative inputs.
+ val negScaleCases = Seq((DecimalType(10, -1), 10), (DecimalType(20, -5),
100000))
+ negScaleCases.foreach { case (negScaleType, unit) =>
+ val strs = Seq(10 * unit, 20 * unit, 30 * unit).map(_.toString)
+ val v = strs.toDF("s").select(col("s").cast(negScaleType).as("v"))
+ // Same-scale Multiply - no alignment needed, must stay native.
+ checkSparkAnswerAndOperator(v.selectExpr("v * v"))
+ checkSparkAnswerAndOperator(v.selectExpr("v * cast(2 as int)"))
+ // UnaryMinus on Decimal(neg) - no alignment, must stay native.
+ checkSparkAnswerAndOperator(v.selectExpr("-v"))
+ // Add, Subtract and Remainder scale-align by `10^(max_scale - s)`,
which is never
+ // negative. They were guarded in an earlier revision of this PR
because the `u8`
+ // underflow in the two width tests made them fail. Those are
rewritten in `i16` here,
+ // so they run natively again. If either rewrite is reverted these
assertions fail.
Review Comment:
You are right that the Scala assertions cannot pin the rewrites with
`overflow-checks = false`. I have reworded that comment rather than leaving it
claiming something it cannot do.
Three tests added, which do run with overflow checks on:
- `test_decimal128_add_sub_width_handles_negative_scale`, covering the
negative-scale cases plus a few non-negative ones to pin the unchanged behavior
- `test_modulo_decimal_negative_scale`, next to `test_modulo_basic_decimal`
- `test_decimal_div_rejects_negative_scale`, next to
`test_decimal_divide_by_zero_returns_spark_error`, covering all three positions
a negative scale can appear in
On sharing the calculation, I did not do it quite as suggested, because the
two sites use different boundaries: `create_binary_expr_with_options` tests `>=
DECIMAL128_MAX_PRECISION` and `create_modulo_expr` tests `>`. A shared
predicate would have changed one of them. So the shared piece is the
calculation, `decimal128_add_sub_width(p1, s1, p2, s2) -> i16`, and each site
keeps its own comparison. That still leaves one copy of the expression and
makes it testable.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
