diff --git a/native/core/src/execution/expressions/arithmetic.rs b/native/core/src/execution/expressions/arithmetic.rs index 4d2b077891..ff275e4978 100644 --- a/native/core/src/execution/expressions/arithmetic.rs +++ b/native/core/src/execution/expressions/arithmetic.rs @@ -45,6 +45,10 @@ impl CheckedBinaryExpr { query_context, } } + + pub(crate) fn child(&self) -> &Arc { + &self.child + } } impl Display for CheckedBinaryExpr { diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 39e7ccc2e5..0fd153ce72 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -101,6 +101,7 @@ use datafusion::physical_expr::expressions::{Literal, StatsType}; use datafusion::physical_expr::window::WindowExpr; use datafusion::physical_expr::LexOrdering; +use crate::execution::expressions::arithmetic::CheckedBinaryExpr; use crate::parquet::parquet_exec::init_datasource_exec; use arrow::array::{ new_empty_array, Array, ArrayRef, BinaryBuilder, BooleanArray, Date32Array, Decimal128Array, @@ -517,18 +518,41 @@ impl PhysicalPlanner { self.create_expr(expr.child.as_ref().unwrap(), Arc::clone(&input_schema))?; let data_type = to_arrow_datatype(expr.datatype.as_ref().unwrap()); let fail_on_error = expr.fail_on_error; + let query_context = spark_expr.expr_id.and_then(|expr_id| { + let registry = &self.query_context_registry; + registry.get(expr_id) + }); // WideDecimalBinaryExpr already handles overflow — skip redundant check - // but only if its output type matches CheckOverflow's declared type - if child.downcast_ref::().is_some() { + // but only if its output type matches CheckOverflow's declared type. A binary + // expression with query context is already wrapped in CheckedBinaryExpr, so + // inspect through that single layer too. + let is_wide_decimal = child.downcast_ref::().is_some() + || child + .downcast_ref::() + .is_some_and(|checked| { + checked + .child() + .downcast_ref::() + .is_some() + }); + if is_wide_decimal { let child_type = child.data_type(&input_schema)?; if child_type == data_type { - return Ok(child); + return if query_context.is_some() + && child.downcast_ref::().is_none() + { + Ok(Arc::new(CheckedBinaryExpr::new(child, query_context))) + } else { + Ok(child) + }; } } // Fuse Cast(Decimal128→Decimal128) + CheckOverflow into single rescale+check - // Only fuse when the Cast target type matches the CheckOverflow output type + // Only fuse when the Cast target type matches the CheckOverflow output type. + // Spark 3.4+ does not currently emit this shape, but keep its errors typed and + // contextualized in case a future serializer makes the fusion reachable. if let Some(cast) = child.downcast_ref::() { if let ( DataType::Decimal128(p_out, s_out), @@ -537,23 +561,33 @@ impl PhysicalPlanner { { let cast_target = cast.data_type(&input_schema)?; if cast_target == data_type { - return Ok(Arc::new(DecimalRescaleCheckOverflow::new( - Arc::clone(&cast.child), - s_in, - *p_out, - *s_out, - fail_on_error, - ))); + let fused: Arc = + Arc::new(DecimalRescaleCheckOverflow::new( + Arc::clone(&cast.child), + s_in, + *p_out, + *s_out, + fail_on_error, + )); + return if query_context.is_some() { + Ok(Arc::new(CheckedBinaryExpr::new(fused, query_context))) + } else { + Ok(fused) + }; } } } - // Look up query context from registry if expr_id is present - let query_context = spark_expr.expr_id.and_then(|expr_id| { - let registry = &self.query_context_registry; - registry.get(expr_id) - }); - + // Generated child protos may not carry an expression id of their own, so retain + // the outer CheckOverflow context as a fallback for bare child SparkErrors. + let child = if query_context.is_some() + && child.downcast_ref::().is_none() + { + Arc::new(CheckedBinaryExpr::new(child, query_context.clone())) + as Arc + } else { + child + }; Ok(Arc::new(CheckOverflow::new( child, data_type, @@ -933,9 +967,14 @@ impl PhysicalPlanner { DataFusionOperator::Multiply => WideDecimalOp::Multiply, _ => unreachable!(), }; - Ok(Arc::new(WideDecimalBinaryExpr::new( + let expr: Arc = Arc::new(WideDecimalBinaryExpr::new( left, right, wide_op, p_out, s_out, eval_mode, - ))) + )); + if query_context.is_some() { + Ok(Arc::new(CheckedBinaryExpr::new(expr, query_context))) + } else { + Ok(expr) + } } ( DataFusionOperator::Divide, @@ -960,13 +999,18 @@ impl PhysicalPlanner { Some(options.check_divide_overflow), eval_mode, )?; - Ok(Arc::new(ScalarFunctionExpr::new( + let expr: Arc = Arc::new(ScalarFunctionExpr::new( func_name, fun_expr, vec![left, right], Arc::new(Field::new(func_name, data_type, true)), Arc::new(ConfigOptions::default()), - ))) + )); + if query_context.is_some() { + Ok(Arc::new(CheckedBinaryExpr::new(expr, query_context))) + } else { + Ok(expr) + } } // Date +/- Int8/Int16/Int32: DataFusion 52's arrow-arith kernels only // support Date32 +/- Interval types, not raw integers. Use the Spark @@ -1033,9 +1077,11 @@ impl PhysicalPlanner { Arc::new(ConfigOptions::default()), )); - // Wrap with CheckedBinaryExpr to add query_context to errors - use crate::execution::expressions::arithmetic::CheckedBinaryExpr; - Ok(Arc::new(CheckedBinaryExpr::new(scalar_expr, query_context))) + if query_context.is_some() { + Ok(Arc::new(CheckedBinaryExpr::new(scalar_expr, query_context))) + } else { + Ok(scalar_expr) + } } else { Ok(Arc::new(BinaryExpr::new(left, op, right))) } diff --git a/native/spark-expr/src/conversion_funcs/cast.rs b/native/spark-expr/src/conversion_funcs/cast.rs index 37fddb8c11..7f8124c6be 100644 --- a/native/spark-expr/src/conversion_funcs/cast.rs +++ b/native/spark-expr/src/conversion_funcs/cast.rs @@ -45,15 +45,17 @@ use arrow::array::{ new_null_array, BinaryBuilder, DictionaryArray, GenericByteArray, ListArray, MapArray, StringArray, StructArray, }; -use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType, Schema}; +use arrow::datatypes::{ + format_decimal_str, ArrowDictionaryKeyType, ArrowNativeType, DataType, Schema, +}; use arrow::datatypes::{Field, Fields, GenericBinaryType}; use arrow::error::ArrowError; use arrow::{ array::{ - cast::AsArray, types::Int32Type, Array, ArrayRef, Int16Array, Int32Array, Int64Array, - Int8Array, OffsetSizeTrait, PrimitiveArray, + cast::AsArray, types::Decimal128Type, types::Int32Type, Array, ArrayRef, Int16Array, + Int32Array, Int64Array, Int8Array, OffsetSizeTrait, PrimitiveArray, }, - compute::{cast_with_options, take, CastOptions}, + compute::{cast_with_options, rescale_decimal, take, CastOptions}, record_batch::RecordBatch, util::display::FormatOptions, }; @@ -339,6 +341,38 @@ pub(crate) fn cast_array( (Utf8 | LargeUtf8, Decimal256(precision, scale)) => { cast_string_to_decimal(&array, to_type, precision, scale, eval_mode) } + (Decimal128(input_precision, input_scale), Decimal128(output_precision, output_scale)) + if eval_mode == EvalMode::Ansi => + { + cast_with_options(&array, to_type, &native_cast_options).map_err(|error| { + array + .as_primitive::() + .iter() + .flatten() + .find(|value| { + rescale_decimal::( + *value, + *input_precision, + *input_scale, + *output_precision, + *output_scale, + ) + .is_none() + }) + .map_or_else( + || error.into(), + |value| SparkError::NumericValueOutOfRange { + value: format_decimal_str( + &value.to_string(), + *input_precision as usize, + *input_scale, + ), + precision: *output_precision, + scale: *output_scale, + }, + ) + }) + } (Int64, Int32) | (Int64, Int16) | (Int64, Int8) @@ -752,8 +786,12 @@ impl PhysicalExpr for Cast { } fn evaluate(&self, batch: &RecordBatch) -> DataFusionResult { - let arg = self.child.evaluate(batch)?; - let result = spark_cast(arg, &self.data_type, &self.cast_options); + // `CometIntegralDivide` builds its inner `CheckOverflow` without an expression id, so a + // bare child `SparkError` deliberately inherits the outer Cast's query context here. + let result = self + .child + .evaluate(batch) + .and_then(|arg| spark_cast(arg, &self.data_type, &self.cast_options)); // If there's an error and we have query_context, wrap it match result { @@ -908,11 +946,62 @@ fn cast_binary_to_string( #[cfg(test)] mod tests { use super::*; - use arrow::array::{BinaryArray, ListArray, NullArray, StringArray}; + use arrow::array::{BinaryArray, Decimal128Array, ListArray, NullArray, StringArray}; use arrow::buffer::OffsetBuffer; use arrow::datatypes::TimestampMicrosecondType; use arrow::datatypes::{Field, Fields}; + #[test] + fn test_cast_decimal_to_decimal_ansi_overflow_returns_spark_error() { + let cases = [ + ( + vec![None, Some(1), Some(-123_456_789)], + DataType::Decimal128(10, 4), + DataType::Decimal128(6, 2), + "-12345.6789", + 6, + 2, + ), + ( + vec![None, Some(1), Some(-999)], + DataType::Decimal128(3, 0), + DataType::Decimal128(3, 2), + "-999", + 3, + 2, + ), + ]; + + for (values, input_type, output_type, expected_value, expected_precision, expected_scale) in + cases + { + let input: ArrayRef = + Arc::new(Decimal128Array::from(values).with_data_type(input_type)); + let error = cast_array( + input, + &output_type, + &SparkCastOptions::new_without_timezone(EvalMode::Ansi, false), + ) + .unwrap_err(); + + match error { + DataFusionError::External(error) => match error.downcast_ref::() { + Some(SparkError::NumericValueOutOfRange { + value, + precision, + scale, + }) => { + assert_eq!(value, expected_value); + assert_eq!(*precision, expected_precision); + assert_eq!(*scale, expected_scale); + } + other => panic!("expected NumericValueOutOfRange, got {other:?}"), + }, + other => panic!("expected external SparkError, got {other:?}"), + } + } + } + #[test] fn test_cast_binary_to_string_replaces_invalid_utf8_jvm_compatibly() { // Invalid bytes are replaced with U+FFFD instead of reinterpreted as an invalid `str`, diff --git a/native/spark-expr/src/conversion_funcs/numeric.rs b/native/spark-expr/src/conversion_funcs/numeric.rs index cb141880d0..667e1346bc 100644 --- a/native/spark-expr/src/conversion_funcs/numeric.rs +++ b/native/spark-expr/src/conversion_funcs/numeric.rs @@ -24,8 +24,8 @@ use arrow::array::{ OffsetSizeTrait, PrimitiveArray, StringBuilder, TimestampMicrosecondBuilder, }; use arrow::datatypes::{ - i256, is_validate_decimal_precision, ArrowPrimitiveType, DataType, Decimal128Type, Float32Type, - Float64Type, Int16Type, Int32Type, Int64Type, Int8Type, + format_decimal_str, i256, is_validate_decimal_precision, ArrowPrimitiveType, DataType, + Decimal128Type, Float32Type, Float64Type, Int16Type, Int32Type, Int64Type, Int8Type, }; use num::{cast::AsPrimitive, ToPrimitive, Zero}; use std::sync::Arc; @@ -559,30 +559,6 @@ macro_rules! cast_decimal_to_int32_up { }}; } -// copied from arrow::dataTypes::Decimal128Type since Decimal128Type::format_decimal can't be called directly -pub(crate) fn format_decimal_str(value_str: &str, precision: usize, scale: i8) -> String { - let (sign, rest) = match value_str.strip_prefix('-') { - Some(stripped) => ("-", stripped), - None => ("", value_str), - }; - let bound = precision.min(rest.len()) + sign.len(); - let value_str = &value_str[0..bound]; - - if scale == 0 { - value_str.to_string() - } else if scale < 0 { - let padding = value_str.len() + scale.unsigned_abs() as usize; - format!("{value_str:0 scale as usize { - // Decimal separator is in the middle of the string - let (whole, decimal) = value_str.split_at(value_str.len() - scale as usize); - format!("{whole}.{decimal}") - } else { - // String has to be padded - format!("{}0.{:0>width$}", sign, rest, width = scale as usize) - } -} - /// Casts a Decimal128 array to string using Java's BigDecimal.toString() semantics, /// which is Spark's LEGACY eval mode behavior. Plain notation when scale >= 0 and /// adjusted_exponent >= -6, otherwise scientific notation (e.g. "0E-18" for zero diff --git a/native/spark-expr/src/error.rs b/native/spark-expr/src/error.rs index bb87915c7b..f77744be17 100644 --- a/native/spark-expr/src/error.rs +++ b/native/spark-expr/src/error.rs @@ -19,3 +19,15 @@ pub use datafusion_comet_common::{ decimal_overflow_error, SparkError, SparkErrorWithContext, SparkResult, }; + +use arrow::error::ArrowError; +use datafusion::common::DataFusionError; + +/// Arrow's `try_*` kernels require closure errors to be `ArrowError`, so a `SparkError` +/// travels through `ExternalError`. Unwrap it again so JNI sees the direct `SparkError`. +pub(crate) fn unwrap_arrow_external_error(error: ArrowError) -> DataFusionError { + match error { + ArrowError::ExternalError(error) => DataFusionError::External(error), + error => error.into(), + } +} diff --git a/native/spark-expr/src/math_funcs/div.rs b/native/spark-expr/src/math_funcs/div.rs index 9b3720a29f..685f5368dd 100644 --- a/native/spark-expr/src/math_funcs/div.rs +++ b/native/spark-expr/src/math_funcs/div.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use crate::error::unwrap_arrow_external_error; use crate::math_funcs::utils::get_precision_scale; use crate::{divide_by_zero_error, integral_divide_overflow_error, EvalMode}; use arrow::array::{Array, Decimal128Array}; @@ -61,9 +62,9 @@ fn quotient_to_i128( ) -> Result { let res = res.to_i128().unwrap_or(i128::MAX); if check_divide_overflow && i64::try_from(res).is_err() { - return Err(ArrowError::ComputeError( - integral_divide_overflow_error().to_string(), - )); + return Err(ArrowError::ExternalError(Box::new( + integral_divide_overflow_error(), + ))); } Ok(res) } @@ -100,7 +101,7 @@ fn spark_decimal_div_internal( let l_exp = ((s2 + s3 + 1) as u32).saturating_sub(s1 as u32); let r_exp = (s1 as u32).saturating_sub((s2 + s3 + 1) as u32); - let result: Decimal128Array = if p1 as u32 + l_exp > DECIMAL128_MAX_PRECISION as u32 + let result = if p1 as u32 + l_exp > DECIMAL128_MAX_PRECISION as u32 || p2 as u32 + r_exp > DECIMAL128_MAX_PRECISION as u32 { let ten = BigInt::from(10); @@ -116,7 +117,7 @@ fn spark_decimal_div_internal( // Spark throws DIVIDE_BY_ZERO for both `/` and `div` when ANSI is enabled, so // the `is_integral_div` guard was wrong and has been removed. if eval_mode == EvalMode::Ansi && r.is_zero() { - return Err(ArrowError::ComputeError(divide_by_zero_error().to_string())); + return Err(ArrowError::ExternalError(Box::new(divide_by_zero_error()))); } // Non-ANSI: zero divisors have already been replaced with null by the // `nullIfWhenPrimitive` wrapper applied in the Scala serde layer, so @@ -131,7 +132,7 @@ fn spark_decimal_div_internal( div + &five } / &ten; quotient_to_i128(&res, check_divide_overflow) - })? + }) } else { let l_mul = 10_i128.pow(l_exp); let r_mul = 10_i128.pow(r_exp); @@ -143,7 +144,7 @@ fn spark_decimal_div_internal( // Spark throws DIVIDE_BY_ZERO for both `/` and `div` when ANSI is enabled, so // the `is_integral_div` guard was wrong and has been removed. if eval_mode == EvalMode::Ansi && r == 0 { - return Err(ArrowError::ComputeError(divide_by_zero_error().to_string())); + return Err(ArrowError::ExternalError(Box::new(divide_by_zero_error()))); } // Non-ANSI: zero divisors have already been replaced with null by the // `nullIfWhenPrimitive` wrapper applied in the Scala serde layer, so @@ -158,8 +159,56 @@ fn spark_decimal_div_internal( div + 5 } / 10; quotient_to_i128(&res, check_divide_overflow) - })? + }) }; + let result: Decimal128Array = result.map_err(unwrap_arrow_external_error)?; let result = result.with_data_type(DataType::Decimal128(p3, s3)); Ok(ColumnarValue::Array(Arc::new(result))) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::SparkError; + + fn decimal(value: i128, precision: u8, scale: i8) -> ColumnarValue { + ColumnarValue::Array(Arc::new( + Decimal128Array::from(vec![Some(value)]) + .with_data_type(DataType::Decimal128(precision, scale)), + )) + } + + fn spark_error(result: Result) -> SparkError { + match result.unwrap_err() { + DataFusionError::External(error) => *error + .downcast::() + .expect("expected external SparkError"), + error => panic!("expected external SparkError, got {error:?}"), + } + } + + #[test] + fn test_decimal_divide_by_zero_returns_spark_error() { + // Exercise both the i128 and BigInt kernels. + for (precision, scale) in [(10, 2), (38, 0)] { + let result = spark_decimal_div( + &[decimal(100, precision, scale), decimal(0, precision, scale)], + &DataType::Decimal128(precision, scale), + EvalMode::Ansi, + ); + assert!(matches!(spark_error(result), SparkError::DivideByZero)); + } + } + + #[test] + fn test_integral_divide_overflow_returns_spark_error() { + let result = quotient_to_i128(&(i64::MAX as i128 + 1), true); + match result.unwrap_err() { + ArrowError::ExternalError(error) => assert!(matches!( + error.downcast_ref::(), + Some(SparkError::IntegralDivideOverflow) + )), + error => panic!("expected external SparkError, got {error:?}"), + } + } +} 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..727685f19d 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 @@ -20,8 +20,10 @@ //! Replaces the pattern `CheckOverflow(Cast(expr, Decimal128(p2,s2)), Decimal128(p2,s2))` //! with a single expression that rescales and validates precision in one pass. +use crate::error::unwrap_arrow_external_error; +use crate::SparkError; use arrow::array::{as_primitive_array, Array, ArrayRef, Decimal128Array}; -use arrow::datatypes::{DataType, Decimal128Type, Schema}; +use arrow::datatypes::{format_decimal_str, DataType, Decimal128Type, Schema}; use arrow::error::ArrowError; use arrow::record_batch::RecordBatch; use datafusion::common::{DataFusionError, ScalarValue}; @@ -110,20 +112,33 @@ fn precision_bound(precision: u8) -> i128 { #[inline] fn rescale_and_check( value: i128, - delta: i8, - scale_factor: i128, + input_scale: i8, + scale_factor: Option, bound: i128, + output_precision: u8, + output_scale: i8, fail_on_error: bool, ) -> Result { + let overflow_error = || { + let unscaled = value.to_string(); + // Preserve every digit of the value that is already known to overflow. + let digits = unscaled.trim_start_matches('-').len(); + ArrowError::ExternalError(Box::new(SparkError::NumericValueOutOfRange { + value: format_decimal_str(&unscaled, digits, input_scale), + precision: output_precision, + scale: output_scale, + })) + }; + let delta = output_scale as i16 - input_scale as i16; + let rescaled = if delta > 0 { // Scale up: multiply. Check for overflow. - match value.checked_mul(scale_factor) { + match scale_factor.and_then(|factor| value.checked_mul(factor)) { Some(v) => v, + None if value == 0 => 0, None => { if fail_on_error { - return Err(ArrowError::ComputeError( - "Decimal overflow during rescale".to_string(), - )); + return Err(overflow_error()); } return Ok(i128::MAX); // sentinel } @@ -131,10 +146,14 @@ fn rescale_and_check( } else if delta < 0 { // Scale down with HALF_UP rounding // divisor = 10^(-delta), half = divisor / 2 - let divisor = scale_factor; // already 10^abs(delta) - let half = divisor / 2; - let sign = value.signum(); - (value + sign * half) / divisor + match scale_factor { + Some(divisor) => { + let half = divisor / 2; + let sign = value.signum(); + (value + sign * half) / divisor + } + None => 0, + } } else { value }; @@ -142,9 +161,7 @@ fn rescale_and_check( // Precision check if rescaled.abs() > bound { if fail_on_error { - return Err(ArrowError::ComputeError( - "Decimal overflow: value does not fit in precision".to_string(), - )); + return Err(overflow_error()); } Ok(i128::MAX) // sentinel for null_if_overflow_precision } else { @@ -170,17 +187,9 @@ impl PhysicalExpr for DecimalRescaleCheckOverflow { fn evaluate(&self, batch: &RecordBatch) -> datafusion::common::Result { let arg = self.child.evaluate(batch)?; - let delta = self.output_scale - self.input_scale; + let delta = self.output_scale as i16 - self.input_scale as i16; let abs_delta = delta.unsigned_abs(); - // If abs_delta > 38, the scale factor overflows i128. In that case, - // any non-zero value will overflow the output precision, so we treat - // it as an immediate overflow condition. - if abs_delta > 38 { - return Err(DataFusionError::Execution(format!( - "DecimalRescaleCheckOverflow: scale delta {delta} exceeds maximum supported range" - ))); - } - let scale_factor = 10i128.pow(abs_delta as u32); + let scale_factor = (abs_delta <= 38).then(|| 10i128.pow(abs_delta as u32)); let bound = precision_bound(self.output_precision); let fail_on_error = self.fail_on_error; let p_out = self.output_precision; @@ -194,8 +203,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, + self.input_scale, + scale_factor, + bound, + p_out, + s_out, + fail_on_error, + ) + }) + .map_err(unwrap_arrow_external_error)?; let result = if !fail_on_error && result.values().contains(&i128::MAX) { // The rescale pass writes i128::MAX as an overflow sentinel for values that @@ -218,8 +236,16 @@ 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, + self.input_scale, + scale_factor, + bound, + p_out, + s_out, + fail_on_error, + ) + .map_err(unwrap_arrow_external_error)?; if r == i128::MAX { None } else { @@ -270,6 +296,29 @@ mod tests { use arrow::record_batch::RecordBatch; use datafusion::physical_expr::expressions::Column; + fn assert_numeric_value_out_of_range( + error: DataFusionError, + expected_value: &str, + expected_precision: u8, + expected_scale: i8, + ) { + match error { + DataFusionError::External(error) => match error.downcast_ref::() { + Some(SparkError::NumericValueOutOfRange { + value, + precision, + scale, + }) => { + assert_eq!(value, expected_value); + assert_eq!(*precision, expected_precision); + assert_eq!(*scale, expected_scale); + } + other => panic!("expected NumericValueOutOfRange, got {other:?}"), + }, + other => panic!("expected external SparkError, got {other:?}"), + } + } + fn make_batch(values: Vec>, precision: u8, scale: i8) -> RecordBatch { let arr = Decimal128Array::from(values).with_data_type(DataType::Decimal128(precision, scale)); @@ -344,9 +393,9 @@ mod tests { #[test] fn test_overflow_error_in_ansi_mode() { - let batch = make_batch(vec![Some(10)], 38, 0); - let result = eval_expr(&batch, 0, 3, 2, true); - assert!(result.is_err()); + let batch = make_batch(vec![Some(-1000)], 10, 2); + let error = eval_expr(&batch, 2, 3, 2, true).unwrap_err(); + assert_numeric_value_out_of_range(error, "-10.00", 3, 2); } #[test] @@ -491,26 +540,34 @@ mod tests { #[test] fn test_scalar_overflow_ansi_returns_error() { - // fail_on_error=true must propagate the error, not silently return None let schema = Schema::new(vec![Field::new("col", DataType::Decimal128(38, 0), true)]); let batch = RecordBatch::new_empty(Arc::new(schema)); + let value = 10i128.pow(38) - 1; let expr = DecimalRescaleCheckOverflow::new( - Arc::new(ScalarChild(Some(10), 38, 0)), + Arc::new(ScalarChild(Some(value), 38, 0)), 0, - 3, - 2, - true, // fail_on_error = true + 38, + 1, + true, ); - let result = expr.evaluate(&batch); - assert!(result.is_err()); // must be error, not Ok(None) + let error = expr.evaluate(&batch).unwrap_err(); + assert_numeric_value_out_of_range(error, &value.to_string(), 38, 1); } #[test] - fn test_large_scale_delta_returns_error() { - // delta = output_scale - input_scale = 38 - (-1) = 39 - // 10i128.pow(39) would overflow, so we must reject gracefully - let batch = make_batch(vec![Some(1)], 38, -1); - let result = eval_expr(&batch, -1, 38, 38, false); - assert!(result.is_err()); + fn test_large_scale_delta() { + let scale_up = make_batch(vec![Some(1), Some(0), None], 38, -1); + let result = eval_expr(&scale_up, -1, 38, 38, false).unwrap(); + let result = result.as_primitive::(); + assert!(result.is_null(0)); + assert_eq!(result.value(1), 0); + assert!(result.is_null(2)); + + let error = eval_expr(&scale_up, -1, 38, 38, true).unwrap_err(); + assert_numeric_value_out_of_range(error, "10", 38, 38); + + let scale_down = make_batch(vec![Some(1)], 38, 38); + let result = eval_expr(&scale_down, 38, 38, -1, false).unwrap(); + assert_eq!(result.as_primitive::().value(0), 0); } } 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..e439e9a9f3 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 @@ -20,10 +20,11 @@ //! Instead of building a 4-node expression tree (Cast→BinaryExpr→Cast→Cast), this performs //! i256 intermediate arithmetic in a single expression, producing only one output array. +use crate::error::unwrap_arrow_external_error; use crate::math_funcs::utils::get_precision_scale; -use crate::EvalMode; +use crate::{EvalMode, SparkError}; use arrow::array::{Array, ArrayRef, AsArray, Decimal128Array}; -use arrow::datatypes::{i256, DataType, Decimal128Type, Schema}; +use arrow::datatypes::{format_decimal_str, i256, DataType, Decimal128Type, Schema}; use arrow::error::ArrowError; use arrow::record_batch::RecordBatch; use datafusion::common::Result; @@ -215,7 +216,7 @@ impl PhysicalExpr for WideDecimalBinaryExpr { let bound = max_for_precision(p_out); let neg_bound = i256::ZERO.wrapping_sub(bound); - let result: Decimal128Array = match op { + let result: std::result::Result = match op { WideDecimalOp::Add | WideDecimalOp::Subtract => { let max_scale = std::cmp::max(s1, s2); let l_scale_up = i256_pow10((max_scale - s1) as u32); @@ -249,8 +250,8 @@ impl PhysicalExpr for WideDecimalBinaryExpr { } else { raw }; - check_overflow_and_convert(result, bound, neg_bound, eval_mode) - })? + check_overflow_and_convert(result, bound, neg_bound, p_out, s_out, eval_mode) + }) } WideDecimalOp::Multiply => { let natural_scale = s1 + s2; @@ -276,10 +277,11 @@ impl PhysicalExpr for WideDecimalBinaryExpr { } else { raw }; - check_overflow_and_convert(result, bound, neg_bound, eval_mode) - })? + check_overflow_and_convert(result, bound, neg_bound, p_out, s_out, eval_mode) + }) } }; + let result = result.map_err(unwrap_arrow_external_error)?; let result = if eval_mode != EvalMode::Ansi { result.null_if_overflow_precision(p_out) @@ -336,11 +338,25 @@ fn check_overflow_and_convert( result: i256, bound: i256, neg_bound: i256, + precision: u8, + scale: i8, eval_mode: EvalMode, ) -> Result { if result > bound || result < neg_bound { if eval_mode == EvalMode::Ansi { - return Err(ArrowError::ComputeError("Arithmetic overflow".to_string())); + let unscaled = result.to_string(); + // Arrow's formatter truncates to its precision argument. This value is already + // known to overflow, so pass its actual digit count to preserve every digit. + // Spark reports the pre-toPrecision value instead; see + // https://github.com/apache/datafusion-comet/issues/5211. + let digits = unscaled.trim_start_matches('-').len(); + return Err(ArrowError::ExternalError(Box::new( + SparkError::NumericValueOutOfRange { + value: format_decimal_str(&unscaled, digits, scale), + precision, + scale, + }, + ))); } // Sentinel value — will be nullified by null_if_overflow_precision Ok(i128::MAX) @@ -355,6 +371,7 @@ mod tests { use arrow::array::Decimal128Array; use arrow::datatypes::{Field, Schema}; use arrow::record_batch::RecordBatch; + use datafusion::common::DataFusionError; use datafusion::physical_expr::expressions::Column; fn make_batch( @@ -507,10 +524,42 @@ mod tests { } #[test] - fn test_overflow_ansi_mode_returns_error() { - let batch = make_batch(vec![Some(5)], 38, 0, vec![Some(5)], 38, 0); - let result = eval_expr(&batch, WideDecimalOp::Add, 1, 0, EvalMode::Ansi); - assert!(result.is_err()); + fn test_overflow_ansi_mode_returns_spark_error() { + let cases = [ + ( + make_batch(vec![Some(5)], 38, 1, vec![Some(5)], 38, 1), + WideDecimalOp::Add, + 1, + 1, + "1.0", + ), + ( + make_batch(vec![Some(99)], 2, 1, vec![Some(10)], 2, 1), + WideDecimalOp::Multiply, + 1, + 0, + "10", + ), + ]; + + for (batch, op, precision, scale, expected_value) in cases { + let result = eval_expr(&batch, op, precision, scale, EvalMode::Ansi); + match result { + Err(DataFusionError::External(error)) => match error.downcast_ref::() { + Some(SparkError::NumericValueOutOfRange { + value, + precision: actual_precision, + scale: actual_scale, + }) => { + assert_eq!(value, expected_value); + assert_eq!(*actual_precision, precision); + assert_eq!(*actual_scale, scale); + } + other => panic!("expected NumericValueOutOfRange, got {other:?}"), + }, + other => panic!("expected external SparkError, got {other:?}"), + } + } } #[test] diff --git a/spark/src/test/scala/org/apache/comet/CometCastSuite.scala b/spark/src/test/scala/org/apache/comet/CometCastSuite.scala index fa5178cdfd..e9ef9d397c 100644 --- a/spark/src/test/scala/org/apache/comet/CometCastSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCastSuite.scala @@ -1602,7 +1602,7 @@ class CometCastSuite extends CometTestBase with AdaptiveSparkPlanHelper { spark.sparkContext.parallelize(rowData), StructType(Seq(StructField("a", DataTypes.createDecimalType(10, 4))))) - castTest(df, DecimalType(6, 2)) + castTest(df, DecimalType(6, 2), expectAnsiFailure = true) } test("cast between decimals with higher precision than source") { @@ -2417,25 +2417,19 @@ class CometCastSuite extends CometTestBase with AdaptiveSparkPlanHelper { val cometMessage = if (cometException.getCause != null) cometException.getCause.getMessage else cometException.getMessage - // this if branch should only check decimal to decimal cast and errors when output precision, scale causes overflow. - if (df.schema("a").dataType.typeName.contains("decimal") && toType.typeName - .contains("decimal") && sparkMessage.contains("cannot be represented as")) { - assert(cometMessage.contains("too large to store")) + if (CometSparkSessionExtensions.isSpark40Plus) { + // for Spark 4 we expect to sparkException carries the message + assert(sparkMessage.contains("SQLSTATE")) + // we compare a subset of the error message. Comet grabs the query + // context eagerly so it displays the call site at the + // line of code where the cast method was called, whereas spark grabs the context + // lazily and displays the call site at the line of code where the error is checked. + assert( + sparkMessage.startsWith( + cometMessage.substring(0, math.min(40, cometMessage.length)))) } else { - if (CometSparkSessionExtensions.isSpark40Plus) { - // for Spark 4 we expect to sparkException carries the message - assert(sparkMessage.contains("SQLSTATE")) - // we compare a subset of the error message. Comet grabs the query - // context eagerly so it displays the call site at the - // line of code where the cast method was called, whereas spark grabs the context - // lazily and displays the call site at the line of code where the error is checked. - assert( - sparkMessage.startsWith( - cometMessage.substring(0, math.min(40, cometMessage.length)))) - } else { - // for Spark 3.4 we expect to reproduce the error message exactly - assert(cometMessage == sparkMessage) - } + // for Spark 3.4 we expect to reproduce the error message exactly + assert(cometMessage == sparkMessage) } } } diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala index dbe252c7da..262797db40 100644 --- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala @@ -24,6 +24,7 @@ import java.time.{Duration, Period} import scala.util.Random import org.apache.hadoop.fs.Path +import org.apache.spark.SparkThrowable import org.apache.spark.sql.{Column, CometTestBase, DataFrame, Row} import org.apache.spark.sql.catalyst.expressions.{Alias, Cast, FromUnixTime, Literal, StructsToJson, TruncDate, TruncTimestamp} import org.apache.spark.sql.catalyst.optimizer.{ConvertToLocalRelation, OptimizeIn, SimplifyExtractValueOps} @@ -46,6 +47,19 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { val DIVIDE_BY_ZERO_EXCEPTION_MSG = """Division by zero. Use `try_divide` to tolerate divisor being 0 and return NULL instead""" + private def arithmeticError(error: Throwable): SparkThrowable = + Iterator + .iterate[Throwable](error)(_.getCause) + .takeWhile(_ != null) + .collectFirst { + // SparkArithmeticException is private[spark] in Spark 3.4, so this cross-version test + // cannot pattern match on its type directly. + case error: SparkThrowable + if error.getClass.getName == "org.apache.spark.SparkArithmeticException" => + error + } + .getOrElse(fail(s"Expected SparkArithmeticException, got $error")) + // Temporary test to verify checkSparkAnswer failure output labels Comet/Spark correctly. ignore("check output labels on mismatch") { val cometDf = Seq((1, "apple"), (2, "banana"), (3, "cherry")).toDF("id", "fruit") @@ -205,6 +219,29 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + test("ANSI decimal divide by zero raises a Spark error") { + withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") { + withTable("decimal_div_zero") { + sql("CREATE TABLE decimal_div_zero (a DECIMAL(10, 2), b DECIMAL(10, 2)) USING PARQUET") + sql("INSERT INTO decimal_div_zero VALUES (1.00, 0.00)") + + Seq("a / b", "a div b").foreach { expression => + val df = sql(s"SELECT $expression FROM decimal_div_zero") + checkCometOperators(stripAQEPlan(df.queryExecution.executedPlan)) + checkSparkAnswerMaybeThrows(df) match { + case (Some(sparkException), Some(cometException)) => + val expected = arithmeticError(sparkException) + val actual = arithmeticError(cometException) + assert(actual.getErrorClass == expected.getErrorClass) + assert(actual.getSqlState == expected.getSqlState) + assert(actual.getQueryContext.exists(_.fragment().contains(expression))) + case errors => fail(s"Expected Spark and Comet divide-by-zero errors, got $errors") + } + } + } + } + } + test("Integral Division Overflow Handling Matches Spark Behavior") { withSQLConf(SQLConf.ANSI_ENABLED.key -> "false") { withTable("t1") { @@ -1401,14 +1438,21 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { // 1.1e19 * 1.1e19 = 1.21e38 overflows DECIMAL(38,0). With ANSI mode on, both Spark and // Comet must throw — Comet must not panic or silently return null. Spark reports // NUMERIC_VALUE_OUT_OF_RANGE; Comet's WideDecimalBinaryExpr catches the overflow first - // and surfaces it as an arithmetic overflow error. + // and must surface the same structured Spark error. withSQLConf(CometConf.COMET_ENABLED.key -> "true", SQLConf.ANSI_ENABLED.key -> "true") { withParquetTable(Seq((BigDecimal("11000000000000000000"), 0)), "tbl") { val res = sql("SELECT _1 * _1 FROM tbl") + checkCometOperators(stripAQEPlan(res.queryExecution.executedPlan)) checkSparkAnswerMaybeThrows(res) match { case (Some(sparkExc), Some(cometExc)) => - assert(sparkExc.getMessage.contains("NUMERIC_VALUE_OUT_OF_RANGE")) - assert(cometExc.getMessage.toLowerCase.contains("overflow")) + val expected = arithmeticError(sparkExc) + val actual = arithmeticError(cometExc) + // Spark formats its pre-toPrecision Decimal, while Comet formats the rescaled i256 + // value (https://github.com/apache/datafusion-comet/issues/5211). This regression + // covers the structured error fields and query context. + assert(actual.getErrorClass == expected.getErrorClass) + assert(actual.getSqlState == expected.getSqlState) + assert(actual.getQueryContext.exists(_.fragment().contains("_1 * _1"))) case _ => fail("Expected exception for decimal overflow in ANSI mode") }