Skip to content

Commit 02d6a3b

Browse files
committed
fix: add defensive checks for CheckOverflow bypass and multiply scale-up
- Validate WideDecimalBinaryExpr output type matches CheckOverflow data_type before bypassing the overflow check - Handle s_out > natural_scale (scale-up) in multiply path for consistency with add/subtract
1 parent 1141b39 commit 02d6a3b

2 files changed

Lines changed: 17 additions & 5 deletions

File tree

native/core/src/execution/planner.rs

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -383,12 +383,16 @@ impl PhysicalPlanner {
383383
let fail_on_error = expr.fail_on_error;
384384

385385
// WideDecimalBinaryExpr already handles overflow — skip redundant check
386+
// but only if its output type matches CheckOverflow's declared type
386387
if child
387388
.as_any()
388389
.downcast_ref::<WideDecimalBinaryExpr>()
389390
.is_some()
390391
{
391-
return Ok(child);
392+
let child_type = child.data_type(&input_schema)?;
393+
if child_type == data_type {
394+
return Ok(child);
395+
}
392396
}
393397

394398
// Fuse Cast(Decimal128→Decimal128) + CheckOverflow into single rescale+check

native/spark-expr/src/math_funcs/wide_decimal_binary_expr.rs

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -247,17 +247,25 @@ impl PhysicalExpr for WideDecimalBinaryExpr {
247247
}
248248
WideDecimalOp::Multiply => {
249249
let natural_scale = s1 + s2;
250-
let need_rescale = s_out < natural_scale;
251-
let rescale_divisor = if need_rescale {
252-
i256_pow10((natural_scale - s_out) as u32)
250+
let scale_diff = natural_scale as i16 - s_out as i16;
251+
let (need_scale_down, need_scale_up) = (scale_diff > 0, scale_diff < 0);
252+
let rescale_divisor = if need_scale_down {
253+
i256_pow10(scale_diff as u32)
254+
} else {
255+
i256::ONE
256+
};
257+
let scale_up_factor = if need_scale_up {
258+
i256_pow10((-scale_diff) as u32)
253259
} else {
254260
i256::ONE
255261
};
256262

257263
arrow::compute::kernels::arity::try_binary(left, right, |l, r| {
258264
let raw = i256::from_i128(l).wrapping_mul(i256::from_i128(r));
259-
let result = if need_rescale {
265+
let result = if need_scale_down {
260266
div_round_half_up(raw, rescale_divisor)
267+
} else if need_scale_up {
268+
raw.wrapping_mul(scale_up_factor)
261269
} else {
262270
raw
263271
};

0 commit comments

Comments
 (0)