perf: track decimal overflow without rescanning results

native/spark-expr/benches/wide_decimal.rs

@@ -19,7 +19,7 @@
 //! for Decimal128 arithmetic that requires wider intermediate precision.
 
 use arrow::array::builder::Decimal128Builder;
-use arrow::array::RecordBatch;
+use arrow::array::{Array, Decimal128Array, RecordBatch};
 use arrow::datatypes::{DataType, Field, Schema};
 use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
 use datafusion::logical_expr::Operator;
@@ -28,6 +28,7 @@ use datafusion::physical_expr::PhysicalExpr;
 use datafusion_comet_spark_expr::{
     Cast, EvalMode, SparkCastOptions, WideDecimalBinaryExpr, WideDecimalOp,
 };
+use std::hint::black_box;
 use std::sync::Arc;
 
 const BATCH_SIZE: usize = 8192;
@@ -49,6 +50,30 @@ fn make_decimal_batch(p1: u8, s1: i8, p2: u8, s2: i8) -> RecordBatch {
     RecordBatch::try_new(Arc::new(schema), vec![Arc::new(left), Arc::new(right)]).unwrap()
 }
 
+fn make_overflow_batch(null_every: usize, overflow_every: usize) -> RecordBatch {
+    let left: Decimal128Array = (0..BATCH_SIZE)
+        .map(|i| {
+            if null_every != 0 && i % null_every == 0 {
+                None
+            } else if overflow_every != 0 && (i + 1) % overflow_every == 0 {
+                Some(10_000_000_000)
+            } else {
+                Some((i as i128 % 100_000) * 100)
+            }
+        })
+        .collect::<Decimal128Array>()
+        .with_precision_and_scale(38, 2)
+        .unwrap();
+    let right = Decimal128Array::from_value(0, BATCH_SIZE)
+        .with_precision_and_scale(38, 2)
+        .unwrap();
+    let schema = Schema::new(vec![
+        Field::new("left", left.data_type().clone(), true),
+        Field::new("right", right.data_type().clone(), false),
+    ]);
+    RecordBatch::try_new(Arc::new(schema), vec![Arc::new(left), Arc::new(right)]).unwrap()
+}
+
 /// Old approach: Cast(Decimal128->Decimal256) both sides, BinaryExpr, Cast(Decimal256->Decimal128).
 fn build_old_expr(
     p1: u8,
@@ -80,16 +105,16 @@ fn build_old_expr(
 }
 
 /// New approach: single fused WideDecimalBinaryExpr.
-fn build_new_expr(op: WideDecimalOp, p_out: u8, s_out: i8) -> Arc<dyn PhysicalExpr> {
+fn build_new_expr(
+    op: WideDecimalOp,
+    p_out: u8,
+    s_out: i8,
+    eval_mode: EvalMode,
+) -> Arc<dyn PhysicalExpr> {
     let left_col: Arc<dyn PhysicalExpr> = Arc::new(Column::new("left", 0));
     let right_col: Arc<dyn PhysicalExpr> = Arc::new(Column::new("right", 1));
     Arc::new(WideDecimalBinaryExpr::new(
-        left_col,
-        right_col,
-        op,
-        p_out,
-        s_out,
-        EvalMode::Legacy,
+        left_col, right_col, op, p_out, s_out, eval_mode,
     ))
 }
 
@@ -116,15 +141,15 @@ fn criterion_benchmark(c: &mut Criterion) {
     {
         let batch = make_decimal_batch(38, 10, 38, 10);
         let old = build_old_expr(38, 10, 38, 10, Operator::Plus, DataType::Decimal128(38, 10));
-        let new = build_new_expr(WideDecimalOp::Add, 38, 10);
+        let new = build_new_expr(WideDecimalOp::Add, 38, 10, EvalMode::Legacy);
         bench_case(&mut group, "add_same_scale", &batch, &old, &new);
     }
 
     // Case 2: Add with different scales - Decimal128(38,6) + Decimal128(38,4) -> Decimal128(38,6)
     {
         let batch = make_decimal_batch(38, 6, 38, 4);
         let old = build_old_expr(38, 6, 38, 4, Operator::Plus, DataType::Decimal128(38, 6));
-        let new = build_new_expr(WideDecimalOp::Add, 38, 6);
+        let new = build_new_expr(WideDecimalOp::Add, 38, 6, EvalMode::Legacy);
         bench_case(&mut group, "add_diff_scale", &batch, &old, &new);
     }
 
@@ -140,7 +165,7 @@ fn criterion_benchmark(c: &mut Criterion) {
             Operator::Multiply,
             DataType::Decimal128(38, 6),
         );
-        let new = build_new_expr(WideDecimalOp::Multiply, 38, 6);
+        let new = build_new_expr(WideDecimalOp::Multiply, 38, 6, EvalMode::Legacy);
         bench_case(&mut group, "multiply", &batch, &old, &new);
     }
 
@@ -155,11 +180,38 @@ fn criterion_benchmark(c: &mut Criterion) {
             Operator::Minus,
             DataType::Decimal128(38, 18),
         );
-        let new = build_new_expr(WideDecimalOp::Subtract, 38, 18);
+        let new = build_new_expr(WideDecimalOp::Subtract, 38, 18, EvalMode::Legacy);
         bench_case(&mut group, "subtract", &batch, &old, &new);
     }
 
     group.finish();
+
+    let no_overflow = make_overflow_batch(0, 0);
+    let no_overflow_nulls = make_overflow_batch(17, 0);
+    let sparse_overflow = make_overflow_batch(0, 17);
+    let dense_overflow = make_overflow_batch(0, 2);
+    let overflow_at_end = make_overflow_batch(0, BATCH_SIZE);
+    let legacy = build_new_expr(WideDecimalOp::Add, 10, 2, EvalMode::Legacy);
+    let ansi = build_new_expr(WideDecimalOp::Add, 10, 2, EvalMode::Ansi);
+
+    c.bench_function("wide_decimal: no overflow", |b| {
+        b.iter(|| black_box(legacy.evaluate(black_box(&no_overflow)).unwrap()))
+    });
+    c.bench_function("wide_decimal: no overflow, nulls", |b| {
+        b.iter(|| black_box(legacy.evaluate(black_box(&no_overflow_nulls)).unwrap()))
+    });
+    c.bench_function("wide_decimal: sparse overflow", |b| {
+        b.iter(|| black_box(legacy.evaluate(black_box(&sparse_overflow)).unwrap()))
+    });
+    c.bench_function("wide_decimal: dense overflow", |b| {
+        b.iter(|| black_box(legacy.evaluate(black_box(&dense_overflow)).unwrap()))
+    });
+    c.bench_function("wide_decimal: overflow at end of batch", |b| {
+        b.iter(|| black_box(legacy.evaluate(black_box(&overflow_at_end)).unwrap()))
+    });
+    c.bench_function("wide_decimal: ansi no overflow", |b| {
+        b.iter(|| black_box(ansi.evaluate(black_box(&no_overflow)).unwrap()))
+    });
 }
 
 criterion_group!(benches, criterion_benchmark);

native/spark-expr/src/math_funcs/internal/decimal_rescale_check.rs

@@ -27,6 +27,7 @@ use arrow::record_batch::RecordBatch;
 use datafusion::common::{DataFusionError, ScalarValue};
 use datafusion::logical_expr::ColumnarValue;
 use datafusion::physical_expr::PhysicalExpr;
+use std::cell::Cell;
 use std::hash::Hash;
 use std::{
     fmt::{Display, Formatter},
@@ -105,15 +106,16 @@ fn precision_bound(precision: u8) -> i128 {
 }
 
 /// Rescale a single i128 value by the given delta (output_scale - input_scale)
-/// and check precision bounds. Returns `Ok(value)` or `Ok(i128::MAX)` as sentinel
-/// for overflow in legacy mode, or `Err` in ANSI mode.
+/// and check precision bounds. In legacy mode, records overflow and returns
+/// `Ok(i128::MAX)` as a sentinel; in ANSI mode, returns `Err`.
 #[inline]
 fn rescale_and_check(
     value: i128,
     delta: i8,
     scale_factor: i128,
     bound: i128,
     fail_on_error: bool,
+    overflowed: &Cell<bool>,
 ) -> Result<i128, ArrowError> {
     let rescaled = if delta > 0 {
         // Scale up: multiply. Check for overflow.
@@ -125,6 +127,7 @@ fn rescale_and_check(
                         "Decimal overflow during rescale".to_string(),
                     ));
                 }
+                overflowed.set(true);
                 return Ok(i128::MAX); // sentinel
             }
         }
@@ -146,6 +149,7 @@ fn rescale_and_check(
                 "Decimal overflow: value does not fit in precision".to_string(),
             ));
         }
+        overflowed.set(true);
         Ok(i128::MAX) // sentinel for null_if_overflow_precision
     } else {
         Ok(rescaled)
@@ -185,6 +189,7 @@ impl PhysicalExpr for DecimalRescaleCheckOverflow {
         let fail_on_error = self.fail_on_error;
         let p_out = self.output_precision;
         let s_out = self.output_scale;
+        let overflowed = Cell::new(false);
 
         match arg {
             ColumnarValue::Array(array)
@@ -194,16 +199,17 @@ impl PhysicalExpr for DecimalRescaleCheckOverflow {
 
                 let result: Decimal128Array =
                     arrow::compute::kernels::arity::try_unary(decimal_array, |value| {
-                        rescale_and_check(value, delta, scale_factor, bound, fail_on_error)
+                        rescale_and_check(
+                            value,
+                            delta,
+                            scale_factor,
+                            bound,
+                            fail_on_error,
+                            &overflowed,
+                        )
                     })?;
 
-                let result = if !fail_on_error && result.values().contains(&i128::MAX) {
-                    // The rescale pass writes i128::MAX as an overflow sentinel for values that
-                    // do not fit the output precision. Only when a sentinel is present do we need
-                    // the extra null-masking pass (which allocates a new array); `contains`
-                    // short-circuits at the first sentinel, so the common no-overflow case skips
-                    // that allocation entirely. ANSI mode raises on overflow and never produces a
-                    // sentinel, so it also skips this pass.
+                let result = if overflowed.get() {
                     result.null_if_overflow_precision(p_out)
                 } else {
                     result
@@ -218,8 +224,15 @@ impl PhysicalExpr for DecimalRescaleCheckOverflow {
             ColumnarValue::Scalar(ScalarValue::Decimal128(v, _precision, _scale)) => {
                 let new_v = match v {
                     Some(val) => {
-                        let r = rescale_and_check(val, delta, scale_factor, bound, fail_on_error)
-                            .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?;
+                        let r = rescale_and_check(
+                            val,
+                            delta,
+                            scale_factor,
+                            bound,
+                            fail_on_error,
+                            &overflowed,
+                        )
+                        .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?;
                         if r == i128::MAX {
                             None
                         } else {
@@ -364,8 +377,7 @@ mod tests {
 
     #[test]
     fn test_all_values_overflow_legacy() {
-        // Every value overflows, so the sentinel sits at index 0: `contains` finds it immediately
-        // and the masking pass nulls the whole array.
+        // Every value overflows, so the masking pass nulls the whole array.
         let batch = make_batch(vec![Some(10_000), Some(20_000), Some(30_000)], 10, 2);
         let result = eval_expr(&batch, 2, 4, 2, false).unwrap();
         let arr = result.as_primitive::<Decimal128Type>();

native/spark-expr/src/math_funcs/wide_decimal_binary_expr.rs

@@ -29,6 +29,7 @@ use arrow::record_batch::RecordBatch;
 use datafusion::common::Result;
 use datafusion::logical_expr::ColumnarValue;
 use datafusion::physical_expr::PhysicalExpr;
+use std::cell::Cell;
 use std::fmt::{Display, Formatter};
 use std::hash::Hash;
 use std::sync::Arc;
@@ -214,6 +215,7 @@ impl PhysicalExpr for WideDecimalBinaryExpr {
 
         let bound = max_for_precision(p_out);
         let neg_bound = i256::ZERO.wrapping_sub(bound);
+        let overflowed = Cell::new(false);
 
         let result: Decimal128Array = match op {
             WideDecimalOp::Add | WideDecimalOp::Subtract => {
@@ -249,7 +251,7 @@ impl PhysicalExpr for WideDecimalBinaryExpr {
                     } else {
                         raw
                     };
-                    check_overflow_and_convert(result, bound, neg_bound, eval_mode)
+                    check_overflow_and_convert(result, bound, neg_bound, eval_mode, &overflowed)
                 })?
             }
             WideDecimalOp::Multiply => {
@@ -276,12 +278,12 @@ impl PhysicalExpr for WideDecimalBinaryExpr {
                     } else {
                         raw
                     };
-                    check_overflow_and_convert(result, bound, neg_bound, eval_mode)
+                    check_overflow_and_convert(result, bound, neg_bound, eval_mode, &overflowed)
                 })?
             }
         };
 
-        let result = if eval_mode != EvalMode::Ansi {
+        let result = if overflowed.get() {
             result.null_if_overflow_precision(p_out)
         } else {
             result
@@ -329,20 +331,22 @@ impl PhysicalExpr for WideDecimalBinaryExpr {
 }
 
 /// Check if the i256 result fits in the output precision. In Ansi mode, return an error
-/// on overflow. In Legacy/Try mode, return i128::MAX as a sentinel value that will be
-/// nullified by `null_if_overflow_precision`.
+/// on overflow. In Legacy/Try mode, record the overflow and return i128::MAX as a sentinel
+/// value that will be nullified by `null_if_overflow_precision`.
 #[inline]
 fn check_overflow_and_convert(
     result: i256,
     bound: i256,
     neg_bound: i256,
     eval_mode: EvalMode,
+    overflowed: &Cell<bool>,
 ) -> Result<i128, ArrowError> {
     if result > bound || result < neg_bound {
         if eval_mode == EvalMode::Ansi {
             return Err(ArrowError::ComputeError("Arithmetic overflow".to_string()));
         }
         // Sentinel value — will be nullified by null_if_overflow_precision
+        overflowed.set(true);
         Ok(i128::MAX)
     } else {
         Ok(result.to_i128().unwrap())
@@ -506,6 +510,23 @@ mod tests {
         assert!(arr.is_null(0));
     }
 
+    #[test]
+    fn test_overflow_with_nulls_legacy_mode() {
+        let batch = make_batch(
+            vec![Some(4), Some(5), None],
+            38,
+            0,
+            vec![Some(5), Some(5), Some(1)],
+            38,
+            0,
+        );
+        let result = eval_expr(&batch, WideDecimalOp::Add, 1, 0, EvalMode::Legacy).unwrap();
+        let arr = result.as_primitive::<Decimal128Type>();
+        assert_eq!(arr.value(0), 9);
+        assert!(arr.is_null(1));
+        assert!(arr.is_null(2));
+    }
+
     #[test]
     fn test_overflow_ansi_mode_returns_error() {
         let batch = make_batch(vec![Some(5)], 38, 0, vec![Some(5)], 38, 0);
@@ -620,6 +641,28 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_scalar_scalar_overflow_returns_null_scalar() {
+        use datafusion::common::ScalarValue;
+        use datafusion::physical_expr::expressions::Literal;
+
+        let value = ScalarValue::Decimal128(Some(5), 38, 0);
+        let expr = WideDecimalBinaryExpr::new(
+            Arc::new(Literal::new(value.clone())),
+            Arc::new(Literal::new(value)),
+            WideDecimalOp::Multiply,
+            1,
+            0,
+            EvalMode::Legacy,
+        );
+        let batch = RecordBatch::new_empty(Arc::new(Schema::empty()));
+
+        assert!(matches!(
+            expr.evaluate(&batch).unwrap(),
+            ColumnarValue::Scalar(ScalarValue::Decimal128(None, 1, 0))
+        ));
+    }
+
     /// Companion test: when at least one input is an Array, the result must remain an Array.
     /// Guards against over-eager scalar-unwrapping in the fix.
     #[test]