perf: optimize CheckOverflow with a shared no-overflow fast path (ANSI and non-ANSI)

docs/source/contributor-guide/expression-audits/math_funcs.md

@@ -41,6 +41,12 @@
 
 - Spark 3.4.3, 3.5.8, 4.0.1, 4.1.1 (audited 2026-05-27): `Divide(left, right, evalMode)`. Non-ANSI mode wraps the divisor in `If(EqualTo(right, 0), null, right)` so DataFusion never throws. Decimal output is wrapped in `CheckOverflow(failOnError = ANSI)`; ANSI surfaces `NUMERIC_VALUE_OUT_OF_RANGE`, non-ANSI returns NULL.
 
+## CheckOverflow (internal)
+
+Internal decimal wrapper emitted around every decimal `+ - * /`, `sum`, and `avg` result to null out (non-ANSI) or raise on (ANSI) values that exceed the declared precision. Native impl: `math_funcs/internal/checkoverflow.rs`.
+
+- Performance (tuned 2026-07-15, PR #4937): both the ANSI and non-ANSI paths share a no-overflow fast path built on `is_valid_decimal_precision`, a small inlined bounds check scanned with `all` (short-circuits at the first overflow). When nothing overflows (the common shape) the input buffers are reused via `to_data()` (cheap Arc metadata clone) instead of allocating through `null_if_overflow_precision` (non-ANSI) or running the heavier per-value `validate_decimal_precision` (ANSI). The ANSI path only falls back to `validate_decimal_precision` when an overflow is present, to build the precise Spark error. ~10% faster on the no-overflow shape, ~17% with nulls, and ~69% faster for ANSI no-overflow (down to parity with non-ANSI); overflow shapes unchanged. Benchmark: `benches/check_overflow.rs`.
+
 ## abs
 
 - Spark 3.4.3, 3.5.8, 4.0.1, 4.1.1 (audited 2026-05-27): `Abs(child, failOnError)` over `NumericType` plus the two interval types. `failOnError` (ANSI) is propagated to the native `abs` UDF, which throws `ARITHMETIC_OVERFLOW` on `Int.MinValue` / `Long.MinValue` / Decimal MIN. `DayTimeIntervalType` and `YearMonthIntervalType` fall back to Spark. Spark 4.0 / 4.1 do the `NullIntolerant` -> `nullIntolerant: Boolean` refactor; behaviour unchanged.

native/spark-expr/Cargo.toml

@@ -139,3 +139,7 @@ harness = false
 [[bench]]
 name = "to_json"
 harness = false
+
+[[bench]]
+name = "check_overflow"
+harness = false

native/spark-expr/benches/check_overflow.rs

@@ -0,0 +1,91 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::{Array, Decimal128Array};
+use arrow::datatypes::{DataType, Field, Schema};
+use arrow::record_batch::RecordBatch;
+use criterion::{criterion_group, criterion_main, Criterion};
+use datafusion::physical_expr::PhysicalExpr;
+use datafusion::physical_plan::expressions::Column;
+use datafusion_comet_spark_expr::CheckOverflow;
+use std::hint::black_box;
+use std::sync::Arc;
+
+// Input arithmetic result is a wide Decimal128(38, 2); CheckOverflow narrows it to the
+// declared Decimal128(10, 2), so any value with |v| >= 10^10 overflows.
+const IN_PRECISION: u8 = 38;
+const TARGET_PRECISION: u8 = 10;
+const SCALE: i8 = 2;
+const OVERFLOW_VALUE: i128 = 1_000_000_000_000; // 10^12, does not fit precision 10
+
+/// Build a Decimal128(38, 2) column of `size` rows.
+/// - `null_every`: null every Nth row (0 = no nulls).
+/// - `overflow_every`: make every Nth value overflow the target precision (0 = none).
+fn build(size: usize, null_every: usize, overflow_every: usize) -> RecordBatch {
+    let values: Decimal128Array = (0..size)
+        .map(|i| {
+            if null_every != 0 && i % null_every == 0 {
+                None
+            } else if overflow_every != 0 && i % overflow_every == 0 {
+                Some(OVERFLOW_VALUE)
+            } else {
+                Some((i as i128 % 100_000) * 100)
+            }
+        })
+        .collect::<Decimal128Array>()
+        .with_precision_and_scale(IN_PRECISION, SCALE)
+        .unwrap();
+
+    let schema = Schema::new(vec![Field::new("d", values.data_type().clone(), true)]);
+    RecordBatch::try_new(Arc::new(schema), vec![Arc::new(values)]).unwrap()
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let size = 8192;
+    let target = DataType::Decimal128(TARGET_PRECISION, SCALE);
+    let col = || Arc::new(Column::new("d", 0)) as Arc<dyn PhysicalExpr>;
+
+    // Non-ANSI (fail_on_error = false): overflowing values become null.
+    let legacy = CheckOverflow::new(col(), target.clone(), false, None, None);
+    // ANSI (fail_on_error = true): overflow raises; only benched on non-overflowing data.
+    let ansi = CheckOverflow::new(col(), target.clone(), true, None, None);
+
+    // The common TPC-DS shape: decimal arithmetic result that fits the declared precision.
+    let no_overflow = build(size, 0, 0);
+    let no_overflow_nulls = build(size, 17, 0);
+    let sparse_overflow = build(size, 0, 17);
+    let dense_overflow = build(size, 0, 2);
+
+    c.bench_function("check_overflow: no overflow", |b| {
+        b.iter(|| black_box(legacy.evaluate(black_box(&no_overflow)).unwrap()))
+    });
+    c.bench_function("check_overflow: no overflow, with nulls", |b| {
+        b.iter(|| black_box(legacy.evaluate(black_box(&no_overflow_nulls)).unwrap()))
+    });
+    c.bench_function("check_overflow: sparse overflow", |b| {
+        b.iter(|| black_box(legacy.evaluate(black_box(&sparse_overflow)).unwrap()))
+    });
+    c.bench_function("check_overflow: dense overflow", |b| {
+        b.iter(|| black_box(legacy.evaluate(black_box(&dense_overflow)).unwrap()))
+    });
+    c.bench_function("check_overflow: ansi no overflow", |b| {
+        b.iter(|| black_box(ansi.evaluate(black_box(&no_overflow)).unwrap()))
+    });
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

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

@@ -119,49 +119,47 @@ impl PhysicalExpr for CheckOverflow {
 
                 let decimal_array = as_primitive_array::<Decimal128Type>(&array);
 
-                let casted_array = if self.fail_on_error {
-                    // Returning error if overflow - convert decimal overflow to SparkError
-                    decimal_array
-                        .validate_decimal_precision(*precision)
-                        .map_err(|e| {
-                            if matches!(e, arrow::error::ArrowError::InvalidArgumentError(_))
-                                && e.to_string().contains("too large to store in a Decimal128") {
-                                // Find the first overflowing value
-                                let overflow_value = decimal_array
-                                    .iter()
-                                    .find(|v| {
-                                        if let Some(val) = v {
-                                            arrow::array::types::Decimal128Type::validate_decimal_precision(
-                                                *val, *precision, *scale
-                                            ).is_err()
-                                        } else {
-                                            false
-                                        }
-                                    })
-                                    .and_then(|v| v)
-                                    .unwrap_or(0);
-
-                                let spark_error = crate::error::decimal_overflow_error(overflow_value, *precision, *scale);
-
-                                // Wrap with query_context if present
-                                if let Some(ctx) = &self.query_context {
-                                    DataFusionError::External(Box::new(
-                                        crate::SparkErrorWithContext::with_context(spark_error, Arc::clone(ctx))
-                                    ))
-                                } else {
-                                    DataFusionError::External(Box::new(spark_error))
-                                }
-                            } else {
-                                DataFusionError::ArrowError(Box::new(e), None)
-                            }
-                        })?;
-                    decimal_array
+                // Fast path shared by both ANSI and non-ANSI: `is_valid_decimal_precision` is a
+                // small, inlined bounds check and `all` short-circuits at the first overflow. When
+                // nothing overflows (the common shape for decimal arithmetic in TPC-DS) we reuse the
+                // input buffers via `to_data()`, which only clones cheap Arc metadata. This avoids
+                // the heavier per-value `validate_decimal_precision` scan (ANSI) or the allocating
+                // `null_if_overflow_precision` (non-ANSI) below.
+                let no_overflow = decimal_array
+                    .iter()
+                    .flatten()
+                    .all(|v| Decimal128Type::is_valid_decimal_precision(v, *precision));
+
+                let casted_array = if no_overflow {
+                    Decimal128Array::from(decimal_array.to_data())
+                } else if self.fail_on_error {
+                    // ANSI mode with a genuine overflow. The fast-path scan already proved an
+                    // overflow exists, so locate the first offending value and raise the precise
+                    // Spark error directly. `is_valid_decimal_precision` checks both the upper and
+                    // lower precision bounds, so this catches negative (underflow) overflow as well
+                    // as positive. This branch only runs on the error path, which aborts the query.
+                    let overflow_value = decimal_array
+                        .iter()
+                        .flatten()
+                        .find(|v| !Decimal128Type::is_valid_decimal_precision(*v, *precision))
+                        .unwrap_or(0);
+                    let spark_error =
+                        crate::error::decimal_overflow_error(overflow_value, *precision, *scale);
+                    return Err(match &self.query_context {
+                        Some(ctx) => DataFusionError::External(Box::new(
+                            crate::SparkErrorWithContext::with_context(
+                                spark_error,
+                                Arc::clone(ctx),
+                            ),
+                        )),
+                        None => DataFusionError::External(Box::new(spark_error)),
+                    });
                 } else {
-                    // Overflowing gets null value
-                    &decimal_array.null_if_overflow_precision(*precision)
+                    // Non-ANSI: overflowing values become null.
+                    decimal_array.null_if_overflow_precision(*precision)
                 };
 
-                let new_array = Decimal128Array::from(casted_array.into_data())
+                let new_array = casted_array
                     .with_precision_and_scale(*precision, *scale)
                     .map(|a| Arc::new(a) as ArrayRef)
                     .map_err(|e| {
@@ -381,4 +379,134 @@ mod tests {
             other => panic!("unexpected: {other:?}"),
         }
     }
+
+    // --- array path ---
+
+    fn array_batch(values: Vec<Option<i128>>, in_precision: u8, scale: i8) -> RecordBatch {
+        let arr = values
+            .into_iter()
+            .collect::<Decimal128Array>()
+            .with_precision_and_scale(in_precision, scale)
+            .unwrap();
+        let schema = Schema::new(vec![Field::new("d", arr.data_type().clone(), true)]);
+        RecordBatch::try_new(Arc::new(schema), vec![Arc::new(arr)]).unwrap()
+    }
+
+    fn array_check_overflow(target_precision: u8, scale: i8, fail_on_error: bool) -> CheckOverflow {
+        CheckOverflow::new(
+            Arc::new(datafusion::physical_plan::expressions::Column::new("d", 0)),
+            DataType::Decimal128(target_precision, scale),
+            fail_on_error,
+            None,
+            None,
+        )
+    }
+
+    fn eval_array(expr: &CheckOverflow, batch: &RecordBatch) -> Decimal128Array {
+        match expr.evaluate(batch).unwrap() {
+            ColumnarValue::Array(a) => a
+                .as_any()
+                .downcast_ref::<Decimal128Array>()
+                .unwrap()
+                .clone(),
+            other => panic!("expected array, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn test_array_no_overflow_legacy_preserves_values_and_type() {
+        // No value overflows precision 3, so the fast path reuses the input; values, nulls,
+        // and the target precision/scale must all be preserved.
+        let batch = array_batch(vec![Some(999), Some(12), None, Some(5)], 38, 0);
+        let out = eval_array(&array_check_overflow(3, 0, false), &batch);
+        assert_eq!(out.data_type(), &DataType::Decimal128(3, 0));
+        assert_eq!(
+            out.iter().collect::<Vec<_>>(),
+            vec![Some(999), Some(12), None, Some(5)]
+        );
+    }
+
+    #[test]
+    fn test_array_overflow_nulled_legacy() {
+        // 1000 does not fit precision 3 → nulled; other values and existing nulls kept.
+        let batch = array_batch(vec![Some(999), Some(1000), None, Some(5)], 38, 0);
+        let out = eval_array(&array_check_overflow(3, 0, false), &batch);
+        assert_eq!(out.data_type(), &DataType::Decimal128(3, 0));
+        assert_eq!(
+            out.iter().collect::<Vec<_>>(),
+            vec![Some(999), None, None, Some(5)]
+        );
+    }
+
+    #[test]
+    fn test_array_no_overflow_ansi_ok() {
+        let batch = array_batch(vec![Some(999), None, Some(5)], 38, 0);
+        let out = eval_array(&array_check_overflow(3, 0, true), &batch);
+        assert_eq!(
+            out.iter().collect::<Vec<_>>(),
+            vec![Some(999), None, Some(5)]
+        );
+    }
+
+    #[test]
+    fn test_array_overflow_ansi_errors() {
+        let batch = array_batch(vec![Some(999), Some(1000)], 38, 0);
+        let result = array_check_overflow(3, 0, true).evaluate(&batch);
+        assert!(result.is_err(), "expected error on overflow in ANSI mode");
+    }
+
+    #[test]
+    fn test_array_negative_overflow_nulled_legacy() {
+        // -1000 is below the precision-3 lower bound (-999) → nulled; other values kept.
+        // Guards the negative (underflow) bound, which is a distinct branch from positive overflow.
+        let batch = array_batch(vec![Some(-1000), Some(5)], 38, 0);
+        let out = eval_array(&array_check_overflow(3, 0, false), &batch);
+        assert_eq!(out.iter().collect::<Vec<_>>(), vec![None, Some(5)]);
+    }
+
+    #[test]
+    fn test_array_negative_overflow_ansi_errors() {
+        // ANSI mode must raise on a negative overflow, not only a positive one. The previous
+        // implementation string-matched "too large" and silently missed the "too small" branch.
+        let batch = array_batch(vec![Some(5), Some(-1000)], 38, 0);
+        let result = array_check_overflow(3, 0, true).evaluate(&batch);
+        assert!(
+            result.is_err(),
+            "expected error on negative overflow in ANSI mode"
+        );
+    }
+
+    #[test]
+    fn test_array_all_null_reuses_input_and_preserves_mask() {
+        // The fast-path scan is `flatten().all(...)`, which returns true on an all-null batch
+        // (empty after flatten). The input must be reused unchanged with its null mask intact.
+        let batch = array_batch(vec![None, None, None], 38, 0);
+        let out = eval_array(&array_check_overflow(3, 0, false), &batch);
+        assert_eq!(out.data_type(), &DataType::Decimal128(3, 0));
+        assert_eq!(out.iter().collect::<Vec<_>>(), vec![None, None, None]);
+    }
+
+    #[test]
+    fn test_array_all_overflow_nulled_legacy() {
+        // Every value overflows precision 3 → every slot nulled in legacy mode.
+        let batch = array_batch(vec![Some(1000), Some(5000), Some(-2000)], 38, 0);
+        let out = eval_array(&array_check_overflow(3, 0, false), &batch);
+        assert_eq!(out.iter().collect::<Vec<_>>(), vec![None, None, None]);
+    }
+
+    #[test]
+    fn test_array_all_overflow_ansi_errors() {
+        let batch = array_batch(vec![Some(1000), Some(5000)], 38, 0);
+        let result = array_check_overflow(3, 0, true).evaluate(&batch);
+        assert!(result.is_err(), "expected error on overflow in ANSI mode");
+    }
+
+    #[test]
+    fn test_array_boundary_precision_max_passes_and_over_by_one_overflows() {
+        // 999 is exactly the max for precision 3 and must pass; 9999 is over and must be nulled.
+        // Pins the off-by-one on the precision bound.
+        let batch = array_batch(vec![Some(999), Some(9999)], 38, 0);
+        let out = eval_array(&array_check_overflow(3, 0, false), &batch);
+        assert_eq!(out.iter().collect::<Vec<_>>(), vec![Some(999), None]);
+    }
 }