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 0d91010f1c8..0f30d3d2d72 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala @@ -138,9 +138,10 @@ trait MathBase { * `returnType`. Rebalances a flattened chain into an `O(log n)`-depth tree so deep `a + b + * ...` chains serialize to a shallow proto instead of a left-deep one that overflows protobuf's * recursion limit when the plan is re-parsed. Only safe for exactly-associative chains -- - * callers gate via [[isAssociativeAndRebalanceable]]. The flattened leaves all share the - * chain's type (Spark coerces operands to it, with casts acting as flatten boundaries), so a - * single `returnType` / `evalMode` is correct for every inner node. + * callers gate via [[isAssociativeAndRebalanceable]]. Spark coerces operands to the chain's + * type. Flattening stops at casts and at nested `Add` or `Multiply` expressions with a + * different `evalMode`, so every rebuilt inner node can use the chain's `returnType` and + * `evalMode`. */ def createBalancedMathExpr( operands: Seq[Expression], @@ -190,7 +191,7 @@ object CometAdd extends CometExpressionSerde[Add] with MathBase { // proto stays shallow and doesn't overflow protobuf's recursion limit when re-parsed. val operands = flattenAssociative( expr, - { case _: Add => true; case _ => false }, + { case a: Add => a.evalMode == expr.evalMode; case _ => false }, { case a: Add => (a.left, a.right) }) createBalancedMathExpr( operands, @@ -246,7 +247,7 @@ object CometMultiply extends CometExpressionSerde[Multiply] with MathBase { // proto stays shallow and doesn't overflow protobuf's recursion limit when re-parsed. val operands = flattenAssociative( expr, - { case _: Multiply => true; case _ => false }, + { case m: Multiply => m.evalMode == expr.evalMode; case _ => false }, { case m: Multiply => (m.left, m.right) }) createBalancedMathExpr( operands, diff --git a/spark/src/test/resources/sql-tests/expressions/math/arithmetic.sql b/spark/src/test/resources/sql-tests/expressions/math/arithmetic.sql index edce7d1b2b1..e8a3ee71879 100644 --- a/spark/src/test/resources/sql-tests/expressions/math/arithmetic.sql +++ b/spark/src/test/resources/sql-tests/expressions/math/arithmetic.sql @@ -45,6 +45,10 @@ INSERT INTO test_arith VALUES (10, 3, 10, 3, 10.5, 3.2, 10.5, 3.2), (0, 1, 0, 1, query SELECT a + b, a - b, a * b, a / b, a % b FROM test_arith +-- Keep the inner TRY overflow check when balancing legacy integer arithmetic. +query +SELECT try_add(a, b) + b, try_multiply(a, a) * b FROM test_arith + query SELECT la + lb, la - lb, la * lb, la / lb, la % lb FROM test_arith