From a2f10a54ca137504a5320d1bf05647c26fdee4ca Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 23 Aug 2026 21:08:45 +0800 Subject: [PATCH] refactor: use Arrow cast for decimal rescale check --- .../expression-audits/math_funcs.md | 4 +- .../internal/decimal_rescale_check.rs | 161 ++++-------------- 2 files changed, 30 insertions(+), 135 deletions(-) diff --git a/docs/source/contributor-guide/expression-audits/math_funcs.md b/docs/source/contributor-guide/expression-audits/math_funcs.md index 87bfa27cb86..6a9d7868519 100644 --- a/docs/source/contributor-guide/expression-audits/math_funcs.md +++ b/docs/source/contributor-guide/expression-audits/math_funcs.md @@ -121,9 +121,7 @@ Internal decimal wrapper emitted around every decimal `+ - * /`, `sum`, and `avg ## DecimalRescaleCheckOverflow (internal) -Internal fused expression that rescales a Decimal128 value (changing scale) and checks output precision in one pass, replacing the `CheckOverflow(Cast(expr, Decimal128(p, s)))` pattern used by decimal-to-decimal casts. Native impl: `math_funcs/internal/decimal_rescale_check.rs`. - -- Performance (tuned 2026-07-15, PR [#4938](https://github.com/apache/datafusion-comet/pull/4938)): the legacy path ran `null_if_overflow_precision` (a second full pass that allocates a new array) on every batch to turn overflow sentinels into nulls, even when nothing overflowed. Now that pass runs only when a sentinel is present (`contains(&i128::MAX)`, short-circuiting), so the common no-overflow case skips the allocation. 8 to 26% faster on no-overflow shapes; overflow and ANSI shapes unchanged. Benchmark: `benches/decimal_rescale.rs`. +Internal fused expression that replaces the `CheckOverflow(Cast(expr, Decimal128(p, s)))` pattern used by decimal-to-decimal casts. It delegates rescaling, HALF_UP rounding, and output-precision validation to Arrow's single-pass decimal cast kernel. Arrow writes nulls directly in legacy mode and returns an error in ANSI mode. Native impl: `math_funcs/internal/decimal_rescale_check.rs`; benchmark: `benches/decimal_rescale.rs`. ## e 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 fea23992020..563370d9750 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,11 +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 arrow::array::{as_primitive_array, Array, ArrayRef, Decimal128Array}; -use arrow::datatypes::{DataType, Decimal128Type, Schema}; -use arrow::error::ArrowError; +use arrow::compute::CastOptions; +use arrow::datatypes::{DataType, Schema}; use arrow::record_batch::RecordBatch; -use datafusion::common::{DataFusionError, ScalarValue}; +use datafusion::common::DataFusionError; use datafusion::logical_expr::ColumnarValue; use datafusion::physical_expr::PhysicalExpr; use std::hash::Hash; @@ -93,65 +92,6 @@ impl Display for DecimalRescaleCheckOverflow { } } -/// Maximum absolute value for a given decimal precision: 10^p - 1. -/// Precision must be <= 38 (max for Decimal128). -#[inline] -fn precision_bound(precision: u8) -> i128 { - assert!( - precision <= 38, - "precision_bound: precision {precision} exceeds maximum 38" - ); - 10i128.pow(precision as u32) - 1 -} - -/// 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. -#[inline] -fn rescale_and_check( - value: i128, - delta: i8, - scale_factor: i128, - bound: i128, - fail_on_error: bool, -) -> Result { - let rescaled = if delta > 0 { - // Scale up: multiply. Check for overflow. - match value.checked_mul(scale_factor) { - Some(v) => v, - None => { - if fail_on_error { - return Err(ArrowError::ComputeError( - "Decimal overflow during rescale".to_string(), - )); - } - return Ok(i128::MAX); // sentinel - } - } - } 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 - } else { - value - }; - - // Precision check - if rescaled.abs() > bound { - if fail_on_error { - return Err(ArrowError::ComputeError( - "Decimal overflow: value does not fit in precision".to_string(), - )); - } - Ok(i128::MAX) // sentinel for null_if_overflow_precision - } else { - Ok(rescaled) - } -} - impl PhysicalExpr for DecimalRescaleCheckOverflow { fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { Display::fmt(self, f) @@ -170,72 +110,19 @@ 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 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 { + if !matches!(arg.data_type(), DataType::Decimal128(_, _)) { return Err(DataFusionError::Execution(format!( - "DecimalRescaleCheckOverflow: scale delta {delta} exceeds maximum supported range" + "DecimalRescaleCheckOverflow expects Decimal128, but found {arg:?}" ))); } - let scale_factor = 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; - let s_out = self.output_scale; - - match arg { - ColumnarValue::Array(array) - if matches!(array.data_type(), DataType::Decimal128(_, _)) => - { - let decimal_array = as_primitive_array::(&array); - - let result: Decimal128Array = - arrow::compute::kernels::arity::try_unary(decimal_array, |value| { - rescale_and_check(value, delta, scale_factor, bound, fail_on_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 - // 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. - result.null_if_overflow_precision(p_out) - } else { - result - }; - - let result = result - .with_precision_and_scale(p_out, s_out) - .map(|a| Arc::new(a) as ArrayRef)?; - - Ok(ColumnarValue::Array(result)) - } - 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))?; - if r == i128::MAX { - None - } else { - Some(r) - } - } - None => None, - }; - Ok(ColumnarValue::Scalar(ScalarValue::Decimal128( - new_v, p_out, s_out, - ))) - } - v => Err(DataFusionError::Execution(format!( - "DecimalRescaleCheckOverflow expects Decimal128, but found {v:?}" - ))), - } + + arg.cast_to( + &DataType::Decimal128(self.output_precision, self.output_scale), + Some(&CastOptions { + safe: !self.fail_on_error, + ..Default::default() + }), + ) } fn children(&self) -> Vec<&Arc> { @@ -265,9 +152,10 @@ impl PhysicalExpr for DecimalRescaleCheckOverflow { #[cfg(test)] mod tests { use super::*; - use arrow::array::{AsArray, Decimal128Array}; - use arrow::datatypes::{Field, Schema}; + use arrow::array::{Array, ArrayRef, AsArray, Decimal128Array}; + use arrow::datatypes::{Decimal128Type, Field, Schema}; use arrow::record_batch::RecordBatch; + use datafusion::common::ScalarValue; use datafusion::physical_expr::expressions::Column; fn make_batch(values: Vec>, precision: u8, scale: i8) -> RecordBatch { @@ -322,6 +210,17 @@ mod tests { assert_eq!(arr.value(2), -124); // -1.24 } + #[test] + fn test_scale_down_rounding_overflow() { + let batch = make_batch(vec![Some(994), Some(995), Some(-995)], 3, 1); + + let result = eval_expr(&batch, 1, 2, 0, false).unwrap(); + let arr = result.as_primitive::(); + assert_eq!(arr.iter().collect::>(), vec![Some(99), None, None]); + + assert!(eval_expr(&batch, 1, 2, 0, true).is_err()); + } + #[test] fn test_same_scale_precision_check_only() { // Same scale, just check precision. Value 999 fits in precision 3, 1000 does not. @@ -351,8 +250,7 @@ mod tests { #[test] fn test_overflow_with_nulls_legacy() { - // Mixes valid, overflowing, and null inputs so the sentinel fallback path runs with - // nulls present: overflow and null both yield null, valid values are preserved. + // Overflow and null both yield null, while valid values are preserved. let batch = make_batch(vec![Some(150), Some(10_000), None, Some(250)], 10, 2); let result = eval_expr(&batch, 2, 4, 2, false).unwrap(); let arr = result.as_primitive::(); @@ -364,8 +262,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 whole result is null. 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::();