From 0cbf15d53c79136210827993a9f677872819a299 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sat, 31 Oct 2026 21:06:18 +0000 Subject: [PATCH 1/2] fix: preserve TRY overflow semantics in nested integer arithmetic --- spark/src/main/scala/org/apache/comet/serde/arithmetic.scala | 4 ++-- .../test/resources/sql-tests/expressions/math/arithmetic.sql | 4 ++++ 2 files changed, 6 insertions(+), 2 deletions(-) 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..51cedfef33a 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala @@ -190,7 +190,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 +246,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 From a1491a5e8a3385296b1e123ecefd19ae3d8e58c3 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Mon, 2 Nov 2026 05:40:56 +0000 Subject: [PATCH 2/2] docs: explain evaluation mode boundaries in arithmetic chains --- .../src/main/scala/org/apache/comet/serde/arithmetic.scala | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) 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 51cedfef33a..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],