diff --git a/vortex-array/benches/binary_ops.rs b/vortex-array/benches/binary_ops.rs index 14c4c9fbe3f..3950d286f46 100644 --- a/vortex-array/benches/binary_ops.rs +++ b/vortex-array/benches/binary_ops.rs @@ -190,6 +190,38 @@ fn add_decimal_i128_nullable(bencher: Bencher) { bench_decimal(bencher, lhs, rhs, Operator::Add); } +#[divan::bench] +fn mul_decimal_i64_nonnull(bencher: Bencher) { + let lhs = decimal_i64_nonnull(0).into_array(); + let rhs = decimal_i64_nonnull(1_000_000).into_array(); + + bench_decimal(bencher, lhs, rhs, Operator::Mul); +} + +#[divan::bench] +fn mul_decimal_i128_nullable(bencher: Bencher) { + let lhs = decimal_i128_nullable(0, 7).into_array(); + let rhs = decimal_i128_nullable(1_000_000, 5).into_array(); + + bench_decimal(bencher, lhs, rhs, Operator::Mul); +} + +#[divan::bench] +fn div_decimal_i64_nonnull(bencher: Bencher) { + let lhs = decimal_i64_nonnull(0).into_array(); + let rhs = decimal_i64_nonnull(1_000_000).into_array(); + + bench_decimal(bencher, lhs, rhs, Operator::Div); +} + +#[divan::bench] +fn div_decimal_i128_nullable(bencher: Bencher) { + let lhs = decimal_i128_nullable(0, 7).into_array(); + let rhs = decimal_i128_nullable(1_000_000, 5).into_array(); + + bench_decimal(bencher, lhs, rhs, Operator::Div); +} + #[divan::bench] fn eq_i64_constant(bencher: Bencher) { let lhs = primitive_nonnull(0).into_array(); diff --git a/vortex-array/src/compute/conformance/binary_numeric.rs b/vortex-array/src/compute/conformance/binary_numeric.rs index a0d91715028..1750ffc8e5d 100644 --- a/vortex-array/src/compute/conformance/binary_numeric.rs +++ b/vortex-array/src/compute/conformance/binary_numeric.rs @@ -26,6 +26,8 @@ use std::fmt::Debug; use itertools::Itertools; use num_traits::Bounded; use num_traits::CheckedAdd; +use num_traits::CheckedDiv; +use num_traits::CheckedMul; use num_traits::CheckedSub; use num_traits::Float; use num_traits::Num; @@ -303,8 +305,16 @@ fn test_decimal_binary_numeric_with_scalar( let scalar = Scalar::decimal(value, decimal_dtype, array.dtype().nullability()); - // Decimal Mul/Div are not yet implemented. - for operator in [NumericOperator::Add, NumericOperator::Sub] { + let mut operators = vec![ + NumericOperator::Add, + NumericOperator::Sub, + NumericOperator::Mul, + ]; + if !value.is_zero() { + operators.push(NumericOperator::Div); + } + + for operator in operators { for lhs_is_array in [true, false] { test_decimal_binary_numeric_direction( array, @@ -329,7 +339,7 @@ fn test_decimal_binary_numeric_direction( ctx: &mut ExecutionCtx, ) { let result_decimal_dtype = numeric_op_result_decimal_dtype(decimal_dtype, operator) - .vortex_expect("decimal Add/Sub must have a result dtype"); + .vortex_expect("decimal arithmetic must have a result dtype"); let result_dtype = DType::Decimal(result_decimal_dtype, array.dtype().nullability()); let expected_results = expected_decimal_results( original_values, @@ -382,10 +392,22 @@ fn expected_decimal_results( let (Some(lhs), Some(rhs)) = (lhs.decimal_value(), rhs.decimal_value()) else { return Some(Scalar::null(result_dtype.clone())); }; + let lhs = lhs.as_i256(); + let rhs = rhs.as_i256(); let value = match operator { - NumericOperator::Add => lhs.as_i256().checked_add(&rhs.as_i256()), - NumericOperator::Sub => lhs.as_i256().checked_sub(&rhs.as_i256()), - NumericOperator::Mul | NumericOperator::Div => unreachable!(), + NumericOperator::Add => lhs.checked_add(&rhs), + NumericOperator::Sub => lhs.checked_sub(&rhs), + NumericOperator::Mul => lhs.checked_mul(&rhs), + NumericOperator::Div => { + let scale_power = result_decimal_dtype.scale(); + let factor = + i256::from_i128(10).checked_pow(scale_power.unsigned_abs() as u32)?; + if scale_power >= 0 { + lhs.checked_mul(&factor)?.checked_div(&rhs) + } else { + lhs.checked_div(&rhs.checked_mul(&factor)?) + } + } }?; let value = DecimalValue::try_from_i256(value, result_decimal_dtype).ok()?; Some(Scalar::decimal( diff --git a/vortex-array/src/scalar/typed_view/decimal/arithmetic.rs b/vortex-array/src/scalar/typed_view/decimal/arithmetic.rs index ff3007ca31d..2a47c9cf977 100644 --- a/vortex-array/src/scalar/typed_view/decimal/arithmetic.rs +++ b/vortex-array/src/scalar/typed_view/decimal/arithmetic.rs @@ -108,7 +108,7 @@ pub(crate) fn checked_decimal_numeric( result: DecimalDType, op: NumericOperator, ) -> Option { - let work = work_decimal_dtype(input, result, op); + let work = decimal_numeric_work_dtype(input, result, op); match_each_decimal_value_type!(DecimalType::smallest_decimal_value_type(&work), |W| { let value = checked_at_width::(lhs.cast::()?, rhs.cast::()?, result.scale(), op)?; DecimalValue::from(value).normalize(result) @@ -120,7 +120,7 @@ pub(crate) fn checked_decimal_numeric( /// The result precision covers Add, Sub and Mul, whose intermediates are the result itself. Div /// scales the dividend (or the divisor, for a negative result scale) by `10^result_scale` before /// dividing, so it needs room for `p + |result_scale|` digits. -fn work_decimal_dtype( +pub(crate) fn decimal_numeric_work_dtype( input: DecimalDType, result: DecimalDType, op: NumericOperator, diff --git a/vortex-array/src/scalar/typed_view/decimal/mod.rs b/vortex-array/src/scalar/typed_view/decimal/mod.rs index dd7219cf8dd..149fc8aeadf 100644 --- a/vortex-array/src/scalar/typed_view/decimal/mod.rs +++ b/vortex-array/src/scalar/typed_view/decimal/mod.rs @@ -8,6 +8,7 @@ mod dvalue; mod scalar; pub(crate) use arithmetic::decimal_numeric_result_dtype; +pub(crate) use arithmetic::decimal_numeric_work_dtype; pub use dvalue::DecimalValue; pub use scalar::DecimalScalar; diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/decimal.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/decimal.rs index a380aaa7996..fa9ffb0b5f1 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/decimal.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/decimal.rs @@ -4,20 +4,28 @@ //! Native execution of the arithmetic operators over decimal arrays. //! //! Both operands share a logical [`DecimalDType`] (equal precision and scale). Add and Sub apply -//! directly to the unscaled stored integers and are exact at that shared scale. The result reserves -//! one additional precision digit for a carry, capped at Vortex's maximum decimal precision. +//! directly to the unscaled stored integers and are exact at that shared scale. Mul takes the raw +//! product, which the doubled result scale leaves correctly scaled, and Div rescales the dividend +//! (or the divisor, for a negative result scale) before integer division. Result precision and +//! scale follow Arrow's rules — see [`decimal_numeric_result_dtype`]. //! -//! Lanes execute in a working width chosen so that in-precision inputs cannot spuriously -//! overflow an intermediate value. An operation that overflows the result precision on a valid -//! lane is an error; invalid lanes never error. +//! Lanes execute in a working width wide enough that in-precision inputs cannot spuriously +//! overflow an intermediate, then narrow to the result's own storage width. Every lane is still +//! checked at that width: [`DecimalArray`] does not validate its stored values against the +//! declared precision, so an out-of-precision value can reach a kernel and must not be able to +//! overflow it. An operation that overflows the result precision on a valid lane is an error; +//! invalid lanes never error. + +use std::ops::Mul; use num_traits::CheckedAdd; +use num_traits::CheckedDiv; +use num_traits::CheckedMul; use num_traits::CheckedSub; use vortex_buffer::Buffer; use vortex_compute::lane_kernels::LaneZip; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_mask::Mask; @@ -33,27 +41,16 @@ use crate::arrays::decimal::DecimalArrayExt; use crate::dtype::BigCast; use crate::dtype::DType; use crate::dtype::DecimalDType; +use crate::dtype::DecimalType; use crate::dtype::NativeDecimalType; use crate::match_each_decimal_value_type; use crate::scalar::DecimalValue; use crate::scalar::NumericOperator; use crate::scalar::Scalar; use crate::scalar::decimal_numeric_result_dtype; +use crate::scalar::decimal_numeric_work_dtype; use crate::validity::Validity; -/// Derive the result decimal dtype for a numeric operation over same-typed operands. -pub(crate) fn result_decimal_dtype( - input: DecimalDType, - op: NumericOperator, -) -> VortexResult { - match op { - NumericOperator::Add | NumericOperator::Sub => decimal_numeric_result_dtype(input, op), - NumericOperator::Mul | NumericOperator::Div => { - vortex_bail!("numeric operator {op} is not yet supported for decimal arrays") - } - } -} - /// Execute a numeric operation between two decimal arrays sharing a decimal dtype. pub(super) fn execute_numeric_decimal( lhs: &ArrayRef, @@ -66,7 +63,7 @@ pub(super) fn execute_numeric_decimal( .as_decimal_opt() .vortex_expect("inputs are both decimals"); - let result_decimal_dtype = result_decimal_dtype(*decimal_dtype, op)?; + let result_decimal_dtype = decimal_numeric_result_dtype(*decimal_dtype, op)?; let result_dtype = DType::Decimal( result_decimal_dtype, lhs.dtype().nullability() | rhs.dtype().nullability(), @@ -89,20 +86,30 @@ pub(super) fn execute_numeric_decimal( let validity = lhs.validity().and(rhs.validity())?; let valid_rows = validity.execute_mask(len, ctx)?; - match_each_decimal_value_type!( - DecimalType::smallest_decimal_value_type(&result_decimal_dtype), - |W| { - execute_decimal_at_width::( - &lhs, - &rhs, - op, - result_decimal_dtype, - &result_dtype, - validity, - &valid_rows, - ) + let work_dtype = decimal_numeric_work_dtype(*decimal_dtype, result_decimal_dtype, op); + match_each_decimal_value_type!(DecimalType::smallest_decimal_value_type(&work_dtype), |W| { + let constants = DecimalOpConstants::::new(result_decimal_dtype, op)?; + macro_rules! execute_typed { + ($Op:ty) => { + execute_decimal_typed::( + &lhs, + &rhs, + result_decimal_dtype, + &result_dtype, + validity, + &valid_rows, + &constants, + ) + }; } - ) + + match op { + NumericOperator::Add => execute_typed!(CheckedDecimalAdd), + NumericOperator::Sub => execute_typed!(CheckedDecimalSub), + NumericOperator::Mul => execute_typed!(CheckedDecimalMul), + NumericOperator::Div => execute_typed!(CheckedDecimalDiv), + } + }) } fn is_null_constant(array: &ArrayRef) -> bool { @@ -180,7 +187,7 @@ struct DecimalValueBounds { impl DecimalValueBounds { fn new(dtype: DecimalDType) -> Self { - let precision = dtype.precision() as usize; + let precision = usize::from(dtype.precision()); Self { lower_bound: W::MIN_BY_PRECISION[precision], upper_bound: W::MAX_BY_PRECISION[precision], @@ -193,74 +200,121 @@ impl DecimalValueBounds { } } -/// A checked fixed-point decimal operation on unscaled values at working width `W`. +/// Per-execution constants for a decimal operation at working width `W`, hoisted out of the +/// lane loop. +struct DecimalOpConstants { + bounds: DecimalValueBounds, + /// Arrow's division rescaling factors. Both are one for every other operator. + lhs_scale_factor: W, + rhs_scale_factor: W, +} + +impl DecimalOpConstants +where + W: NativeDecimalType + CheckedMul, +{ + fn new(result: DecimalDType, op: NumericOperator) -> VortexResult { + let one = ::from(1_i8).vortex_expect("one fits every decimal working width"); + let (lhs_scale_factor, rhs_scale_factor) = if op == NumericOperator::Div { + // Arrow scales the quotient by 10^(result_scale - lhs_scale + rhs_scale). Both + // Vortex operands share a dtype, so this simplifies to 10^result_scale. A negative + // exponent scales the divisor instead of the dividend. + let exponent = >::from(result.scale().unsigned_abs()); + if result.scale() >= 0 { + (decimal_scale_factor::(exponent)?, one) + } else { + (one, decimal_scale_factor::(exponent)?) + } + } else { + (one, one) + }; + + Ok(Self { + bounds: DecimalValueBounds::new(result), + lhs_scale_factor, + rhs_scale_factor, + }) + } +} + +fn decimal_scale_factor(exp: u32) -> VortexResult +where + W: NativeDecimalType + CheckedMul, +{ + let ten = ::from(10_i8).vortex_expect("ten fits every decimal working width"); + let mut factor = + ::from(1_i8).vortex_expect("one fits every decimal working width"); + for _ in 0..exp { + factor = factor.checked_mul(&ten).ok_or_else(|| { + vortex_err!( + InvalidArgument: + "decimal scale factor 10^{exp} cannot be represented at the working width" + ) + })?; + } + Ok(factor) +} + +/// A checked decimal operation on unscaled values at working width `W`. trait CheckedDecimalOp { const ERROR: &'static str; - fn apply(lhs: W, rhs: W, bounds: &DecimalValueBounds) -> Option + fn apply(lhs: W, rhs: W, constants: &DecimalOpConstants) -> Option where - W: NativeDecimalType + CheckedAdd + CheckedSub; + W: NativeDecimalType + CheckedAdd + CheckedSub + CheckedMul + CheckedDiv + Mul; } struct CheckedDecimalAdd; struct CheckedDecimalSub; +struct CheckedDecimalMul; + +struct CheckedDecimalDiv; + impl CheckedDecimalOp for CheckedDecimalAdd { const ERROR: &'static str = "decimal overflow in checked add"; - fn apply(lhs: W, rhs: W, bounds: &DecimalValueBounds) -> Option + fn apply(lhs: W, rhs: W, constants: &DecimalOpConstants) -> Option where - W: NativeDecimalType + CheckedAdd + CheckedSub, + W: NativeDecimalType + CheckedAdd + CheckedSub + CheckedMul + CheckedDiv + Mul, { - bounds.in_precision(lhs.checked_add(&rhs)?) + constants.bounds.in_precision(lhs.checked_add(&rhs)?) } } impl CheckedDecimalOp for CheckedDecimalSub { const ERROR: &'static str = "decimal overflow in checked sub"; - fn apply(lhs: W, rhs: W, bounds: &DecimalValueBounds) -> Option + fn apply(lhs: W, rhs: W, constants: &DecimalOpConstants) -> Option where - W: NativeDecimalType + CheckedAdd + CheckedSub, + W: NativeDecimalType + CheckedAdd + CheckedSub + CheckedMul + CheckedDiv + Mul, { - bounds.in_precision(lhs.checked_sub(&rhs)?) + constants.bounds.in_precision(lhs.checked_sub(&rhs)?) } } -fn execute_decimal_at_width( - lhs: &DecimalOperand, - rhs: &DecimalOperand, - op: NumericOperator, - result_decimal_dtype: DecimalDType, - result_dtype: &DType, - validity: Validity, - valid_rows: &Mask, -) -> VortexResult -where - W: NativeDecimalType + CheckedAdd + CheckedSub, - DecimalValue: From, -{ - macro_rules! execute_typed { - ($Op:ty) => { - execute_decimal_typed::( - lhs, - rhs, - result_decimal_dtype, - result_dtype, - validity, - valid_rows, - ) - }; +impl CheckedDecimalOp for CheckedDecimalMul { + const ERROR: &'static str = "decimal overflow in checked mul"; + + fn apply(lhs: W, rhs: W, constants: &DecimalOpConstants) -> Option + where + W: NativeDecimalType + CheckedAdd + CheckedSub + CheckedMul + CheckedDiv + Mul, + { + constants.bounds.in_precision(lhs.checked_mul(&rhs)?) } +} + +impl CheckedDecimalOp for CheckedDecimalDiv { + const ERROR: &'static str = "decimal overflow or division by zero in checked div"; - match op { - NumericOperator::Add => execute_typed!(CheckedDecimalAdd), - NumericOperator::Sub => execute_typed!(CheckedDecimalSub), - NumericOperator::Mul | NumericOperator::Div => vortex_bail!( - "numeric operator {:?} is not yet supported for decimal arrays", - op - ), + fn apply(lhs: W, rhs: W, constants: &DecimalOpConstants) -> Option + where + W: NativeDecimalType + CheckedAdd + CheckedSub + CheckedMul + CheckedDiv + Mul, + { + let lhs = lhs.checked_mul(&constants.lhs_scale_factor)?; + let rhs = rhs.checked_mul(&constants.rhs_scale_factor)?; + constants.bounds.in_precision(lhs.checked_div(&rhs)?) } } @@ -271,25 +325,25 @@ fn execute_decimal_typed( result_dtype: &DType, validity: Validity, valid_rows: &Mask, + constants: &DecimalOpConstants, ) -> VortexResult where - W: NativeDecimalType + CheckedAdd + CheckedSub, + W: NativeDecimalType + CheckedAdd + CheckedSub + CheckedMul + CheckedDiv + Mul, DecimalValue: From, Op: CheckedDecimalOp, { let len = lhs.len(); - let bounds = DecimalValueBounds::::new(result_decimal_dtype); let values = match (lhs, rhs) { (DecimalOperand::Array { values: lhs, .. }, DecimalOperand::Array { values: rhs, .. }) => { - checked_decimal_arrays::(lhs, rhs, &bounds, valid_rows) + checked_decimal_arrays::(lhs, rhs, constants, valid_rows) } (DecimalOperand::Array { values: lhs, .. }, DecimalOperand::Constant { value, .. }) => { let rhs = typed_constant::(value); match_each_decimal_value_type!(lhs.values_type(), |L| { let lhs = lhs.buffer::(); checked_lanes(lhs.as_slice(), valid_rows, |lhs| { - Op::apply(::from(lhs)?, rhs, &bounds) + Op::apply(::from(lhs)?, rhs, constants) }) }) } @@ -298,7 +352,7 @@ where match_each_decimal_value_type!(rhs.values_type(), |R| { let rhs = rhs.buffer::(); checked_lanes(rhs.as_slice(), valid_rows, |rhs| { - Op::apply(lhs, ::from(rhs)?, &bounds) + Op::apply(lhs, ::from(rhs)?, constants) }) }) } @@ -308,14 +362,13 @@ where ) => { let lhs = typed_constant::(lhs); let rhs = typed_constant::(rhs); - let value = Op::apply(lhs, rhs, &bounds) + let value = Op::apply(lhs, rhs, constants) .ok_or_else(|| vortex_err!(InvalidArgument: "{}", Op::ERROR))?; + let value = DecimalValue::from(value) + .normalize(result_decimal_dtype) + .vortex_expect("bounds-checked result fits the result precision"); return Ok(ConstantArray::new( - Scalar::decimal( - DecimalValue::from(value), - result_decimal_dtype, - result_dtype.nullability(), - ), + Scalar::decimal(value, result_decimal_dtype, result_dtype.nullability()), len, ) .into_array()); @@ -323,22 +376,48 @@ where } .map_err(|_lane| vortex_err!(InvalidArgument: "{}", Op::ERROR))?; - Ok(DecimalArray::new( + Ok(decimal_array_narrowed( values, result_decimal_dtype, validity.union_nullability(result_dtype.nullability()), - ) - .into_array()) + )) +} + +/// Build the result array, narrowing to the dtype's own storage width when the working width is +/// wider than it. Only division picks a working width above the result precision, and only for a +/// negative result scale, so this copies in a corner case rather than on the common path. +fn decimal_array_narrowed( + values: Buffer, + decimal_dtype: DecimalDType, + validity: Validity, +) -> ArrayRef { + let target = DecimalType::smallest_decimal_value_type(&decimal_dtype); + if target == W::DECIMAL_TYPE { + return DecimalArray::new(values, decimal_dtype, validity).into_array(); + } + + match_each_decimal_value_type!(target, |O| { + let narrowed: Buffer = values + .as_slice() + .iter() + .copied() + .map(|value| { + ::from(value) + .vortex_expect("precision-checked decimal result must fit the output width") + }) + .collect(); + DecimalArray::new(narrowed, decimal_dtype, validity).into_array() + }) } fn checked_decimal_arrays( lhs: &DecimalArray, rhs: &DecimalArray, - bounds: &DecimalValueBounds, + constants: &DecimalOpConstants, valid_rows: &Mask, ) -> Result, usize> where - W: NativeDecimalType + CheckedAdd + CheckedSub, + W: NativeDecimalType + CheckedAdd + CheckedSub + CheckedMul + CheckedDiv + Mul, Op: CheckedDecimalOp, { debug_assert_eq!(lhs.len(), rhs.len()); @@ -353,7 +432,7 @@ where Op::apply( ::from(lhs)?, ::from(rhs)?, - bounds, + constants, ) }, ) diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs index 7211ae9ed6a..6622e08f82b 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs @@ -13,8 +13,6 @@ mod primitive; mod tests; use decimal::execute_numeric_decimal; -use decimal::result_decimal_dtype; -pub(crate) use decimal::result_decimal_dtype as numeric_op_result_decimal_dtype; pub(crate) use primitive::PrimitiveOperand; use primitive::execute_numeric_primitive; use vortex_error::VortexResult; @@ -26,6 +24,7 @@ use crate::ExecutionCtx; use crate::IntoArray; use crate::dtype::DType; use crate::scalar::NumericOperator; +pub(crate) use crate::scalar::decimal_numeric_result_dtype as numeric_op_result_decimal_dtype; /// Execute a numeric operation between two arrays. pub(crate) fn execute_numeric( @@ -74,9 +73,10 @@ fn build_empty_result( let nullability = lhs.dtype().nullability() | rhs.dtype().nullability(); let result_dtype = match lhs.dtype() { DType::Primitive(..) => lhs.dtype().with_nullability(nullability), - DType::Decimal(decimal_dtype, _) => { - DType::Decimal(result_decimal_dtype(*decimal_dtype, op)?, nullability) - } + DType::Decimal(decimal_dtype, _) => DType::Decimal( + numeric_op_result_decimal_dtype(*decimal_dtype, op)?, + nullability, + ), _ => unreachable!("dtype is either Primitive or Decimal"), }; diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs index 811797109f8..7ee98ae717e 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs @@ -6,7 +6,7 @@ use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_error::VortexResult; -use super::result_decimal_dtype; +use super::numeric_op_result_decimal_dtype as result_decimal_dtype; use crate::ArrayRef; use crate::Columnar; use crate::IntoArray; @@ -378,6 +378,12 @@ fn decimal_constant(value: impl Into, dtype: DecimalDType, len: us #[rstest] #[case::add(NumericOperator::Add, [150i64, 225], [1050i64, 1225])] #[case::sub(NumericOperator::Sub, [150i64, 225], [750i64, 775])] +#[case::mul(NumericOperator::Mul, [150i64, 225], [135_000i64, 225_000])] +#[case::div( + NumericOperator::Div, + [150i64, 225], + [6_000_000i64, 4_444_444] +)] fn test_decimal_array_array( #[case] op: NumericOperator, #[case] rhs: [i64; 2], @@ -457,6 +463,50 @@ fn test_decimal_value_outside_working_width_errors() { assert!(decimal_binary(lhs, rhs, Operator::Add).is_err()); } +#[test] +fn test_decimal_div_negative_result_scale() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + // A negative result scale scales the divisor rather than the dividend: 5e12 / 2e8 is 25_000, + // truncated to 2 at the decimal(6, -4) result scale. The i64 working width is also wider than + // the i32 result storage, so this narrows on the way out. + let dtype = DecimalDType::new(10, -8); + let lhs = DecimalArray::from_iter::([50_000], dtype).into_array(); + let rhs = DecimalArray::from_iter::([2], dtype).into_array(); + + let result = decimal_binary(lhs, rhs, Operator::Div)?; + assert_arrays_eq!( + result, + DecimalArray::from_iter::([2], DecimalDType::new(6, -4)), + &mut ctx + ); + Ok(()) +} + +#[test] +fn test_decimal_mul_value_outside_precision_errors() { + // `DecimalArray::new` does not validate stored values against the declared precision, so Mul + // cannot assume its inputs are in-precision: 500 * 500 is 250_000, well past the 99_999 that + // the decimal(5, 0) result can represent. + let dtype = DecimalDType::new(2, 0); + let value = i256::from_i128(500); + let lhs = DecimalArray::new(buffer![value], dtype, Validity::NonNullable).into_array(); + let rhs = DecimalArray::new(buffer![value], dtype, Validity::NonNullable).into_array(); + + assert!(decimal_binary(lhs, rhs, Operator::Mul).is_err()); +} + +#[test] +fn test_decimal_mul_value_outside_working_width_errors() { + // 50_000 * 50_000 overflows the i32 working width chosen for a decimal(5, 0) result, which + // an unchecked multiply would wrap in release and panic on in debug. + let dtype = DecimalDType::new(2, 0); + let value = i256::from_i128(50_000); + let lhs = DecimalArray::new(buffer![value], dtype, Validity::NonNullable).into_array(); + let rhs = DecimalArray::new(buffer![value], dtype, Validity::NonNullable).into_array(); + + assert!(decimal_binary(lhs, rhs, Operator::Mul).is_err()); +} + #[test] fn test_decimal_value_outside_working_width_on_null_lane_ignored() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); @@ -498,6 +548,39 @@ fn test_decimal_overflow_on_null_lane_ignored() { ); } +#[test] +fn test_decimal_divide_by_zero_on_null_lane_ignored() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DecimalDType::new(10, 2); + let lhs = DecimalArray::new( + buffer![100i64, 1_000], + dtype, + Validity::from_iter([false, true]), + ) + .into_array(); + let rhs = DecimalArray::from_iter::([0, 200], dtype).into_array(); + + let result = decimal_binary(lhs, rhs, Operator::Div)?; + assert_arrays_eq!( + result, + DecimalArray::from_option_iter::( + [None, Some(5_000_000)], + result_decimal_dtype(dtype, NumericOperator::Div)?, + ), + &mut ctx + ); + Ok(()) +} + +#[test] +fn test_decimal_divide_by_zero_on_valid_lane_errors() { + let dtype = DecimalDType::new(10, 2); + let lhs = DecimalArray::from_iter::([100], dtype).into_array(); + let rhs = DecimalArray::from_iter::([0], dtype).into_array(); + + assert!(decimal_binary(lhs, rhs, Operator::Div).is_err()); +} + #[test] fn test_decimal_add_reserves_carry_digit() { let mut ctx = array_session().create_execution_ctx(); @@ -513,6 +596,47 @@ fn test_decimal_add_reserves_carry_digit() { ); } +#[test] +fn test_decimal_mul_widens_before_multiplying() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DecimalDType::new(38, 0); + let max = ::MAX_BY_PRECISION[38]; + let widened_max = i256::from_i128(max); + let result_dtype = result_decimal_dtype(dtype, NumericOperator::Mul)?; + + let result = decimal_binary( + DecimalArray::from_iter::([max], dtype).into_array(), + DecimalArray::from_iter::([max], dtype).into_array(), + Operator::Mul, + )?; + assert_arrays_eq!( + result, + DecimalArray::from_iter::([widened_max * widened_max], result_dtype), + &mut ctx + ); + Ok(()) +} + +#[test] +fn test_decimal_mul_above_result_precision_errors() { + let dtype = DecimalDType::new(39, 0); + let one = i256::from_i128(1); + let ten_to_38 = ::MAX_BY_PRECISION[38] + one; + let value = ten_to_38 * i256::from_i128(2); + let product = value * value; + + // The product fits the native i256 width but not the capped precision-76 result. + assert!(product > ::MAX_BY_PRECISION[76]); + assert!( + decimal_binary( + DecimalArray::from_iter::([value], dtype).into_array(), + DecimalArray::from_iter::([value], dtype).into_array(), + Operator::Mul, + ) + .is_err() + ); +} + #[rstest] #[case::precision_2( DecimalArray::from_iter::([10, 20], DecimalDType::new(2, 0)), @@ -638,9 +762,16 @@ fn test_decimal_nullable_constant_preserves_nullable_output() -> VortexResult<() } #[rstest] -#[case::null_lhs(true)] -#[case::null_rhs(false)] -fn test_decimal_null_constant_yields_all_null(#[case] null_lhs: bool) -> VortexResult<()> { +#[case::add_null_lhs(true, NumericOperator::Add)] +#[case::add_null_rhs(false, NumericOperator::Add)] +#[case::mul_null_lhs(true, NumericOperator::Mul)] +#[case::mul_null_rhs(false, NumericOperator::Mul)] +#[case::div_null_lhs(true, NumericOperator::Div)] +#[case::div_null_rhs(false, NumericOperator::Div)] +fn test_decimal_null_constant_yields_all_null( + #[case] null_lhs: bool, + #[case] op: NumericOperator, +) -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let dtype = DecimalDType::new(10, 2); let values = DecimalArray::from_iter::([100, 200], dtype).into_array(); @@ -655,16 +786,11 @@ fn test_decimal_null_constant_yields_all_null(#[case] null_lhs: bool) -> VortexR (values, null_constant) }; - let result = lhs - .binary(rhs, Operator::Add)? - .execute::(&mut ctx)?; + let result = lhs.binary(rhs, op.into())?.execute::(&mut ctx)?; assert!(matches!(&result, Columnar::Constant(_))); assert_arrays_eq!( result.into_array(), - DecimalArray::from_option_iter::( - [None, None], - result_decimal_dtype(dtype, NumericOperator::Add)?, - ), + DecimalArray::from_option_iter::([None, None], result_decimal_dtype(dtype, op)?,), &mut ctx ); Ok(()) @@ -691,34 +817,45 @@ fn test_decimal_constant_wider_than_array_storage() -> VortexResult<()> { Ok(()) } -#[test] -fn test_decimal_empty() -> VortexResult<()> { +#[rstest] +#[case::add(NumericOperator::Add)] +#[case::sub(NumericOperator::Sub)] +#[case::mul(NumericOperator::Mul)] +#[case::div(NumericOperator::Div)] +fn test_decimal_empty(#[case] op: NumericOperator) -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let dtype = DecimalDType::new(10, 2); let empty = DecimalArray::from_iter::([], dtype).into_array(); - let result = decimal_binary(empty.clone(), empty, Operator::Add)?; + let result = decimal_binary(empty.clone(), empty, op.into())?; assert_arrays_eq!( result, - DecimalArray::from_iter::([], result_decimal_dtype(dtype, NumericOperator::Add)?,), + DecimalArray::from_iter::([], result_decimal_dtype(dtype, op)?,), &mut ctx ); Ok(()) } -#[test] -fn test_decimal_constant_constant_folds() -> VortexResult<()> { +#[rstest] +#[case::add(NumericOperator::Add, 200)] +#[case::sub(NumericOperator::Sub, 100)] +#[case::mul(NumericOperator::Mul, 7_500)] +#[case::div(NumericOperator::Div, 3_000_000)] +fn test_decimal_constant_constant_folds( + #[case] op: NumericOperator, + #[case] expected: i128, +) -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let dtype = DecimalDType::new(10, 2); let lhs = decimal_constant(150i64, dtype, 3); let rhs = decimal_constant(50i64, dtype, 3); - let result = decimal_binary(lhs, rhs, Operator::Add)?; + let result = decimal_binary(lhs, rhs, op.into())?; assert_arrays_eq!( result, - DecimalArray::from_iter::( - [200, 200, 200], - result_decimal_dtype(dtype, NumericOperator::Add)?, + DecimalArray::from_iter::( + [i256::from_i128(expected); 3], + result_decimal_dtype(dtype, op)?, ), &mut ctx ); @@ -726,11 +863,22 @@ fn test_decimal_constant_constant_folds() -> VortexResult<()> { } #[rstest] -#[case::mul(Operator::Mul)] -#[case::div(Operator::Div)] -fn test_decimal_mul_div_unsupported(#[case] op: Operator) { - let dtype = DecimalDType::new(10, 2); - let values = DecimalArray::from_iter::([100], dtype).into_array(); +#[case::add(NumericOperator::Add, DecimalDType::new(11, 2))] +#[case::sub(NumericOperator::Sub, DecimalDType::new(11, 2))] +#[case::mul(NumericOperator::Mul, DecimalDType::new(21, 4))] +#[case::div(NumericOperator::Div, DecimalDType::new(16, 6))] +fn test_decimal_result_dtype( + #[case] op: NumericOperator, + #[case] expected: DecimalDType, +) -> VortexResult<()> { + assert_eq!( + result_decimal_dtype(DecimalDType::new(10, 2), op)?, + expected + ); + Ok(()) +} - assert!(decimal_binary(values.clone(), values, op).is_err()); +#[test] +fn test_decimal_mul_result_scale_overflow_errors() { + assert!(result_decimal_dtype(DecimalDType::new(40, 40), NumericOperator::Mul).is_err()); }