diff --git a/native/spark-expr/benches/wide_decimal.rs b/native/spark-expr/benches/wide_decimal.rs index ec932ae68f..b38ddbfaa0 100644 --- a/native/spark-expr/benches/wide_decimal.rs +++ b/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::() + .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 { +fn build_new_expr( + op: WideDecimalOp, + p_out: u8, + s_out: i8, + eval_mode: EvalMode, +) -> Arc { let left_col: Arc = Arc::new(Column::new("left", 0)); let right_col: Arc = 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,7 +141,7 @@ 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); } @@ -124,7 +149,7 @@ fn criterion_benchmark(c: &mut Criterion) { { 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); diff --git a/native/spark-expr/src/math_funcs/internal/decimal_rescale_check.rs b/native/spark-expr/src/math_funcs/internal/decimal_rescale_check.rs index fea2399202..7787fa17bd 100644 --- a/native/spark-expr/src/math_funcs/internal/decimal_rescale_check.rs +++ b/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,8 +106,8 @@ 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, @@ -114,6 +115,7 @@ fn rescale_and_check( scale_factor: i128, bound: i128, fail_on_error: bool, + overflowed: &Cell, ) -> Result { 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::(); diff --git a/native/spark-expr/src/math_funcs/wide_decimal_binary_expr.rs b/native/spark-expr/src/math_funcs/wide_decimal_binary_expr.rs index ca4869357e..8d40e8140f 100644 --- a/native/spark-expr/src/math_funcs/wide_decimal_binary_expr.rs +++ b/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, ) -> Result { 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::(); + 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]