fix: preserve TRY overflow semantics in nested integer arithmetic

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,

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