Skip to content

perf: use Arrow cast for decimal rescale check #5440

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sunchao merged 2 commits into from
Aug 23, 2026
+30 −135
Merged

perf: use Arrow cast for decimal rescale check #5440

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a2f10a5
refactor: use Arrow cast for decimal rescale check
peterxcli Aug 23, 2026
ae8ecf7
Merge branch 'main' into refactor/decimal-rescale-arrow-cast
sunchao Aug 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 1 addition & 3 deletions docs/source/contributor-guide/expression-audits/math_funcs.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
161 changes: 29 additions & 132 deletions native/spark-expr/src/math_funcs/internal/decimal_rescale_check.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<i128, ArrowError> {
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)
Expand All @@ -170,72 +110,19 @@ impl PhysicalExpr for DecimalRescaleCheckOverflow {

fn evaluate(&self, batch: &RecordBatch) -> datafusion::common::Result<ColumnarValue> {
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::<Decimal128Type>(&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<dyn PhysicalExpr>> {
Expand Down Expand Up @@ -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<Option<i128>>, precision: u8, scale: i8) -> RecordBatch {
Expand Down Expand Up @@ -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::<Decimal128Type>();
assert_eq!(arr.iter().collect::<Vec<_>>(), 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.
Expand Down Expand Up @@ -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::<Decimal128Type>();
Expand All @@ -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::<Decimal128Type>();
Expand Down
Loading