From 2f2e00dcdc5af392cefb333f9627fcb20fb9dcec Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 10 Jul 2026 10:33:21 -0700 Subject: [PATCH 1/5] Split binary numeric execution into a module tree Mechanical move mirroring the compare split: dtype-generic lane machinery stays in numeric/mod.rs (bounds relaxed to T: Default so wider decimal types can reuse it), primitive-specific code moves to numeric/primitive.rs, tests to numeric/tests.rs. Signed-off-by: Matt Katz --- .../src/scalar_fn/fns/binary/numeric/mod.rs | 189 +++++++++ .../{numeric.rs => numeric/primitive.rs} | 386 +----------------- .../src/scalar_fn/fns/binary/numeric/tests.rs | 207 ++++++++++ 3 files changed, 407 insertions(+), 375 deletions(-) create mode 100644 vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs rename vortex-array/src/scalar_fn/fns/binary/{numeric.rs => numeric/primitive.rs} (66%) create mode 100644 vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs new file mode 100644 index 00000000000..702ec47adf0 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs @@ -0,0 +1,189 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Native execution of the arithmetic operators (Add/Sub/Mul/Div) of the [`Binary`] scalar +//! function. There is no Arrow fallback. +//! +//! [`Binary`]: super::Binary + +mod primitive; +#[cfg(test)] +mod tests; + +pub(crate) use primitive::PrimitiveOperand; +use vortex_buffer::BitBuffer; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::scalar::NumericOperator; + +/// Execute a numeric operation between two arrays. +pub(crate) fn execute_numeric( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: NumericOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + primitive::execute_numeric_primitive(lhs, rhs, op, ctx) +} + +/// The values produced by a checked lane loop, plus whether any lane failed. +struct CheckedValues { + values: Buffer, + failed: bool, +} + +impl CheckedValues { + fn zeroed(len: usize) -> Self { + Self { + values: Buffer::::zeroed(len), + failed: false, + } + } + + fn failed(len: usize) -> Self { + Self { + values: Buffer::::zeroed(len), + failed: true, + } + } +} + +// Checked one-pass ops delay early exit until the end of a small block. This +// keeps the loop generic while avoiding a branch-driven exit decision on every +// lane; it is deliberately independent of mask density or input length. +const CHECKED_BLOCK_LANES: usize = 16; + +fn checked_all_lanes(len: usize, mut checked_at: F) -> CheckedValues +where + T: Default, + F: FnMut(usize) -> Option, +{ + let mut values = BufferMut::::with_capacity(len); + let mut base = 0; + + while base + CHECKED_BLOCK_LANES <= len { + let mut block_failed = false; + for idx in base..base + CHECKED_BLOCK_LANES { + match checked_at(idx) { + Some(value) => { + // SAFETY: the buffer is allocated with capacity `len`, and + // this loop pushes at most one value for each `idx`. + unsafe { values.push_unchecked(value) }; + } + None => { + block_failed = true; + // SAFETY: the buffer is allocated with capacity `len`, and + // this loop pushes at most one value for each `idx`. + unsafe { values.push_unchecked(T::default()) }; + } + } + } + + if block_failed { + return CheckedValues::failed(len); + } + base += CHECKED_BLOCK_LANES; + } + + for idx in base..len { + let Some(value) = checked_at(idx) else { + return CheckedValues::failed(len); + }; + // SAFETY: the buffer is allocated with capacity `len`, and this loop + // pushes at most one value for each `idx`. + unsafe { values.push_unchecked(value) }; + } + + CheckedValues { + values: values.freeze(), + failed: false, + } +} + +fn checked_valid_lanes( + len: usize, + valid_bits: &BitBuffer, + mut checked_at: F, +) -> CheckedValues +where + T: Default, + F: FnMut(usize) -> Option, +{ + let mut values = BufferMut::::zeroed(len); + let mut failed = false; + { + let values = values.as_mut_slice(); + for_each_valid_idx(len, valid_bits, |idx| { + let Some(value) = checked_at(idx) else { + failed = true; + return false; + }; + values[idx] = value; + true + }); + } + + CheckedValues { + values: values.freeze(), + failed, + } +} + +fn any_valid_error(len: usize, valid_bits: &BitBuffer, is_error: F) -> bool +where + F: Fn(usize) -> bool, +{ + !for_each_valid_idx(len, valid_bits, |idx| !is_error(idx)) +} + +fn for_each_valid_idx(len: usize, valid_bits: &BitBuffer, mut f: F) -> bool +where + F: FnMut(usize) -> bool, +{ + debug_assert_eq!(len, valid_bits.len()); + + for (word_idx, valid_word) in valid_bits.chunks().iter_padded().enumerate() { + if valid_word == 0 { + continue; + } + + let offset = word_idx * 64; + let lanes = len.saturating_sub(offset).min(64); + if lanes == 64 && valid_word == u64::MAX { + for bit_idx in 0..64 { + if !f(offset + bit_idx) { + return false; + } + } + continue; + } + + let mut valid_word = if lanes == 64 { + valid_word + } else { + valid_word & ((1u64 << lanes) - 1) + }; + while valid_word != 0 { + let bit_idx = valid_word.trailing_zeros() as usize; + if !f(offset + bit_idx) { + return false; + } + valid_word &= valid_word - 1; + } + } + + true +} + +fn check_numeric_errors(failed: bool, error: &'static str) -> VortexResult<()> { + if failed { + return Err(vortex_err!(InvalidArgument: "{}", error)); + } + + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs similarity index 66% rename from vortex-array/src/scalar_fn/fns/binary/numeric.rs rename to vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs index ac76d2e4934..f1ef0b09163 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs @@ -10,6 +10,11 @@ use vortex_error::vortex_err; use vortex_mask::AllOr; use vortex_mask::Mask; +use super::CheckedValues; +use super::any_valid_error; +use super::check_numeric_errors; +use super::checked_all_lanes; +use super::checked_valid_lanes; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; @@ -104,8 +109,8 @@ impl CheckedPrimitiveOp for CheckedDiv { } } -/// Execute a numeric operation between two arrays. -pub(crate) fn execute_numeric( +/// Execute a numeric operation between two primitive-typed arrays. +pub(super) fn execute_numeric_primitive( lhs: &ArrayRef, rhs: &ArrayRef, op: NumericOperator, @@ -214,7 +219,7 @@ where /// A primitive binary-operator operand: a materialized buffer, a non-null constant, or an /// all-null constant. -pub(super) enum PrimitiveOperand { +pub(crate) enum PrimitiveOperand { Array { values: Buffer, validity: Validity, @@ -228,7 +233,7 @@ pub(super) enum PrimitiveOperand { } impl PrimitiveOperand { - pub(super) fn try_new(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + pub(crate) fn try_new(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { if let Some(constant) = array.as_opt::() { return Ok( match constant.scalar().as_primitive().try_typed_value::()? { @@ -252,14 +257,14 @@ impl PrimitiveOperand { Ok(Self::Array { values, validity }) } - pub(super) fn len(&self) -> usize { + pub(crate) fn len(&self) -> usize { match self { Self::Array { values, .. } => values.len(), Self::Constant { len, .. } | Self::Null(len) => *len, } } - pub(super) fn validity(&self) -> Validity { + pub(crate) fn validity(&self) -> Validity { match self { Self::Array { validity, .. } => validity.clone(), Self::Constant { validity, .. } => validity.clone(), @@ -268,27 +273,6 @@ impl PrimitiveOperand { } } -struct CheckedValues { - values: Buffer, - failed: bool, -} - -impl CheckedValues { - fn zeroed(len: usize) -> Self { - Self { - values: Buffer::::zeroed(len), - failed: false, - } - } - - fn failed(len: usize) -> Self { - Self { - values: Buffer::::zeroed(len), - failed: true, - } - } -} - fn checked_array_array(lhs: &[T], rhs: &[T], valid_rows: &Mask) -> CheckedValues where T: NativePType, @@ -497,133 +481,6 @@ where checked_valid_lanes(rhs.len(), valid_bits, |idx| Op::checked(lhs, rhs[idx])) } -// Checked one-pass ops delay early exit until the end of a small block. This -// keeps the loop generic while avoiding a branch-driven exit decision on every -// lane; it is deliberately independent of mask density or input length. -const CHECKED_BLOCK_LANES: usize = 16; - -fn checked_all_lanes(len: usize, mut checked_at: F) -> CheckedValues -where - T: NativePType, - F: FnMut(usize) -> Option, -{ - let mut values = BufferMut::::with_capacity(len); - let mut base = 0; - - while base + CHECKED_BLOCK_LANES <= len { - let mut block_failed = false; - for idx in base..base + CHECKED_BLOCK_LANES { - match checked_at(idx) { - Some(value) => { - // SAFETY: the buffer is allocated with capacity `len`, and - // this loop pushes at most one value for each `idx`. - unsafe { values.push_unchecked(value) }; - } - None => { - block_failed = true; - // SAFETY: the buffer is allocated with capacity `len`, and - // this loop pushes at most one value for each `idx`. - unsafe { values.push_unchecked(T::default()) }; - } - } - } - - if block_failed { - return CheckedValues::failed(len); - } - base += CHECKED_BLOCK_LANES; - } - - for idx in base..len { - let Some(value) = checked_at(idx) else { - return CheckedValues::failed(len); - }; - // SAFETY: the buffer is allocated with capacity `len`, and this loop - // pushes at most one value for each `idx`. - unsafe { values.push_unchecked(value) }; - } - - CheckedValues { - values: values.freeze(), - failed: false, - } -} - -fn checked_valid_lanes( - len: usize, - valid_bits: &BitBuffer, - mut checked_at: F, -) -> CheckedValues -where - T: NativePType, - F: FnMut(usize) -> Option, -{ - let mut values = BufferMut::::zeroed(len); - let mut failed = false; - { - let values = values.as_mut_slice(); - for_each_valid_idx(len, valid_bits, |idx| { - let Some(value) = checked_at(idx) else { - failed = true; - return false; - }; - values[idx] = value; - true - }); - } - - CheckedValues { - values: values.freeze(), - failed, - } -} - -fn any_valid_error(len: usize, valid_bits: &BitBuffer, is_error: F) -> bool -where - F: Fn(usize) -> bool, -{ - !for_each_valid_idx(len, valid_bits, |idx| !is_error(idx)) -} - -fn for_each_valid_idx(len: usize, valid_bits: &BitBuffer, mut f: F) -> bool -where - F: FnMut(usize) -> bool, -{ - debug_assert_eq!(len, valid_bits.len()); - - for (word_idx, valid_word) in valid_bits.chunks().iter_padded().enumerate() { - if valid_word == 0 { - continue; - } - - let offset = word_idx * 64; - let lanes = len.saturating_sub(offset).min(64); - if lanes == 64 && valid_word == u64::MAX { - for bit_idx in 0..64 { - if !f(offset + bit_idx) { - return false; - } - } - continue; - } - - let mut valid_word = if lanes == 64 { - valid_word - } else { - valid_word & ((1u64 << lanes) - 1) - }; - while valid_word != 0 { - let bit_idx = valid_word.trailing_zeros() as usize; - if !f(offset + bit_idx) { - return false; - } - valid_word &= valid_word - 1; - } - } - - true -} - trait CheckedArithmetic: NativePType { const DIV_CHECKS_IN_VALUE_LOOP: bool; @@ -913,224 +770,3 @@ impl_checked_signed!(i16, widening_mul: i32); impl_checked_signed!(i32, widening_mul: i64); impl_checked_signed!(i64, overflowing_mul); impl_checked_float!(f16, f32, f64); - -fn check_numeric_errors(failed: bool, error: &'static str) -> VortexResult<()> { - if failed { - return Err(vortex_err!(InvalidArgument: "{}", error)); - } - - Ok(()) -} -#[cfg(test)] -mod test { - use vortex_buffer::buffer; - use vortex_error::VortexResult; - - use crate::ArrayRef; - use crate::IntoArray; - use crate::RecursiveCanonical; - use crate::VortexSessionExecute; - use crate::array_session; - use crate::arrays::ConstantArray; - use crate::arrays::PrimitiveArray; - use crate::assert_arrays_eq; - use crate::builtins::ArrayBuiltins; - use crate::scalar::Scalar; - use crate::scalar_fn::fns::operators::Operator; - use crate::validity::Validity; - - fn sub_scalar(array: &ArrayRef, scalar: impl Into) -> VortexResult { - array - .binary( - ConstantArray::new(scalar, array.len()).into_array(), - Operator::Sub, - ) - .and_then(|a| { - a.execute::(&mut array_session().create_execution_ctx()) - }) - .map(|a| a.0.into_array()) - } - - #[test] - fn test_scalar_subtract_unsigned() { - let mut ctx = array_session().create_execution_ctx(); - let values = buffer![1u16, 2, 3].into_array(); - let result = sub_scalar(&values, 1u16).unwrap(); - assert_arrays_eq!(result, PrimitiveArray::from_iter([0u16, 1, 2]), &mut ctx); - } - - #[test] - fn test_scalar_subtract_signed() { - let mut ctx = array_session().create_execution_ctx(); - let values = buffer![1i64, 2, 3].into_array(); - let result = sub_scalar(&values, -1i64).unwrap(); - assert_arrays_eq!(result, PrimitiveArray::from_iter([2i64, 3, 4]), &mut ctx); - } - - #[test] - fn test_scalar_subtract_nullable() { - let mut ctx = array_session().create_execution_ctx(); - let values = PrimitiveArray::from_option_iter([Some(1u16), Some(2), None, Some(3)]); - let result = sub_scalar(&values.into_array(), Some(1u16)).unwrap(); - assert_arrays_eq!( - result, - PrimitiveArray::from_option_iter([Some(0u16), Some(1), None, Some(2)]), - &mut ctx - ); - } - - #[test] - fn test_scalar_subtract_float() { - let mut ctx = array_session().create_execution_ctx(); - let values = buffer![1.0f64, 2.0, 3.0].into_array(); - let result = sub_scalar(&values, -1f64).unwrap(); - assert_arrays_eq!( - result, - PrimitiveArray::from_iter([2.0f64, 3.0, 4.0]), - &mut ctx - ); - } - - #[test] - fn test_scalar_subtract_float_underflow_is_ok() { - let values = buffer![f32::MIN, 2.0, 3.0].into_array(); - let _results = sub_scalar(&values, 1.0f32).unwrap(); - let _results = sub_scalar(&values, f32::MAX).unwrap(); - } - - #[test] - fn test_float_divide_by_zero_is_ok() { - let mut ctx = array_session().create_execution_ctx(); - let values = buffer![1.0f64, -1.0].into_array(); - let result = values - .binary( - ConstantArray::new(0.0f64, values.len()).into_array(), - Operator::Div, - ) - .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())) - .unwrap(); - - assert_arrays_eq!( - result, - PrimitiveArray::from_iter([f64::INFINITY, f64::NEG_INFINITY]), - &mut ctx - ); - } - - #[test] - fn test_integer_overflow_errors() { - let values = buffer![u8::MAX].into_array(); - let result = values - .binary( - ConstantArray::new(1u8, values.len()).into_array(), - Operator::Add, - ) - .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())); - - assert!(result.is_err()); - } - - #[test] - fn test_integer_divide_by_zero_errors() { - let values = buffer![1i32].into_array(); - let result = values - .binary( - ConstantArray::new(0i32, values.len()).into_array(), - Operator::Div, - ) - .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())); - - assert!(result.is_err()); - } - - #[test] - fn test_integer_divide_overflow_errors() { - let values = buffer![i64::MIN].into_array(); - let result = values - .binary( - ConstantArray::new(-1i64, values.len()).into_array(), - Operator::Div, - ) - .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())); - - assert!(result.is_err()); - } - - #[test] - fn test_integer_divide_errors_ignore_null_lanes() { - let mut ctx = array_session().create_execution_ctx(); - let lhs = PrimitiveArray::new(buffer![10i32, 10], Validity::from_iter([false, true])) - .into_array(); - let rhs = buffer![0i32, 2].into_array(); - let result = lhs - .binary(rhs, Operator::Div) - .and_then(|a| { - a.execute::(&mut array_session().create_execution_ctx()) - }) - .map(|a| a.0.into_array()) - .unwrap(); - - assert_arrays_eq!( - result, - PrimitiveArray::from_option_iter([None, Some(5i32)]), - &mut ctx - ); - } - - #[test] - fn test_integer_errors_ignore_null_lanes() { - let mut ctx = array_session().create_execution_ctx(); - let values = PrimitiveArray::new(buffer![u8::MAX, 1], Validity::from_iter([false, true])) - .into_array(); - let result = values - .binary( - ConstantArray::new(1u8, values.len()).into_array(), - Operator::Add, - ) - .and_then(|a| { - a.execute::(&mut array_session().create_execution_ctx()) - }) - .map(|a| a.0.into_array()) - .unwrap(); - - assert_arrays_eq!( - result, - PrimitiveArray::from_option_iter([None, Some(2u8)]), - &mut ctx - ); - } - - #[test] - fn test_integer_array_array_errors_on_valid_lanes() { - let lhs = PrimitiveArray::new( - buffer![u8::MAX, 1, u8::MAX], - Validity::from_iter([false, true, true]), - ) - .into_array(); - let rhs = buffer![1u8, 1, 1].into_array(); - let result = lhs - .binary(rhs, Operator::Add) - .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())); - - assert!(result.is_err()); - } - - #[test] - fn test_present_nullable_constant_preserves_nullable_output() { - let mut ctx = array_session().create_execution_ctx(); - let values = buffer![1u8, 2].into_array(); - let result = values - .binary( - ConstantArray::new(Some(1u8), values.len()).into_array(), - Operator::Add, - ) - .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())) - .unwrap(); - - assert_arrays_eq!( - result, - PrimitiveArray::from_option_iter([Some(2u8), Some(3)]), - &mut ctx - ); - } -} diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs new file mode 100644 index 00000000000..00f613c23b6 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs @@ -0,0 +1,207 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::buffer; +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::RecursiveCanonical; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::ConstantArray; +use crate::arrays::PrimitiveArray; +use crate::assert_arrays_eq; +use crate::builtins::ArrayBuiltins; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::operators::Operator; +use crate::validity::Validity; + +fn sub_scalar(array: &ArrayRef, scalar: impl Into) -> VortexResult { + array + .binary( + ConstantArray::new(scalar, array.len()).into_array(), + Operator::Sub, + ) + .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())) + .map(|a| a.0.into_array()) +} + +#[test] +fn test_scalar_subtract_unsigned() { + let mut ctx = array_session().create_execution_ctx(); + let values = buffer![1u16, 2, 3].into_array(); + let result = sub_scalar(&values, 1u16).unwrap(); + assert_arrays_eq!(result, PrimitiveArray::from_iter([0u16, 1, 2]), &mut ctx); +} + +#[test] +fn test_scalar_subtract_signed() { + let mut ctx = array_session().create_execution_ctx(); + let values = buffer![1i64, 2, 3].into_array(); + let result = sub_scalar(&values, -1i64).unwrap(); + assert_arrays_eq!(result, PrimitiveArray::from_iter([2i64, 3, 4]), &mut ctx); +} + +#[test] +fn test_scalar_subtract_nullable() { + let mut ctx = array_session().create_execution_ctx(); + let values = PrimitiveArray::from_option_iter([Some(1u16), Some(2), None, Some(3)]); + let result = sub_scalar(&values.into_array(), Some(1u16)).unwrap(); + assert_arrays_eq!( + result, + PrimitiveArray::from_option_iter([Some(0u16), Some(1), None, Some(2)]), + &mut ctx + ); +} + +#[test] +fn test_scalar_subtract_float() { + let mut ctx = array_session().create_execution_ctx(); + let values = buffer![1.0f64, 2.0, 3.0].into_array(); + let result = sub_scalar(&values, -1f64).unwrap(); + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([2.0f64, 3.0, 4.0]), + &mut ctx + ); +} + +#[test] +fn test_scalar_subtract_float_underflow_is_ok() { + let values = buffer![f32::MIN, 2.0, 3.0].into_array(); + let _results = sub_scalar(&values, 1.0f32).unwrap(); + let _results = sub_scalar(&values, f32::MAX).unwrap(); +} + +#[test] +fn test_float_divide_by_zero_is_ok() { + let mut ctx = array_session().create_execution_ctx(); + let values = buffer![1.0f64, -1.0].into_array(); + let result = values + .binary( + ConstantArray::new(0.0f64, values.len()).into_array(), + Operator::Div, + ) + .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())) + .unwrap(); + + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([f64::INFINITY, f64::NEG_INFINITY]), + &mut ctx + ); +} + +#[test] +fn test_integer_overflow_errors() { + let values = buffer![u8::MAX].into_array(); + let result = values + .binary( + ConstantArray::new(1u8, values.len()).into_array(), + Operator::Add, + ) + .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())); + + assert!(result.is_err()); +} + +#[test] +fn test_integer_divide_by_zero_errors() { + let values = buffer![1i32].into_array(); + let result = values + .binary( + ConstantArray::new(0i32, values.len()).into_array(), + Operator::Div, + ) + .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())); + + assert!(result.is_err()); +} + +#[test] +fn test_integer_divide_overflow_errors() { + let values = buffer![i64::MIN].into_array(); + let result = values + .binary( + ConstantArray::new(-1i64, values.len()).into_array(), + Operator::Div, + ) + .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())); + + assert!(result.is_err()); +} + +#[test] +fn test_integer_divide_errors_ignore_null_lanes() { + let mut ctx = array_session().create_execution_ctx(); + let lhs = + PrimitiveArray::new(buffer![10i32, 10], Validity::from_iter([false, true])).into_array(); + let rhs = buffer![0i32, 2].into_array(); + let result = lhs + .binary(rhs, Operator::Div) + .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())) + .map(|a| a.0.into_array()) + .unwrap(); + + assert_arrays_eq!( + result, + PrimitiveArray::from_option_iter([None, Some(5i32)]), + &mut ctx + ); +} + +#[test] +fn test_integer_errors_ignore_null_lanes() { + let mut ctx = array_session().create_execution_ctx(); + let values = + PrimitiveArray::new(buffer![u8::MAX, 1], Validity::from_iter([false, true])).into_array(); + let result = values + .binary( + ConstantArray::new(1u8, values.len()).into_array(), + Operator::Add, + ) + .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())) + .map(|a| a.0.into_array()) + .unwrap(); + + assert_arrays_eq!( + result, + PrimitiveArray::from_option_iter([None, Some(2u8)]), + &mut ctx + ); +} + +#[test] +fn test_integer_array_array_errors_on_valid_lanes() { + let lhs = PrimitiveArray::new( + buffer![u8::MAX, 1, u8::MAX], + Validity::from_iter([false, true, true]), + ) + .into_array(); + let rhs = buffer![1u8, 1, 1].into_array(); + let result = lhs + .binary(rhs, Operator::Add) + .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())); + + assert!(result.is_err()); +} + +#[test] +fn test_present_nullable_constant_preserves_nullable_output() { + let mut ctx = array_session().create_execution_ctx(); + let values = buffer![1u8, 2].into_array(); + let result = values + .binary( + ConstantArray::new(Some(1u8), values.len()).into_array(), + Operator::Add, + ) + .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())) + .unwrap(); + + assert_arrays_eq!( + result, + PrimitiveArray::from_option_iter([Some(2u8), Some(3)]), + &mut ctx + ); +} From 2bc7b81f6aaa2942d385ce59571ee2dba409bdcd Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 10 Jul 2026 17:56:41 -0700 Subject: [PATCH 2/5] Add native decimal Add/Sub kernels Add and Sub over decimal arrays sharing a decimal dtype, applied to the unscaled stored integers (exact at a shared scale) in a working width wide enough that in-precision inputs cannot spuriously overflow. Precision violations error only on valid lanes, matching primitive semantics. Mul and Div need rescaling and are gated off until follow-up PRs. widened_buffer moves to arrays/decimal so compare and numeric share it. Signed-off-by: Matt Katz --- vortex-array/src/arrays/decimal/utils.rs | 19 + .../scalar_fn/fns/binary/compare/decimal.rs | 18 +- .../scalar_fn/fns/binary/compare/nested.rs | 5 +- vortex-array/src/scalar_fn/fns/binary/mod.rs | 5 +- .../scalar_fn/fns/binary/numeric/decimal.rs | 328 ++++++++++++++++++ .../src/scalar_fn/fns/binary/numeric/mod.rs | 17 +- .../scalar_fn/fns/binary/numeric/primitive.rs | 7 - .../src/scalar_fn/fns/binary/numeric/tests.rs | 224 ++++++++++++ 8 files changed, 595 insertions(+), 28 deletions(-) create mode 100644 vortex-array/src/scalar_fn/fns/binary/numeric/decimal.rs diff --git a/vortex-array/src/arrays/decimal/utils.rs b/vortex-array/src/arrays/decimal/utils.rs index a1eafde0fe3..eb307cc53a0 100644 --- a/vortex-array/src/arrays/decimal/utils.rs +++ b/vortex-array/src/arrays/decimal/utils.rs @@ -3,11 +3,30 @@ use itertools::Itertools; use itertools::MinMaxResult; +use vortex_buffer::Buffer; use vortex_error::VortexExpect; use crate::arrays::DecimalArray; +use crate::arrays::decimal::DecimalArrayExt; use crate::dtype::DecimalType; +use crate::dtype::NativeDecimalType; use crate::dtype::i256; +use crate::match_each_decimal_value_type; + +/// Return the array's unscaled values widened to `W`, which must be at least as wide as the +/// array's storage type. +pub(crate) fn widened_buffer(array: &DecimalArray) -> Buffer { + if array.values_type() == W::DECIMAL_TYPE { + return array.buffer::(); + } + match_each_decimal_value_type!(array.values_type(), |T| { + array + .buffer::() + .iter() + .map(|v| W::from(*v).vortex_expect("widening decimal cast must succeed")) + .collect() + }) +} macro_rules! try_downcast { ($array:expr, from: $src:ty, to: $($dst:ty),*) => {{ diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/decimal.rs b/vortex-array/src/scalar_fn/fns/binary/compare/decimal.rs index 0825d0674af..eea6589645e 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/decimal.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/decimal.rs @@ -10,8 +10,6 @@ //! [`DecimalDType`]: crate::dtype::DecimalDType use vortex_buffer::BitBuffer; -use vortex_buffer::Buffer; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -21,6 +19,7 @@ use crate::IntoArray; use crate::arrays::BoolArray; use crate::arrays::Constant; use crate::arrays::DecimalArray; +use crate::arrays::decimal::widened_buffer; use crate::dtype::NativeDecimalType; use crate::dtype::Nullability; use crate::dtype::i256; @@ -120,21 +119,6 @@ fn compare_decimal_values( }) } -/// Return the array's unscaled values widened to `W`, which must be at least as wide as the -/// array's storage type. -pub(super) fn widened_buffer(array: &DecimalArray) -> Buffer { - if array.values_type() == W::DECIMAL_TYPE { - return array.buffer::(); - } - match_each_decimal_value_type!(array.values_type(), |T| { - array - .buffer::() - .iter() - .map(|v| W::from(*v).vortex_expect("widening decimal cast must succeed")) - .collect() - }) -} - fn compare_decimal_constant( array: &DecimalArray, constant: DecimalValue, diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs b/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs index cbc38d6b23e..05e95c03c57 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs @@ -31,6 +31,7 @@ use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::VarBinViewArray; +use crate::arrays::decimal::widened_buffer; use crate::arrays::extension::ExtensionArrayExt; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::arrays::listview::ListViewArrayExt; @@ -151,8 +152,8 @@ fn build_values_comparator( let rhs = rhs.clone().execute::(ctx)?; let common = lhs.values_type().max(rhs.values_type()); crate::match_each_decimal_value_type!(common, |W| { - let lhs = super::decimal::widened_buffer::(&lhs); - let rhs = super::decimal::widened_buffer::(&rhs); + let lhs = widened_buffer::(&lhs); + let rhs = widened_buffer::(&rhs); Box::new(move |i: usize, j: usize| lhs[i].cmp(&rhs[j])) as RowComparator }) } diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index 361c4fab240..7532801853c 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -118,7 +118,10 @@ impl ScalarFnVTable for Binary { let rhs = &arg_dtypes[1]; if operator.is_arithmetic() { - if lhs.is_primitive() && lhs.eq_ignore_nullability(rhs) { + // Decimal Mul/Div need rescaling support and are not yet implemented. + let decimal_supported = + lhs.is_decimal() && matches!(operator, Operator::Add | Operator::Sub); + if (lhs.is_primitive() || decimal_supported) && lhs.eq_ignore_nullability(rhs) { return Ok(lhs.with_nullability(lhs.nullability() | rhs.nullability())); } vortex_bail!( diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/decimal.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/decimal.rs new file mode 100644 index 00000000000..b95c557048f --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/decimal.rs @@ -0,0 +1,328 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Native execution of the arithmetic operators over decimal arrays. +//! +//! Both operands share a logical [`DecimalDType`] (equal precision and scale) and the result +//! keeps that dtype: fixed-point arithmetic at the shared scale. Add and Sub apply directly to +//! the unscaled stored integers and are exact; Mul and Div require rescaling and are not yet +//! implemented. +//! +//! 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. + +use num_traits::CheckedAdd; +use num_traits::CheckedSub; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_mask::AllOr; +use vortex_mask::Mask; + +use super::CheckedValues; +use super::check_numeric_errors; +use super::checked_all_lanes; +use super::checked_valid_lanes; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::Constant; +use crate::arrays::ConstantArray; +use crate::arrays::DecimalArray; +use crate::arrays::decimal::DecimalArrayExt; +use crate::arrays::decimal::widened_buffer; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::DecimalType; +use crate::dtype::NativeDecimalType; +use crate::dtype::i256; +use crate::match_each_decimal_value_type; +use crate::scalar::DecimalValue; +use crate::scalar::NumericOperator; +use crate::scalar::Scalar; +use crate::validity::Validity; + +/// Execute a numeric operation between two decimal arrays sharing a decimal dtype. +pub(super) fn execute_numeric_decimal( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: NumericOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let DType::Decimal(decimal_dtype, _) = lhs.dtype() else { + vortex_bail!("expected a decimal dtype, got {}", lhs.dtype()); + }; + let decimal_dtype = *decimal_dtype; + let result_dtype = lhs + .dtype() + .with_nullability(lhs.dtype().nullability() | rhs.dtype().nullability()); + + let lhs = DecimalOperand::try_new(lhs, ctx)?; + let rhs = DecimalOperand::try_new(rhs, ctx)?; + let len = lhs.len(); + if len != rhs.len() { + vortex_bail!( + "numeric operator requires equal lengths, got {} and {}", + len, + rhs.len() + ); + } + + let validity = lhs.validity().and(rhs.validity())?; + let valid_rows = validity.execute_mask(len, ctx)?; + + let work = working_type(&lhs, &rhs, decimal_dtype); + match_each_decimal_value_type!(work, |W| { + match op { + NumericOperator::Add => execute_decimal_typed::( + &lhs, + &rhs, + decimal_dtype, + &result_dtype, + validity, + &valid_rows, + ), + NumericOperator::Sub => execute_decimal_typed::( + &lhs, + &rhs, + decimal_dtype, + &result_dtype, + validity, + &valid_rows, + ), + NumericOperator::Mul | NumericOperator::Div => vortex_bail!( + "numeric operator {:?} is not yet supported for decimal arrays", + op + ), + } + }) +} + +/// A decimal binary-operator operand: a canonical decimal array, a non-null constant, or an +/// all-null constant. +enum DecimalOperand { + Array { + values: DecimalArray, + validity: Validity, + }, + Constant { + value: DecimalValue, + len: usize, + validity: Validity, + }, + Null(usize), +} + +impl DecimalOperand { + fn try_new(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + if let Some(constant) = array.as_opt::() { + return Ok(match constant.scalar().as_decimal().decimal_value() { + Some(value) => Self::Constant { + value, + len: array.len(), + validity: if constant.scalar().dtype().is_nullable() { + Validity::AllValid + } else { + Validity::NonNullable + }, + }, + None => Self::Null(array.len()), + }); + } + + let values = array.clone().execute::(ctx)?; + let validity = values.validity()?; + Ok(Self::Array { values, validity }) + } + + fn len(&self) -> usize { + match self { + Self::Array { values, .. } => values.len(), + Self::Constant { len, .. } | Self::Null(len) => *len, + } + } + + fn validity(&self) -> Validity { + match self { + Self::Array { validity, .. } | Self::Constant { validity, .. } => validity.clone(), + Self::Null(_) => Validity::AllInvalid, + } + } +} + +/// Choose the lane width: wide enough for the operand storage and for every intermediate value +/// that in-precision inputs can produce, capped at 256 bits. +fn working_type(lhs: &DecimalOperand, rhs: &DecimalOperand, dtype: DecimalDType) -> DecimalType { + // The sum or difference of two p-digit values has at most p + 1 digits. + let needed_digits = (dtype.precision() + 1).min(::MAX_PRECISION); + let needed = DecimalType::smallest_decimal_value_type(&DecimalDType::new(needed_digits, 0)); + + let operand_type = |operand: &DecimalOperand| match operand { + DecimalOperand::Array { values, .. } => values.values_type(), + DecimalOperand::Constant { value, .. } => smallest_value_type(value), + DecimalOperand::Null(_) => DecimalType::I8, + }; + + needed.max(operand_type(lhs)).max(operand_type(rhs)) +} + +/// The smallest decimal value type that can represent `value`, regardless of its stored width. +fn smallest_value_type(value: &DecimalValue) -> DecimalType { + if value.cast::().is_some() { + DecimalType::I8 + } else if value.cast::().is_some() { + DecimalType::I16 + } else if value.cast::().is_some() { + DecimalType::I32 + } else if value.cast::().is_some() { + DecimalType::I64 + } else if value.cast::().is_some() { + DecimalType::I128 + } else { + DecimalType::I256 + } +} + +/// Per-execution constants for checked decimal lane operations at working width `W`. +struct DecimalOpPlan { + /// Inclusive stored-value bounds implied by the result precision. + prec_min: W, + prec_max: W, +} + +impl DecimalOpPlan { + fn new(dtype: DecimalDType) -> Self { + let precision = dtype.precision() as usize; + Self { + prec_min: W::MIN_BY_PRECISION[precision], + prec_max: W::MAX_BY_PRECISION[precision], + } + } + + /// Bounds-check a candidate result against the result precision. + #[inline(always)] + fn in_precision(&self, value: W) -> Option { + (self.prec_min <= value && value <= self.prec_max).then_some(value) + } +} + +/// A checked fixed-point decimal operation on unscaled values at working width `W`. +trait CheckedDecimalOp { + const ERROR: &'static str; + + fn checked(lhs: W, rhs: W, plan: &DecimalOpPlan) -> Option + where + W: NativeDecimalType + CheckedAdd + CheckedSub; +} + +struct DecimalAdd; + +struct DecimalSub; + +impl CheckedDecimalOp for DecimalAdd { + const ERROR: &'static str = "decimal overflow in checked add"; + + #[inline(always)] + fn checked(lhs: W, rhs: W, plan: &DecimalOpPlan) -> Option + where + W: NativeDecimalType + CheckedAdd + CheckedSub, + { + plan.in_precision(lhs.checked_add(&rhs)?) + } +} + +impl CheckedDecimalOp for DecimalSub { + const ERROR: &'static str = "decimal overflow in checked sub"; + + #[inline(always)] + fn checked(lhs: W, rhs: W, plan: &DecimalOpPlan) -> Option + where + W: NativeDecimalType + CheckedAdd + CheckedSub, + { + plan.in_precision(lhs.checked_sub(&rhs)?) + } +} + +fn execute_decimal_typed( + lhs: &DecimalOperand, + rhs: &DecimalOperand, + decimal_dtype: DecimalDType, + result_dtype: &DType, + validity: Validity, + valid_rows: &Mask, +) -> VortexResult +where + W: NativeDecimalType + CheckedAdd + CheckedSub, + DecimalValue: From, + Op: CheckedDecimalOp, +{ + let len = lhs.len(); + let plan = DecimalOpPlan::::new(decimal_dtype); + + let checked = match (lhs, rhs) { + (DecimalOperand::Array { values: lhs, .. }, DecimalOperand::Array { values: rhs, .. }) => { + let lhs = widened_buffer::(lhs); + let rhs = widened_buffer::(rhs); + checked_decimal_lanes(len, valid_rows, |idx| { + Op::checked(lhs[idx], rhs[idx], &plan) + }) + } + (DecimalOperand::Array { values: lhs, .. }, DecimalOperand::Constant { value, .. }) => { + let lhs = widened_buffer::(lhs); + let rhs = typed_constant::(value); + checked_decimal_lanes(len, valid_rows, |idx| Op::checked(lhs[idx], rhs, &plan)) + } + (DecimalOperand::Constant { value, .. }, DecimalOperand::Array { values: rhs, .. }) => { + let lhs = typed_constant::(value); + let rhs = widened_buffer::(rhs); + checked_decimal_lanes(len, valid_rows, |idx| Op::checked(lhs, rhs[idx], &plan)) + } + ( + DecimalOperand::Constant { value: lhs, .. }, + DecimalOperand::Constant { value: rhs, .. }, + ) => { + let lhs = typed_constant::(lhs); + let rhs = typed_constant::(rhs); + let value = Op::checked(lhs, rhs, &plan) + .ok_or_else(|| vortex_err!(InvalidArgument: "{}", Op::ERROR))?; + return Ok(ConstantArray::new( + Scalar::decimal( + DecimalValue::from(value), + decimal_dtype, + result_dtype.nullability(), + ), + len, + ) + .into_array()); + } + (DecimalOperand::Null(_), _) | (_, DecimalOperand::Null(_)) => CheckedValues::zeroed(len), + }; + check_numeric_errors(checked.failed, Op::ERROR)?; + + Ok(DecimalArray::new( + checked.values, + decimal_dtype, + validity.union_nullability(result_dtype.nullability()), + ) + .into_array()) +} + +fn checked_decimal_lanes(len: usize, valid_rows: &Mask, checked_at: F) -> CheckedValues +where + W: NativeDecimalType, + F: FnMut(usize) -> Option, +{ + match valid_rows.bit_buffer() { + AllOr::All => checked_all_lanes(len, checked_at), + AllOr::None => CheckedValues::zeroed(len), + AllOr::Some(valid_bits) => checked_valid_lanes(len, valid_bits, checked_at), + } +} + +fn typed_constant(value: &DecimalValue) -> W { + value + .cast::() + .vortex_expect("the working width must be able to represent the constant") +} 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 702ec47adf0..b0b7d754e81 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs @@ -6,6 +6,7 @@ //! //! [`Binary`]: super::Binary +mod decimal; mod primitive; #[cfg(test)] mod tests; @@ -15,10 +16,12 @@ use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_err; use crate::ArrayRef; use crate::ExecutionCtx; +use crate::dtype::DType; use crate::scalar::NumericOperator; /// Execute a numeric operation between two arrays. @@ -28,7 +31,19 @@ pub(crate) fn execute_numeric( op: NumericOperator, ctx: &mut ExecutionCtx, ) -> VortexResult { - primitive::execute_numeric_primitive(lhs, rhs, op, ctx) + if !lhs.dtype().eq_ignore_nullability(rhs.dtype()) { + vortex_bail!( + "numeric operator requires matching types, got {} and {}", + lhs.dtype(), + rhs.dtype() + ); + } + + match lhs.dtype() { + DType::Primitive(..) => primitive::execute_numeric_primitive(lhs, rhs, op, ctx), + DType::Decimal(..) => decimal::execute_numeric_decimal(lhs, rhs, op, ctx), + dtype => vortex_bail!("numeric operator is not supported for dtype {}", dtype), + } } /// The values produced by a checked lane loop, plus whether any lane failed. diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs index f1ef0b09163..327f078bba0 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs @@ -117,13 +117,6 @@ pub(super) fn execute_numeric_primitive( ctx: &mut ExecutionCtx, ) -> VortexResult { let ptype = PType::try_from(lhs.dtype())?; - if !lhs.dtype().eq_ignore_nullability(rhs.dtype()) { - vortex_bail!( - "numeric operator requires matching primitive types, got {} and {}", - lhs.dtype(), - rhs.dtype() - ); - } match_each_native_ptype!(ptype, |T| { match op { 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 00f613c23b6..fedcf0f420c 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use rstest::rstest; use vortex_buffer::buffer; use vortex_error::VortexResult; @@ -10,9 +11,14 @@ use crate::RecursiveCanonical; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::ConstantArray; +use crate::arrays::DecimalArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::Nullability; +use crate::scalar::DecimalValue; use crate::scalar::Scalar; use crate::scalar_fn::fns::operators::Operator; use crate::validity::Validity; @@ -205,3 +211,221 @@ fn test_present_nullable_constant_preserves_nullable_output() { &mut ctx ); } + +// -- Decimal arithmetic -- + +fn decimal_binary(lhs: ArrayRef, rhs: ArrayRef, op: Operator) -> VortexResult { + lhs.binary(rhs, op) + .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())) + .map(|a| a.0.into_array()) +} + +fn decimal_constant(value: impl Into, dtype: DecimalDType, len: usize) -> ArrayRef { + ConstantArray::new( + Scalar::decimal(value.into(), dtype, Nullability::NonNullable), + len, + ) + .into_array() +} + +#[rstest] +#[case::add(Operator::Add, [150i64, 225], [1050i64, 1225])] +#[case::sub(Operator::Sub, [150i64, 225], [750i64, 775])] +fn test_decimal_array_array( + #[case] op: Operator, + #[case] rhs: [i64; 2], + #[case] expected: [i64; 2], +) { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DecimalDType::new(10, 2); + let lhs = DecimalArray::from_iter::([900, 1000], dtype).into_array(); + let rhs = DecimalArray::from_iter::(rhs, dtype).into_array(); + + let result = decimal_binary(lhs, rhs, op).unwrap(); + assert_arrays_eq!( + result, + DecimalArray::from_iter::(expected, dtype), + &mut ctx + ); +} + +#[test] +fn test_decimal_mixed_storage_widths() { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DecimalDType::new(10, 2); + let lhs = DecimalArray::from_iter::([100, 250], dtype).into_array(); + let rhs = DecimalArray::from_iter::([200, 250], dtype).into_array(); + + let result = decimal_binary(lhs, rhs, Operator::Add).unwrap(); + assert_arrays_eq!( + result, + DecimalArray::from_iter::([300, 500], dtype), + &mut ctx + ); +} + +#[test] +fn test_decimal_nullable_lanes() { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DecimalDType::new(10, 2); + let lhs = + DecimalArray::from_option_iter::([Some(100), None, Some(300)], dtype).into_array(); + let rhs = DecimalArray::from_iter::([50, 50, 50], dtype).into_array(); + + let result = decimal_binary(lhs, rhs, Operator::Add).unwrap(); + assert_arrays_eq!( + result, + DecimalArray::from_option_iter::([Some(150), None, Some(350)], dtype), + &mut ctx + ); +} + +#[test] +fn test_decimal_overflow_on_valid_lane_errors() { + let dtype = DecimalDType::new(3, 0); + let lhs = DecimalArray::from_iter::([999], dtype).into_array(); + let rhs = DecimalArray::from_iter::([2], dtype).into_array(); + + assert!(decimal_binary(lhs, rhs, Operator::Add).is_err()); +} + +#[test] +fn test_decimal_overflow_on_null_lane_ignored() { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DecimalDType::new(3, 0); + // The null lane holds 999, so adding 500 overflows the precision there but is ignored. + let lhs = DecimalArray::new( + buffer![999i16, 1], + dtype, + Validity::from_iter([false, true]), + ) + .into_array(); + let rhs = decimal_constant(500i16, dtype, 2); + + let result = decimal_binary(lhs, rhs, Operator::Add).unwrap(); + assert_arrays_eq!( + result, + DecimalArray::from_option_iter::([None, Some(501)], dtype), + &mut ctx + ); +} + +/// A value can fit the storage width while violating the dtype precision: 60 + 60 fits an i8 but +/// exceeds precision 2. +#[test] +fn test_decimal_precision_stricter_than_width() { + let dtype = DecimalDType::new(2, 0); + let lhs = DecimalArray::from_iter::([60], dtype).into_array(); + let rhs = DecimalArray::from_iter::([60], dtype).into_array(); + + assert!(decimal_binary(lhs, rhs, Operator::Add).is_err()); +} + +#[test] +fn test_decimal_constant_lhs_non_commutative() { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DecimalDType::new(10, 2); + let lhs = decimal_constant(1000i64, dtype, 2); + let rhs = DecimalArray::from_iter::([250, 400], dtype).into_array(); + + let result = decimal_binary(lhs, rhs, Operator::Sub).unwrap(); + assert_arrays_eq!( + result, + DecimalArray::from_iter::([750, 600], dtype), + &mut ctx + ); +} + +#[test] +fn test_decimal_nullable_constant_preserves_nullable_output() { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DecimalDType::new(10, 2); + let values = DecimalArray::from_iter::([100, 200], dtype).into_array(); + let constant = ConstantArray::new( + Scalar::decimal(DecimalValue::from(50i64), dtype, Nullability::Nullable), + 2, + ) + .into_array(); + + let result = decimal_binary(values, constant, Operator::Add).unwrap(); + assert_arrays_eq!( + result, + DecimalArray::from_option_iter::([Some(150), Some(250)], dtype), + &mut ctx + ); +} + +#[test] +fn test_decimal_null_constant_yields_all_null() { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DecimalDType::new(10, 2); + let values = DecimalArray::from_iter::([100, 200], dtype).into_array(); + let null_constant = ConstantArray::new( + Scalar::null(DType::Decimal(dtype, Nullability::Nullable)), + 2, + ) + .into_array(); + + let result = decimal_binary(values, null_constant, Operator::Add).unwrap(); + assert_arrays_eq!( + result, + DecimalArray::from_option_iter::([None, None], dtype), + &mut ctx + ); +} + +/// A constant stored in a wider variant than the array storage participates through the widened +/// working type. +#[test] +fn test_decimal_constant_wider_than_array_storage() { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DecimalDType::new(20, 0); + let values = DecimalArray::from_iter::([1, 2], dtype).into_array(); + let constant = decimal_constant(10_000_000_000i64, dtype, 2); + + let result = decimal_binary(values, constant, Operator::Add).unwrap(); + assert_arrays_eq!( + result, + DecimalArray::from_iter::([10_000_000_001, 10_000_000_002], dtype), + &mut ctx + ); +} + +#[test] +fn test_decimal_empty() { + 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).unwrap(); + assert_arrays_eq!( + result, + DecimalArray::from_iter::([], dtype), + &mut ctx + ); +} + +#[test] +fn test_decimal_constant_constant_folds() { + 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).unwrap(); + assert_arrays_eq!( + result, + DecimalArray::from_iter::([200, 200, 200], dtype), + &mut ctx + ); +} + +#[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(); + + assert!(decimal_binary(values.clone(), values, op).is_err()); +} From e3fdf45335adea4b4fc80db1f23d6aa279198e95 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 10 Jul 2026 17:58:02 -0700 Subject: [PATCH 3/5] Extend binary numeric conformance and benches to decimal Add/Sub The conformance harness now accepts decimal arrays, checking Add/Sub results element-wise against DecimalScalar::checked_binary_numeric for representative constants. Wire it into the decimal, chunked, and decimal-byte-parts compute tests, and add decimal Add cases to the binary_ops benchmark. Signed-off-by: Matt Katz --- .../src/decimal_byte_parts/compute/mod.rs | 19 +++ vortex-array/benches/binary_ops.rs | 33 ++++++ .../src/arrays/chunked/compute/mod.rs | 9 ++ .../src/arrays/decimal/compute/mod.rs | 49 ++++++++ .../src/compute/conformance/binary_numeric.rs | 112 +++++++++++++++++- 5 files changed, 221 insertions(+), 1 deletion(-) diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs index 0a5477e93cf..6c2d0dabb31 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs @@ -16,6 +16,7 @@ mod tests { use vortex_array::VortexSessionExecute; use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; + use vortex_array::compute::conformance::binary_numeric::test_binary_numeric_array; use vortex_array::compute::conformance::consistency::test_array_consistency; use vortex_array::dtype::DecimalDType; use vortex_buffer::buffer; @@ -74,4 +75,22 @@ mod tests { let ctx = &mut array_session().create_execution_ctx(); test_array_consistency(&array.into_array(), ctx); } + + #[rstest] + #[case::decimal_i32(DecimalByteParts::try_new( + buffer![100i32, 200, 300, 400, 500].into_array(), + DecimalDType::new(10, 2) + ).unwrap())] + #[case::decimal_nullable_i64(DecimalByteParts::try_new( + PrimitiveArray::from_option_iter([Some(1000i64), None, Some(3000), Some(4000), None]).into_array(), + DecimalDType::new(19, 4) + ).unwrap())] + #[case::decimal_negative(DecimalByteParts::try_new( + buffer![-100i32, -200, 300, -400, 500].into_array(), + DecimalDType::new(10, 2) + ).unwrap())] + fn test_decimal_byte_parts_binary_numeric(#[case] array: DecimalBytePartsArray) { + let ctx = &mut array_session().create_execution_ctx(); + test_binary_numeric_array(&array.into_array(), ctx); + } } diff --git a/vortex-array/benches/binary_ops.rs b/vortex-array/benches/binary_ops.rs index d1f664f8645..63c36a72985 100644 --- a/vortex-array/benches/binary_ops.rs +++ b/vortex-array/benches/binary_ops.rs @@ -18,8 +18,10 @@ use vortex_array::VortexSessionExecute; use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DecimalDType; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_session::VortexSession; @@ -136,6 +138,22 @@ fn sub_i64_constant(bencher: Bencher) { bench_primitive(bencher, lhs, rhs, Operator::Sub); } +#[divan::bench] +fn add_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::Add); +} + +#[divan::bench] +fn add_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::Add); +} + #[divan::bench] fn eq_i64_constant(bencher: Bencher) { let lhs = primitive_nonnull(0).into_array(); @@ -172,6 +190,10 @@ fn bench_primitive(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, operator: Ope bench_binary::(bencher, lhs, rhs, operator); } +fn bench_decimal(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, operator: Operator) { + bench_binary::(bencher, lhs, rhs, operator); +} + fn bench_bool(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, operator: Operator) { bench_binary::(bencher, lhs, rhs, operator); } @@ -197,6 +219,17 @@ fn primitive_nonnull(base: i64) -> PrimitiveArray { PrimitiveArray::from_iter((0..LEN as i64).map(|i| base + i)) } +fn decimal_i64_nonnull(base: i64) -> DecimalArray { + DecimalArray::from_iter::((0..LEN as i64).map(|i| base + i), DecimalDType::new(18, 2)) +} + +fn decimal_i128_nullable(base: i128, null_every: usize) -> DecimalArray { + DecimalArray::from_option_iter::( + (0..LEN as i128).map(|i| (!(i as usize).is_multiple_of(null_every)).then_some(base + i)), + DecimalDType::new(38, 2), + ) +} + fn primitive_small_nonnull(offset: i64) -> PrimitiveArray { PrimitiveArray::from_iter((0..LEN as i64).map(|i| ((i + offset) % 1024) + 1)) } diff --git a/vortex-array/src/arrays/chunked/compute/mod.rs b/vortex-array/src/arrays/chunked/compute/mod.rs index f7c9a4e9d61..4bcaa56338e 100644 --- a/vortex-array/src/arrays/chunked/compute/mod.rs +++ b/vortex-array/src/arrays/chunked/compute/mod.rs @@ -21,10 +21,12 @@ mod tests { use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::ChunkedArray; + use crate::arrays::DecimalArray; use crate::arrays::PrimitiveArray; use crate::compute::conformance::binary_numeric::test_binary_numeric_array; use crate::compute::conformance::consistency::test_array_consistency; use crate::dtype::DType; + use crate::dtype::DecimalDType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -158,6 +160,13 @@ mod tests { ], DType::Primitive(PType::I32, Nullability::NonNullable), ).unwrap())] + #[case::chunked_decimal(ChunkedArray::try_new( + vec![ + DecimalArray::from_iter::([100, 250], DecimalDType::new(10, 2)).into_array(), + DecimalArray::from_iter::([300, 400, 500], DecimalDType::new(10, 2)).into_array(), + ], + DType::Decimal(DecimalDType::new(10, 2), Nullability::NonNullable), + ).unwrap())] fn test_chunked_binary_numeric(#[case] array: ChunkedArray) { test_binary_numeric_array( &array.into_array(), diff --git a/vortex-array/src/arrays/decimal/compute/mod.rs b/vortex-array/src/arrays/decimal/compute/mod.rs index b2f21fa8398..468f6aa7910 100644 --- a/vortex-array/src/arrays/decimal/compute/mod.rs +++ b/vortex-array/src/arrays/decimal/compute/mod.rs @@ -17,6 +17,7 @@ mod tests { use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::DecimalArray; + use crate::compute::conformance::binary_numeric::test_binary_numeric_array; use crate::compute::conformance::consistency::test_array_consistency; use crate::dtype::DecimalDType; use crate::validity::Validity; @@ -60,4 +61,52 @@ mod tests { &mut array_session().create_execution_ctx(), ); } + + #[rstest] + #[case::decimal_array(DecimalArray::new( + buffer![100i128, 200i128, 300i128, 400i128, 500i128], + DecimalDType::new(19, 2), + Validity::NonNullable, + ))] + #[case::decimal_nullable(DecimalArray::new( + buffer![1000i128, 2000i128, 3000i128, 4000i128, 5000i128], + DecimalDType::new(19, 3), + Validity::from_iter([true, false, true, true, false]), + ))] + #[case::decimal_small_precision(DecimalArray::new( + buffer![10i128, 20i128, 30i128], + DecimalDType::new(5, 1), + Validity::NonNullable, + ))] + #[case::decimal_narrow_storage(DecimalArray::new( + buffer![10i8, 20i8, 30i8], + DecimalDType::new(5, 1), + Validity::NonNullable, + ))] + #[case::decimal_single(DecimalArray::new( + buffer![42i128], + DecimalDType::new(10, 0), + Validity::NonNullable, + ))] + #[case::decimal_large_scale(DecimalArray::new( + buffer![123456789012345i128, 987654321098765i128], + DecimalDType::new(20, 10), + Validity::NonNullable, + ))] + #[case::decimal_negative(DecimalArray::new( + buffer![-100i128, -200i128, 300i128, -400i128], + DecimalDType::new(10, 2), + Validity::NonNullable, + ))] + #[case::decimal_negative_scale(DecimalArray::new( + buffer![5i64, -3i64, 600i64], + DecimalDType::new(10, -2), + Validity::NonNullable, + ))] + fn test_decimal_binary_numeric(#[case] array: DecimalArray) { + test_binary_numeric_array( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); + } } diff --git a/vortex-array/src/compute/conformance/binary_numeric.rs b/vortex-array/src/compute/conformance/binary_numeric.rs index 1fe409bf648..c909771f948 100644 --- a/vortex-array/src/compute/conformance/binary_numeric.rs +++ b/vortex-array/src/compute/conformance/binary_numeric.rs @@ -40,8 +40,12 @@ use crate::RecursiveCanonical; use crate::arrays::ConstantArray; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::NativeDecimalType; use crate::dtype::NativePType; use crate::dtype::PType; +use crate::dtype::i256; +use crate::scalar::DecimalValue; use crate::scalar::NumericOperator; use crate::scalar::PrimitiveScalar; use crate::scalar::Scalar; @@ -240,12 +244,118 @@ pub fn test_binary_numeric_array(array: &ArrayRef, ctx: &mut ExecutionCtx) { PType::F32 => test_binary_numeric_conformance::(array, ctx), PType::F64 => test_binary_numeric_conformance::(array, ctx), }, + DType::Decimal(decimal_dtype, _) => { + test_binary_numeric_conformance_decimal(array, *decimal_dtype, ctx) + } dtype => vortex_panic!( - "Binary numeric tests are only supported for primitive numeric types, got {dtype}", + "Binary numeric tests are only supported for primitive and decimal types, got {dtype}", ), } } +/// Tests binary numeric operations on a decimal array against the decimal scalar +/// implementation, using a set of representative constants: stored `0`, `1`, `-1`, the decimal +/// `1.0` (when it fits the precision), and the largest stored value for the precision. +fn test_binary_numeric_conformance_decimal( + array: &ArrayRef, + decimal_dtype: DecimalDType, + ctx: &mut ExecutionCtx, +) { + let precision = decimal_dtype.precision() as usize; + let prec_max = ::MAX_BY_PRECISION[precision]; + + let mut constants = vec![ + i256::from_i128(0), + i256::from_i128(1), + i256::from_i128(-1), + prec_max, + ]; + // The decimal value 1.0 (stored 10^s), when representable within the precision. + if decimal_dtype.scale() >= 0 + && let Some(one) = i256::from_i128(10).checked_pow(decimal_dtype.scale() as u32) + && one <= prec_max + { + constants.push(one); + } + + for constant in constants { + let value = DecimalValue::try_from_i256(constant, decimal_dtype) + .vortex_expect("conformance constants fit the precision"); + test_decimal_binary_numeric_with_scalar(array, value, decimal_dtype, ctx); + } +} + +fn test_decimal_binary_numeric_with_scalar( + array: &ArrayRef, + value: DecimalValue, + decimal_dtype: DecimalDType, + ctx: &mut ExecutionCtx, +) { + let canonicalized_array = array + .clone() + .execute::(ctx) + .vortex_expect("Must be able to canonicalise") + .into_array(); + let original_values = to_vec_of_scalar(&canonicalized_array, ctx); + + let scalar = Scalar::decimal(value, decimal_dtype, array.dtype().nullability()); + + // Decimal Mul/Div are not yet implemented. + let operators = vec![NumericOperator::Add, NumericOperator::Sub]; + + for operator in operators { + let rhs_const = ConstantArray::new(scalar.clone(), array.len()).into_array(); + + for lhs_is_array in [true, false] { + let (lhs, rhs) = if lhs_is_array { + (array.clone(), rhs_const.clone()) + } else { + (rhs_const.clone(), array.clone()) + }; + + let result = lhs + .binary(rhs, operator.into()) + .vortex_expect("apply shouldn't fail") + .execute::(ctx) + .map(|c| c.0.into_array()); + + // Skip this operator if the entire operation fails (e.g. an overflowing lane). + let Ok(result) = result else { + continue; + }; + + let actual_values = to_vec_of_scalar(&result, ctx); + let expected_results: Vec> = original_values + .iter() + .map(|x| { + let (lhs, rhs) = if lhs_is_array { + (x.as_decimal(), scalar.as_decimal()) + } else { + (scalar.as_decimal(), x.as_decimal()) + }; + lhs.checked_binary_numeric(&rhs, operator).map(Scalar::from) + }) + .collect(); + + // For elements that didn't overflow, check they match. + for (idx, (actual, expected)) in actual_values.iter().zip(&expected_results).enumerate() + { + if let Some(expected_value) = expected { + assert_eq!( + actual, + expected_value, + "Decimal binary numeric operation failed for encoding {} at index {}: \ + ({array:?})[{idx}] {operator:?} {scalar} (lhs_is_array: {lhs_is_array}) \ + expected {expected_value:?}, got {actual:?}", + array.encoding_id(), + idx, + ); + } + } + } + } +} + /// Tests binary numeric operations with edge case scalar values. /// /// This function tests operations with scalar values: From 1345b5d269ac03d4ff1e699ed4bdac72969e5a2a Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 14 Jul 2026 15:53:52 -0700 Subject: [PATCH 4/5] Improve decimal arithmetic execution Signed-off-by: Matt Katz --- .../src/compute/conformance/binary_numeric.rs | 62 +++-- vortex-array/src/expr/transform/coerce.rs | 47 ++++ .../scalar_fn/fns/binary/numeric/decimal.rs | 259 +++++++++++++----- .../src/scalar_fn/fns/binary/numeric/mod.rs | 17 ++ .../scalar_fn/fns/binary/numeric/primitive.rs | 9 +- .../src/scalar_fn/fns/binary/numeric/tests.rs | 54 ++++ 6 files changed, 353 insertions(+), 95 deletions(-) diff --git a/vortex-array/src/compute/conformance/binary_numeric.rs b/vortex-array/src/compute/conformance/binary_numeric.rs index c909771f948..8899ad89b28 100644 --- a/vortex-array/src/compute/conformance/binary_numeric.rs +++ b/vortex-array/src/compute/conformance/binary_numeric.rs @@ -313,18 +313,6 @@ fn test_decimal_binary_numeric_with_scalar( (rhs_const.clone(), array.clone()) }; - let result = lhs - .binary(rhs, operator.into()) - .vortex_expect("apply shouldn't fail") - .execute::(ctx) - .map(|c| c.0.into_array()); - - // Skip this operator if the entire operation fails (e.g. an overflowing lane). - let Ok(result) = result else { - continue; - }; - - let actual_values = to_vec_of_scalar(&result, ctx); let expected_results: Vec> = original_values .iter() .map(|x| { @@ -337,20 +325,46 @@ fn test_decimal_binary_numeric_with_scalar( }) .collect(); - // For elements that didn't overflow, check they match. + let result = lhs + .binary(rhs, operator.into()) + .vortex_expect("apply shouldn't fail") + .execute::(ctx) + .map(|c| c.0.into_array()); + + let expects_overflow = expected_results.iter().any(Option::is_none); + if expects_overflow { + assert!( + result.is_err(), + "Decimal binary numeric operation should overflow for encoding {}: \ + {operator:?} {scalar} (lhs_is_array: {lhs_is_array})", + array.encoding_id(), + ); + continue; + } + + let result = result.unwrap_or_else(|err| { + vortex_panic!( + "Decimal binary numeric operation unexpectedly failed for encoding {}: \ + {operator:?} {scalar} (lhs_is_array: {lhs_is_array}): {err}", + array.encoding_id(), + ) + }); + + let actual_values = to_vec_of_scalar(&result, ctx); for (idx, (actual, expected)) in actual_values.iter().zip(&expected_results).enumerate() { - if let Some(expected_value) = expected { - assert_eq!( - actual, - expected_value, - "Decimal binary numeric operation failed for encoding {} at index {}: \ - ({array:?})[{idx}] {operator:?} {scalar} (lhs_is_array: {lhs_is_array}) \ - expected {expected_value:?}, got {actual:?}", - array.encoding_id(), - idx, - ); - } + let expected_value = expected + .as_ref() + .vortex_expect("non-overflowing decimal lane must have an expected value"); + assert_eq!( + actual, + expected_value, + "Decimal binary numeric operation failed for encoding {} at index {}: \ + ({array:?})[{idx}] {operator:?} {scalar} (lhs_is_array: {lhs_is_array}) \ + expected {expected_value:?}, got {actual:?}", + array.encoding_id(), + idx, + ); } } } diff --git a/vortex-array/src/expr/transform/coerce.rs b/vortex-array/src/expr/transform/coerce.rs index 1b6e9acd661..7d431163959 100644 --- a/vortex-array/src/expr/transform/coerce.rs +++ b/vortex-array/src/expr/transform/coerce.rs @@ -70,6 +70,7 @@ mod tests { use vortex_error::VortexResult; use crate::dtype::DType; + use crate::dtype::DecimalDType; use crate::dtype::Nullability::NonNullable; use crate::dtype::PType; use crate::dtype::StructFields; @@ -153,6 +154,52 @@ mod tests { Ok(()) } + #[test] + fn decimal_arithmetic_coerces_precision_and_scale() -> VortexResult<()> { + let result_dtype = DType::Decimal(DecimalDType::new(4, 2), NonNullable); + let scope = DType::Struct( + StructFields::new( + ["a", "b"].into(), + vec![ + DType::Decimal(DecimalDType::new(3, 1), NonNullable), + result_dtype.clone(), + ], + ), + NonNullable, + ); + let expr = Binary.new_expr(Operator::Add, [col("a"), col("b")]); + + let coerced = coerce_expression(expr, &scope)?; + + assert!(coerced.child(0).is::()); + assert!(!coerced.child(1).is::()); + assert_eq!(coerced.return_dtype(&scope)?, result_dtype); + Ok(()) + } + + #[test] + fn decimal_arithmetic_coerces_integer_operand() -> VortexResult<()> { + let result_dtype = DType::Decimal(DecimalDType::new(5, 2), NonNullable); + let scope = DType::Struct( + StructFields::new( + ["a", "b"].into(), + vec![ + DType::Decimal(DecimalDType::new(4, 2), NonNullable), + DType::Primitive(PType::I8, NonNullable), + ], + ), + NonNullable, + ); + let expr = Binary.new_expr(Operator::Add, [col("a"), col("b")]); + + let coerced = coerce_expression(expr, &scope)?; + + assert!(coerced.child(0).is::()); + assert!(coerced.child(1).is::()); + assert_eq!(coerced.return_dtype(&scope)?, result_dtype); + Ok(()) + } + #[test] fn boolean_operators_no_coercion() -> VortexResult<()> { let scope = DType::Struct( 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 b95c557048f..7649c540113 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/decimal.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/decimal.rs @@ -32,7 +32,7 @@ use crate::arrays::Constant; use crate::arrays::ConstantArray; use crate::arrays::DecimalArray; use crate::arrays::decimal::DecimalArrayExt; -use crate::arrays::decimal::widened_buffer; +use crate::dtype::BigCast; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::DecimalType; @@ -62,42 +62,88 @@ pub(super) fn execute_numeric_decimal( let lhs = DecimalOperand::try_new(lhs, ctx)?; let rhs = DecimalOperand::try_new(rhs, ctx)?; let len = lhs.len(); - if len != rhs.len() { - vortex_bail!( - "numeric operator requires equal lengths, got {} and {}", - len, - rhs.len() - ); - } + debug_assert_eq!(len, rhs.len()); let validity = lhs.validity().and(rhs.validity())?; let valid_rows = validity.execute_mask(len, ctx)?; - let work = working_type(&lhs, &rhs, decimal_dtype); - match_each_decimal_value_type!(work, |W| { - match op { - NumericOperator::Add => execute_decimal_typed::( - &lhs, - &rhs, - decimal_dtype, - &result_dtype, - validity, - &valid_rows, - ), - NumericOperator::Sub => execute_decimal_typed::( - &lhs, - &rhs, - decimal_dtype, - &result_dtype, - validity, - &valid_rows, - ), - NumericOperator::Mul | NumericOperator::Div => vortex_bail!( - "numeric operator {:?} is not yet supported for decimal arrays", - op - ), - } - }) + let work = working_type(decimal_dtype); + let output = DecimalType::smallest_decimal_value_type(&decimal_dtype); + match (work, output) { + (DecimalType::I8, DecimalType::I8) => execute_decimal_at_widths::( + &lhs, + &rhs, + op, + decimal_dtype, + &result_dtype, + validity, + &valid_rows, + ), + (DecimalType::I16, DecimalType::I8) => execute_decimal_at_widths::( + &lhs, + &rhs, + op, + decimal_dtype, + &result_dtype, + validity, + &valid_rows, + ), + (DecimalType::I16, DecimalType::I16) => execute_decimal_at_widths::( + &lhs, + &rhs, + op, + decimal_dtype, + &result_dtype, + validity, + &valid_rows, + ), + (DecimalType::I32, DecimalType::I32) => execute_decimal_at_widths::( + &lhs, + &rhs, + op, + decimal_dtype, + &result_dtype, + validity, + &valid_rows, + ), + (DecimalType::I64, DecimalType::I64) => execute_decimal_at_widths::( + &lhs, + &rhs, + op, + decimal_dtype, + &result_dtype, + validity, + &valid_rows, + ), + (DecimalType::I128, DecimalType::I128) => execute_decimal_at_widths::( + &lhs, + &rhs, + op, + decimal_dtype, + &result_dtype, + validity, + &valid_rows, + ), + (DecimalType::I256, DecimalType::I128) => execute_decimal_at_widths::( + &lhs, + &rhs, + op, + decimal_dtype, + &result_dtype, + validity, + &valid_rows, + ), + (DecimalType::I256, DecimalType::I256) => execute_decimal_at_widths::( + &lhs, + &rhs, + op, + decimal_dtype, + &result_dtype, + validity, + &valid_rows, + ), + _ => vortex_bail!("unsupported decimal working/output width combination: {work}/{output}"), + } } /// A decimal binary-operator operand: a canonical decimal array, a non-null constant, or an @@ -152,20 +198,14 @@ impl DecimalOperand { } } -/// Choose the lane width: wide enough for the operand storage and for every intermediate value -/// that in-precision inputs can produce, capped at 256 bits. -fn working_type(lhs: &DecimalOperand, rhs: &DecimalOperand, dtype: DecimalDType) -> DecimalType { - // The sum or difference of two p-digit values has at most p + 1 digits. - let needed_digits = (dtype.precision() + 1).min(::MAX_PRECISION); - let needed = DecimalType::smallest_decimal_value_type(&DecimalDType::new(needed_digits, 0)); - - let operand_type = |operand: &DecimalOperand| match operand { - DecimalOperand::Array { values, .. } => values.values_type(), - DecimalOperand::Constant { value, .. } => smallest_value_type(value), - DecimalOperand::Null(_) => DecimalType::I8, - }; - - needed.max(operand_type(lhs)).max(operand_type(rhs)) +/// Choose the smallest lane width that can represent every sum or difference of two valid inputs. +fn working_type(dtype: DecimalDType) -> DecimalType { + let precision = dtype.precision() as usize; + let max = ::MAX_BY_PRECISION[precision]; + let max_result = max + .checked_add(&max) + .vortex_expect("the sum of two valid decimal values must fit in i256"); + smallest_value_type(&DecimalValue::from(max_result)) } /// The smallest decimal value type that can represent `value`, regardless of its stored width. @@ -245,7 +285,45 @@ impl CheckedDecimalOp for DecimalSub { } } -fn execute_decimal_typed( +fn execute_decimal_at_widths( + lhs: &DecimalOperand, + rhs: &DecimalOperand, + op: NumericOperator, + decimal_dtype: DecimalDType, + result_dtype: &DType, + validity: Validity, + valid_rows: &Mask, +) -> VortexResult +where + W: NativeDecimalType + CheckedAdd + CheckedSub, + O: NativeDecimalType, + DecimalValue: From, +{ + match op { + NumericOperator::Add => execute_decimal_typed::( + lhs, + rhs, + decimal_dtype, + result_dtype, + validity, + valid_rows, + ), + NumericOperator::Sub => execute_decimal_typed::( + lhs, + rhs, + decimal_dtype, + result_dtype, + validity, + valid_rows, + ), + NumericOperator::Mul | NumericOperator::Div => vortex_bail!( + "numeric operator {:?} is not yet supported for decimal arrays", + op + ), + } +} + +fn execute_decimal_typed( lhs: &DecimalOperand, rhs: &DecimalOperand, decimal_dtype: DecimalDType, @@ -255,7 +333,8 @@ fn execute_decimal_typed( ) -> VortexResult where W: NativeDecimalType + CheckedAdd + CheckedSub, - DecimalValue: From, + O: NativeDecimalType, + DecimalValue: From, Op: CheckedDecimalOp, { let len = lhs.len(); @@ -263,21 +342,25 @@ where let checked = match (lhs, rhs) { (DecimalOperand::Array { values: lhs, .. }, DecimalOperand::Array { values: rhs, .. }) => { - let lhs = widened_buffer::(lhs); - let rhs = widened_buffer::(rhs); - checked_decimal_lanes(len, valid_rows, |idx| { - Op::checked(lhs[idx], rhs[idx], &plan) - }) + checked_decimal_arrays::(lhs, rhs, &plan, valid_rows) } (DecimalOperand::Array { values: lhs, .. }, DecimalOperand::Constant { value, .. }) => { - let lhs = widened_buffer::(lhs); let rhs = typed_constant::(value); - checked_decimal_lanes(len, valid_rows, |idx| Op::checked(lhs[idx], rhs, &plan)) + match_each_decimal_value_type!(lhs.values_type(), |L| { + let lhs = lhs.buffer::(); + checked_decimal_lanes::(len, valid_rows, |idx| { + Op::checked(cast_work_value::(lhs[idx]), rhs, &plan) + }) + }) } (DecimalOperand::Constant { value, .. }, DecimalOperand::Array { values: rhs, .. }) => { let lhs = typed_constant::(value); - let rhs = widened_buffer::(rhs); - checked_decimal_lanes(len, valid_rows, |idx| Op::checked(lhs, rhs[idx], &plan)) + match_each_decimal_value_type!(rhs.values_type(), |R| { + let rhs = rhs.buffer::(); + checked_decimal_lanes::(len, valid_rows, |idx| { + Op::checked(lhs, cast_work_value::(rhs[idx]), &plan) + }) + }) } ( DecimalOperand::Constant { value: lhs, .. }, @@ -287,6 +370,7 @@ where let rhs = typed_constant::(rhs); let value = Op::checked(lhs, rhs, &plan) .ok_or_else(|| vortex_err!(InvalidArgument: "{}", Op::ERROR))?; + let value = cast_result_value::(value); return Ok(ConstantArray::new( Scalar::decimal( DecimalValue::from(value), @@ -297,7 +381,9 @@ where ) .into_array()); } - (DecimalOperand::Null(_), _) | (_, DecimalOperand::Null(_)) => CheckedValues::zeroed(len), + (DecimalOperand::Null(_), _) | (_, DecimalOperand::Null(_)) => { + CheckedValues::::zeroed(len) + } }; check_numeric_errors(checked.failed, Op::ERROR)?; @@ -309,18 +395,65 @@ where .into_array()) } -fn checked_decimal_lanes(len: usize, valid_rows: &Mask, checked_at: F) -> CheckedValues +fn checked_decimal_arrays( + lhs: &DecimalArray, + rhs: &DecimalArray, + plan: &DecimalOpPlan, + valid_rows: &Mask, +) -> CheckedValues +where + W: NativeDecimalType + CheckedAdd + CheckedSub, + O: NativeDecimalType, + Op: CheckedDecimalOp, +{ + let len = lhs.len(); + debug_assert_eq!(len, rhs.len()); + match_each_decimal_value_type!(lhs.values_type(), |L| { + let lhs = lhs.buffer::(); + match_each_decimal_value_type!(rhs.values_type(), |R| { + let rhs = rhs.buffer::(); + checked_decimal_lanes::(len, valid_rows, |idx| { + Op::checked( + cast_work_value::(lhs[idx]), + cast_work_value::(rhs[idx]), + plan, + ) + }) + }) + }) +} + +fn checked_decimal_lanes( + len: usize, + valid_rows: &Mask, + mut checked_at: F, +) -> CheckedValues where W: NativeDecimalType, + O: NativeDecimalType, F: FnMut(usize) -> Option, { match valid_rows.bit_buffer() { - AllOr::All => checked_all_lanes(len, checked_at), - AllOr::None => CheckedValues::zeroed(len), - AllOr::Some(valid_bits) => checked_valid_lanes(len, valid_bits, checked_at), + AllOr::All => checked_all_lanes(len, |idx| checked_at(idx).map(cast_result_value::)), + AllOr::None => CheckedValues::::zeroed(len), + AllOr::Some(valid_bits) => checked_valid_lanes(len, valid_bits, |idx| { + checked_at(idx).map(cast_result_value::) + }), } } +#[inline(always)] +fn cast_work_value(value: T) -> W { + ::from(value) + .vortex_expect("valid decimal input must fit the arithmetic working width") +} + +#[inline(always)] +fn cast_result_value(value: W) -> O { + ::from(value) + .vortex_expect("precision-checked decimal result must fit the output width") +} + fn typed_constant(value: &DecimalValue) -> W { value .cast::() 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 b0b7d754e81..ebef4a109bf 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs @@ -20,7 +20,9 @@ use vortex_error::vortex_bail; use vortex_error::vortex_err; use crate::ArrayRef; +use crate::Canonical; use crate::ExecutionCtx; +use crate::IntoArray; use crate::dtype::DType; use crate::scalar::NumericOperator; @@ -39,6 +41,21 @@ pub(crate) fn execute_numeric( ); } + if lhs.len() != rhs.len() { + vortex_bail!( + "numeric operator requires equal lengths, got {} and {}", + lhs.len(), + rhs.len() + ); + } + + if lhs.is_empty() { + let result_dtype = lhs + .dtype() + .with_nullability(lhs.dtype().nullability() | rhs.dtype().nullability()); + return Ok(Canonical::empty(&result_dtype).into_array()); + } + match lhs.dtype() { DType::Primitive(..) => primitive::execute_numeric_primitive(lhs, rhs, op, ctx), DType::Decimal(..) => decimal::execute_numeric_decimal(lhs, rhs, op, ctx), diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs index 327f078bba0..b3ccff3dcd8 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs @@ -5,7 +5,6 @@ use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexResult; -use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_mask::AllOr; use vortex_mask::Mask; @@ -145,13 +144,7 @@ where let lhs = PrimitiveOperand::::try_new(lhs, ctx)?; let rhs = PrimitiveOperand::::try_new(rhs, ctx)?; let len = lhs.len(); - if len != rhs.len() { - vortex_bail!( - "numeric operator requires equal lengths, got {} and {}", - len, - rhs.len() - ); - } + debug_assert_eq!(len, rhs.len()); let validity = lhs.validity().and(rhs.validity())?; let valid_rows = validity.execute_mask(len, ctx)?; 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 fedcf0f420c..1d2a34b8d83 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs @@ -17,6 +17,7 @@ use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::DecimalDType; +use crate::dtype::DecimalType; use crate::dtype::Nullability; use crate::scalar::DecimalValue; use crate::scalar::Scalar; @@ -212,6 +213,19 @@ fn test_present_nullable_constant_preserves_nullable_output() { ); } +#[test] +fn test_empty_primitive_constants_do_not_evaluate() -> VortexResult<()> { + let lhs = ConstantArray::new(u8::MAX, 0).into_array(); + let rhs = ConstantArray::new(1u8, 0).into_array(); + + let result = lhs + .binary(rhs, Operator::Add)? + .execute::(&mut array_session().create_execution_ctx())?; + + assert!(result.0.is_empty()); + Ok(()) +} + // -- Decimal arithmetic -- fn decimal_binary(lhs: ArrayRef, rhs: ArrayRef, op: Operator) -> VortexResult { @@ -321,6 +335,46 @@ fn test_decimal_precision_stricter_than_width() { assert!(decimal_binary(lhs, rhs, Operator::Add).is_err()); } +#[rstest] +#[case::precision_2( + DecimalArray::from_iter::([10, 20], DecimalDType::new(2, 0)), + DecimalType::I8, +)] +#[case::precision_18( + DecimalArray::from_iter::([10, 20], DecimalDType::new(18, 0)), + DecimalType::I64, +)] +#[case::precision_38( + DecimalArray::from_iter::([10, 20], DecimalDType::new(38, 0)), + DecimalType::I128, +)] +fn test_decimal_result_uses_logical_storage_width( + #[case] values: DecimalArray, + #[case] expected_type: DecimalType, +) -> VortexResult<()> { + let result = decimal_binary( + values.clone().into_array(), + values.into_array(), + Operator::Add, + )? + .execute::(&mut array_session().create_execution_ctx())?; + + assert_eq!(result.values_type(), expected_type); + Ok(()) +} + +#[test] +fn test_decimal_empty_constants_do_not_evaluate() -> VortexResult<()> { + let dtype = DecimalDType::new(3, 0); + let lhs = decimal_constant(999i16, dtype, 0); + let rhs = decimal_constant(1i16, dtype, 0); + + let result = decimal_binary(lhs, rhs, Operator::Add)?; + + assert!(result.is_empty()); + Ok(()) +} + #[test] fn test_decimal_constant_lhs_non_commutative() { let mut ctx = array_session().create_execution_ctx(); From 3a56ba9eeeb7b7541b3ba5534e734b3899f93697 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 15 Jul 2026 09:21:50 -0700 Subject: [PATCH 5/5] Match Arrow decimal Add/Sub result typing Signed-off-by: Matt Katz --- .../src/compute/conformance/binary_numeric.rs | 15 ++- vortex-array/src/expr/transform/coerce.rs | 28 +--- vortex-array/src/scalar_fn/fns/binary/mod.rs | 24 +++- .../scalar_fn/fns/binary/numeric/decimal.rs | 106 +++++---------- .../src/scalar_fn/fns/binary/numeric/mod.rs | 14 +- .../src/scalar_fn/fns/binary/numeric/tests.rs | 126 +++++++++++++----- 6 files changed, 170 insertions(+), 143 deletions(-) diff --git a/vortex-array/src/compute/conformance/binary_numeric.rs b/vortex-array/src/compute/conformance/binary_numeric.rs index 8899ad89b28..2cc6484bc18 100644 --- a/vortex-array/src/compute/conformance/binary_numeric.rs +++ b/vortex-array/src/compute/conformance/binary_numeric.rs @@ -49,6 +49,7 @@ use crate::scalar::DecimalValue; use crate::scalar::NumericOperator; use crate::scalar::PrimitiveScalar; use crate::scalar::Scalar; +use crate::scalar_fn::fns::binary::decimal_add_sub_result_dtype; fn to_vec_of_scalar(array: &ArrayRef, ctx: &mut ExecutionCtx) -> Vec { // Not fast, but obviously correct @@ -299,6 +300,8 @@ fn test_decimal_binary_numeric_with_scalar( let original_values = to_vec_of_scalar(&canonicalized_array, ctx); let scalar = Scalar::decimal(value, decimal_dtype, array.dtype().nullability()); + let result_decimal_dtype = decimal_add_sub_result_dtype(decimal_dtype); + let result_dtype = DType::Decimal(result_decimal_dtype, array.dtype().nullability()); // Decimal Mul/Div are not yet implemented. let operators = vec![NumericOperator::Add, NumericOperator::Sub]; @@ -321,7 +324,17 @@ fn test_decimal_binary_numeric_with_scalar( } else { (scalar.as_decimal(), x.as_decimal()) }; - lhs.checked_binary_numeric(&rhs, operator).map(Scalar::from) + let (Some(lhs), Some(rhs)) = (lhs.decimal_value(), rhs.decimal_value()) else { + return Some(Scalar::null(result_dtype.clone())); + }; + let value = match operator { + NumericOperator::Add => lhs.checked_add(&rhs), + NumericOperator::Sub => lhs.checked_sub(&rhs), + NumericOperator::Mul | NumericOperator::Div => unreachable!(), + }?; + value.fits_in_precision(result_decimal_dtype).then(|| { + Scalar::decimal(value, result_decimal_dtype, result_dtype.nullability()) + }) }) .collect(); diff --git a/vortex-array/src/expr/transform/coerce.rs b/vortex-array/src/expr/transform/coerce.rs index 7d431163959..781045d652b 100644 --- a/vortex-array/src/expr/transform/coerce.rs +++ b/vortex-array/src/expr/transform/coerce.rs @@ -156,13 +156,14 @@ mod tests { #[test] fn decimal_arithmetic_coerces_precision_and_scale() -> VortexResult<()> { - let result_dtype = DType::Decimal(DecimalDType::new(4, 2), NonNullable); + let common_dtype = DType::Decimal(DecimalDType::new(4, 2), NonNullable); + let result_dtype = DType::Decimal(DecimalDType::new(5, 2), NonNullable); let scope = DType::Struct( StructFields::new( ["a", "b"].into(), vec![ DType::Decimal(DecimalDType::new(3, 1), NonNullable), - result_dtype.clone(), + common_dtype, ], ), NonNullable, @@ -177,29 +178,6 @@ mod tests { Ok(()) } - #[test] - fn decimal_arithmetic_coerces_integer_operand() -> VortexResult<()> { - let result_dtype = DType::Decimal(DecimalDType::new(5, 2), NonNullable); - let scope = DType::Struct( - StructFields::new( - ["a", "b"].into(), - vec![ - DType::Decimal(DecimalDType::new(4, 2), NonNullable), - DType::Primitive(PType::I8, NonNullable), - ], - ), - NonNullable, - ); - let expr = Binary.new_expr(Operator::Add, [col("a"), col("b")]); - - let coerced = coerce_expression(expr, &scope)?; - - assert!(coerced.child(0).is::()); - assert!(coerced.child(1).is::()); - assert_eq!(coerced.return_dtype(&scope)?, result_dtype); - Ok(()) - } - #[test] fn boolean_operators_no_coercion() -> VortexResult<()> { let scope = DType::Struct( diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index 7532801853c..7ee31b11abe 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -17,6 +17,8 @@ use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::ExecutionCtx; use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::MAX_PRECISION; use crate::dtype::Nullability; use crate::expr::and; use crate::expr::expression::Expression; @@ -48,6 +50,14 @@ use crate::scalar::Scalar; #[derive(Clone)] pub struct Binary; +/// Derive the result type for Add/Sub over operands that already share a decimal dtype. +pub(crate) fn decimal_add_sub_result_dtype(input: DecimalDType) -> DecimalDType { + DecimalDType::new( + input.precision().saturating_add(1).min(MAX_PRECISION), + input.scale(), + ) +} + impl ScalarFnVTable for Binary { type Options = Operator; @@ -118,12 +128,18 @@ impl ScalarFnVTable for Binary { let rhs = &arg_dtypes[1]; if operator.is_arithmetic() { - // Decimal Mul/Div need rescaling support and are not yet implemented. - let decimal_supported = - lhs.is_decimal() && matches!(operator, Operator::Add | Operator::Sub); - if (lhs.is_primitive() || decimal_supported) && lhs.eq_ignore_nullability(rhs) { + if lhs.is_primitive() && lhs.eq_ignore_nullability(rhs) { return Ok(lhs.with_nullability(lhs.nullability() | rhs.nullability())); } + if let DType::Decimal(decimal_dtype, _) = lhs + && matches!(operator, Operator::Add | Operator::Sub) + && lhs.eq_ignore_nullability(rhs) + { + return Ok(DType::Decimal( + decimal_add_sub_result_dtype(*decimal_dtype), + lhs.nullability() | rhs.nullability(), + )); + } vortex_bail!( "incompatible types for arithmetic operation: {} {}", lhs, 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 7649c540113..a2421d9c14e 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/decimal.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/decimal.rs @@ -3,10 +3,10 @@ //! Native execution of the arithmetic operators over decimal arrays. //! -//! Both operands share a logical [`DecimalDType`] (equal precision and scale) and the result -//! keeps that dtype: fixed-point arithmetic at the shared scale. Add and Sub apply directly to -//! the unscaled stored integers and are exact; Mul and Div require rescaling and are not yet -//! implemented. +//! 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. Mul +//! and Div require rescaling and are not yet implemented. //! //! 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 @@ -25,6 +25,7 @@ use super::CheckedValues; use super::check_numeric_errors; use super::checked_all_lanes; use super::checked_valid_lanes; +use super::decimal_add_sub_result_dtype; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; @@ -54,10 +55,11 @@ pub(super) fn execute_numeric_decimal( let DType::Decimal(decimal_dtype, _) = lhs.dtype() else { vortex_bail!("expected a decimal dtype, got {}", lhs.dtype()); }; - let decimal_dtype = *decimal_dtype; - let result_dtype = lhs - .dtype() - .with_nullability(lhs.dtype().nullability() | rhs.dtype().nullability()); + let result_decimal_dtype = decimal_add_sub_result_dtype(*decimal_dtype); + let result_dtype = DType::Decimal( + result_decimal_dtype, + lhs.dtype().nullability() | rhs.dtype().nullability(), + ); let lhs = DecimalOperand::try_new(lhs, ctx)?; let rhs = DecimalOperand::try_new(rhs, ctx)?; @@ -67,82 +69,61 @@ pub(super) fn execute_numeric_decimal( let validity = lhs.validity().and(rhs.validity())?; let valid_rows = validity.execute_mask(len, ctx)?; - let work = working_type(decimal_dtype); - let output = DecimalType::smallest_decimal_value_type(&decimal_dtype); - match (work, output) { - (DecimalType::I8, DecimalType::I8) => execute_decimal_at_widths::( + match DecimalType::smallest_decimal_value_type(&result_decimal_dtype) { + DecimalType::I8 => execute_decimal_at_widths::( &lhs, &rhs, op, - decimal_dtype, + result_decimal_dtype, &result_dtype, validity, &valid_rows, ), - (DecimalType::I16, DecimalType::I8) => execute_decimal_at_widths::( + DecimalType::I16 => execute_decimal_at_widths::( &lhs, &rhs, op, - decimal_dtype, + result_decimal_dtype, &result_dtype, validity, &valid_rows, ), - (DecimalType::I16, DecimalType::I16) => execute_decimal_at_widths::( + DecimalType::I32 => execute_decimal_at_widths::( &lhs, &rhs, op, - decimal_dtype, + result_decimal_dtype, &result_dtype, validity, &valid_rows, ), - (DecimalType::I32, DecimalType::I32) => execute_decimal_at_widths::( + DecimalType::I64 => execute_decimal_at_widths::( &lhs, &rhs, op, - decimal_dtype, + result_decimal_dtype, &result_dtype, validity, &valid_rows, ), - (DecimalType::I64, DecimalType::I64) => execute_decimal_at_widths::( + DecimalType::I128 => execute_decimal_at_widths::( &lhs, &rhs, op, - decimal_dtype, + result_decimal_dtype, &result_dtype, validity, &valid_rows, ), - (DecimalType::I128, DecimalType::I128) => execute_decimal_at_widths::( + DecimalType::I256 => execute_decimal_at_widths::( &lhs, &rhs, op, - decimal_dtype, + result_decimal_dtype, &result_dtype, validity, &valid_rows, ), - (DecimalType::I256, DecimalType::I128) => execute_decimal_at_widths::( - &lhs, - &rhs, - op, - decimal_dtype, - &result_dtype, - validity, - &valid_rows, - ), - (DecimalType::I256, DecimalType::I256) => execute_decimal_at_widths::( - &lhs, - &rhs, - op, - decimal_dtype, - &result_dtype, - validity, - &valid_rows, - ), - _ => vortex_bail!("unsupported decimal working/output width combination: {work}/{output}"), } } @@ -198,33 +179,6 @@ impl DecimalOperand { } } -/// Choose the smallest lane width that can represent every sum or difference of two valid inputs. -fn working_type(dtype: DecimalDType) -> DecimalType { - let precision = dtype.precision() as usize; - let max = ::MAX_BY_PRECISION[precision]; - let max_result = max - .checked_add(&max) - .vortex_expect("the sum of two valid decimal values must fit in i256"); - smallest_value_type(&DecimalValue::from(max_result)) -} - -/// The smallest decimal value type that can represent `value`, regardless of its stored width. -fn smallest_value_type(value: &DecimalValue) -> DecimalType { - if value.cast::().is_some() { - DecimalType::I8 - } else if value.cast::().is_some() { - DecimalType::I16 - } else if value.cast::().is_some() { - DecimalType::I32 - } else if value.cast::().is_some() { - DecimalType::I64 - } else if value.cast::().is_some() { - DecimalType::I128 - } else { - DecimalType::I256 - } -} - /// Per-execution constants for checked decimal lane operations at working width `W`. struct DecimalOpPlan { /// Inclusive stored-value bounds implied by the result precision. @@ -289,7 +243,7 @@ fn execute_decimal_at_widths( lhs: &DecimalOperand, rhs: &DecimalOperand, op: NumericOperator, - decimal_dtype: DecimalDType, + result_decimal_dtype: DecimalDType, result_dtype: &DType, validity: Validity, valid_rows: &Mask, @@ -303,7 +257,7 @@ where NumericOperator::Add => execute_decimal_typed::( lhs, rhs, - decimal_dtype, + result_decimal_dtype, result_dtype, validity, valid_rows, @@ -311,7 +265,7 @@ where NumericOperator::Sub => execute_decimal_typed::( lhs, rhs, - decimal_dtype, + result_decimal_dtype, result_dtype, validity, valid_rows, @@ -326,7 +280,7 @@ where fn execute_decimal_typed( lhs: &DecimalOperand, rhs: &DecimalOperand, - decimal_dtype: DecimalDType, + result_decimal_dtype: DecimalDType, result_dtype: &DType, validity: Validity, valid_rows: &Mask, @@ -338,7 +292,7 @@ where Op: CheckedDecimalOp, { let len = lhs.len(); - let plan = DecimalOpPlan::::new(decimal_dtype); + let plan = DecimalOpPlan::::new(result_decimal_dtype); let checked = match (lhs, rhs) { (DecimalOperand::Array { values: lhs, .. }, DecimalOperand::Array { values: rhs, .. }) => { @@ -374,7 +328,7 @@ where return Ok(ConstantArray::new( Scalar::decimal( DecimalValue::from(value), - decimal_dtype, + result_decimal_dtype, result_dtype.nullability(), ), len, @@ -389,7 +343,7 @@ where Ok(DecimalArray::new( checked.values, - decimal_dtype, + result_decimal_dtype, validity.union_nullability(result_dtype.nullability()), ) .into_array()) 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 ebef4a109bf..4c28a9a3570 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs @@ -6,6 +6,8 @@ //! //! [`Binary`]: super::Binary +use super::decimal_add_sub_result_dtype; + mod decimal; mod primitive; #[cfg(test)] @@ -50,9 +52,15 @@ pub(crate) fn execute_numeric( } if lhs.is_empty() { - let result_dtype = lhs - .dtype() - .with_nullability(lhs.dtype().nullability() | rhs.dtype().nullability()); + 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, _) => { + debug_assert!(matches!(op, NumericOperator::Add | NumericOperator::Sub)); + DType::Decimal(decimal_add_sub_result_dtype(*decimal_dtype), nullability) + } + dtype => vortex_bail!("numeric operator is not supported for dtype {}", dtype), + }; return Ok(Canonical::empty(&result_dtype).into_array()); } 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 1d2a34b8d83..b5129cc671a 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs @@ -5,6 +5,7 @@ use rstest::rstest; use vortex_buffer::buffer; use vortex_error::VortexResult; +use super::decimal_add_sub_result_dtype; use crate::ArrayRef; use crate::IntoArray; use crate::RecursiveCanonical; @@ -18,7 +19,9 @@ use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::DecimalType; +use crate::dtype::NativeDecimalType; use crate::dtype::Nullability; +use crate::dtype::i256; use crate::scalar::DecimalValue; use crate::scalar::Scalar; use crate::scalar_fn::fns::operators::Operator; @@ -252,13 +255,14 @@ fn test_decimal_array_array( ) { let mut ctx = array_session().create_execution_ctx(); let dtype = DecimalDType::new(10, 2); + let result_dtype = decimal_add_sub_result_dtype(dtype); let lhs = DecimalArray::from_iter::([900, 1000], dtype).into_array(); let rhs = DecimalArray::from_iter::(rhs, dtype).into_array(); let result = decimal_binary(lhs, rhs, op).unwrap(); assert_arrays_eq!( result, - DecimalArray::from_iter::(expected, dtype), + DecimalArray::from_iter::(expected, result_dtype), &mut ctx ); } @@ -273,7 +277,7 @@ fn test_decimal_mixed_storage_widths() { let result = decimal_binary(lhs, rhs, Operator::Add).unwrap(); assert_arrays_eq!( result, - DecimalArray::from_iter::([300, 500], dtype), + DecimalArray::from_iter::([300, 500], decimal_add_sub_result_dtype(dtype)), &mut ctx ); } @@ -289,16 +293,20 @@ fn test_decimal_nullable_lanes() { let result = decimal_binary(lhs, rhs, Operator::Add).unwrap(); assert_arrays_eq!( result, - DecimalArray::from_option_iter::([Some(150), None, Some(350)], dtype), + DecimalArray::from_option_iter::( + [Some(150), None, Some(350)], + decimal_add_sub_result_dtype(dtype), + ), &mut ctx ); } #[test] -fn test_decimal_overflow_on_valid_lane_errors() { - let dtype = DecimalDType::new(3, 0); - let lhs = DecimalArray::from_iter::([999], dtype).into_array(); - let rhs = DecimalArray::from_iter::([2], dtype).into_array(); +fn test_decimal_max_precision_overflow_on_valid_lane_errors() { + let dtype = DecimalDType::new(76, 0); + let max = ::MAX_BY_PRECISION[76]; + let lhs = DecimalArray::from_iter::([max], dtype).into_array(); + let rhs = DecimalArray::from_iter::([i256::from_i128(1)], dtype).into_array(); assert!(decimal_binary(lhs, rhs, Operator::Add).is_err()); } @@ -306,52 +314,62 @@ fn test_decimal_overflow_on_valid_lane_errors() { #[test] fn test_decimal_overflow_on_null_lane_ignored() { let mut ctx = array_session().create_execution_ctx(); - let dtype = DecimalDType::new(3, 0); - // The null lane holds 999, so adding 500 overflows the precision there but is ignored. - let lhs = DecimalArray::new( - buffer![999i16, 1], - dtype, - Validity::from_iter([false, true]), - ) - .into_array(); - let rhs = decimal_constant(500i16, dtype, 2); + let dtype = DecimalDType::new(76, 0); + let max = ::MAX_BY_PRECISION[76]; + let one = i256::from_i128(1); + // The null lane holds the maximum value, so adding one overflows there but is ignored. + let lhs = DecimalArray::new(buffer![max, one], dtype, Validity::from_iter([false, true])) + .into_array(); + let rhs = decimal_constant(one, dtype, 2); let result = decimal_binary(lhs, rhs, Operator::Add).unwrap(); assert_arrays_eq!( result, - DecimalArray::from_option_iter::([None, Some(501)], dtype), + DecimalArray::from_option_iter::([None, Some(i256::from_i128(2))], dtype,), &mut ctx ); } -/// A value can fit the storage width while violating the dtype precision: 60 + 60 fits an i8 but -/// exceeds precision 2. #[test] -fn test_decimal_precision_stricter_than_width() { +fn test_decimal_add_reserves_carry_digit() { + let mut ctx = array_session().create_execution_ctx(); let dtype = DecimalDType::new(2, 0); let lhs = DecimalArray::from_iter::([60], dtype).into_array(); let rhs = DecimalArray::from_iter::([60], dtype).into_array(); - assert!(decimal_binary(lhs, rhs, Operator::Add).is_err()); + let result = decimal_binary(lhs, rhs, Operator::Add).unwrap(); + assert_arrays_eq!( + result, + DecimalArray::from_iter::([120], DecimalDType::new(3, 0)), + &mut ctx + ); } #[rstest] #[case::precision_2( DecimalArray::from_iter::([10, 20], DecimalDType::new(2, 0)), - DecimalType::I8, + DecimalType::I16, )] #[case::precision_18( DecimalArray::from_iter::([10, 20], DecimalDType::new(18, 0)), - DecimalType::I64, + DecimalType::I128, )] #[case::precision_38( DecimalArray::from_iter::([10, 20], DecimalDType::new(38, 0)), - DecimalType::I128, + DecimalType::I256, )] -fn test_decimal_result_uses_logical_storage_width( +#[case::precision_76( + DecimalArray::from_iter::( + [i256::from_i128(10), i256::from_i128(20)], + DecimalDType::new(76, 0), + ), + DecimalType::I256, +)] +fn test_decimal_result_uses_widened_logical_storage_width( #[case] values: DecimalArray, #[case] expected_type: DecimalType, ) -> VortexResult<()> { + let expected_dtype = decimal_add_sub_result_dtype(values.decimal_dtype()); let result = decimal_binary( values.clone().into_array(), values.into_array(), @@ -360,18 +378,52 @@ fn test_decimal_result_uses_logical_storage_width( .execute::(&mut array_session().create_execution_ctx())?; assert_eq!(result.values_type(), expected_type); + assert_eq!(result.decimal_dtype(), expected_dtype); + Ok(()) +} + +#[test] +fn test_decimal_precision_76_boundary() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DecimalDType::new(76, 0); + let max = ::MAX_BY_PRECISION[76]; + let zero = i256::from_i128(0); + let one = i256::from_i128(1); + + let result = decimal_binary( + DecimalArray::from_iter::([max], dtype).into_array(), + DecimalArray::from_iter::([zero], dtype).into_array(), + Operator::Add, + )?; + assert_arrays_eq!( + result, + DecimalArray::from_iter::([max], dtype), + &mut ctx + ); + + let overflow = decimal_binary( + DecimalArray::from_iter::([max], dtype).into_array(), + DecimalArray::from_iter::([one], dtype).into_array(), + Operator::Add, + ); + assert!(overflow.is_err()); Ok(()) } #[test] fn test_decimal_empty_constants_do_not_evaluate() -> VortexResult<()> { - let dtype = DecimalDType::new(3, 0); - let lhs = decimal_constant(999i16, dtype, 0); - let rhs = decimal_constant(1i16, dtype, 0); + let dtype = DecimalDType::new(76, 0); + let max = ::MAX_BY_PRECISION[76]; + let lhs = decimal_constant(max, dtype, 0); + let rhs = decimal_constant(i256::from_i128(1), dtype, 0); let result = decimal_binary(lhs, rhs, Operator::Add)?; assert!(result.is_empty()); + assert_eq!( + result.dtype(), + &DType::Decimal(dtype, Nullability::NonNullable) + ); Ok(()) } @@ -385,7 +437,7 @@ fn test_decimal_constant_lhs_non_commutative() { let result = decimal_binary(lhs, rhs, Operator::Sub).unwrap(); assert_arrays_eq!( result, - DecimalArray::from_iter::([750, 600], dtype), + DecimalArray::from_iter::([750, 600], decimal_add_sub_result_dtype(dtype),), &mut ctx ); } @@ -404,7 +456,10 @@ fn test_decimal_nullable_constant_preserves_nullable_output() { let result = decimal_binary(values, constant, Operator::Add).unwrap(); assert_arrays_eq!( result, - DecimalArray::from_option_iter::([Some(150), Some(250)], dtype), + DecimalArray::from_option_iter::( + [Some(150), Some(250)], + decimal_add_sub_result_dtype(dtype), + ), &mut ctx ); } @@ -423,7 +478,7 @@ fn test_decimal_null_constant_yields_all_null() { let result = decimal_binary(values, null_constant, Operator::Add).unwrap(); assert_arrays_eq!( result, - DecimalArray::from_option_iter::([None, None], dtype), + DecimalArray::from_option_iter::([None, None], decimal_add_sub_result_dtype(dtype),), &mut ctx ); } @@ -440,7 +495,10 @@ fn test_decimal_constant_wider_than_array_storage() { let result = decimal_binary(values, constant, Operator::Add).unwrap(); assert_arrays_eq!( result, - DecimalArray::from_iter::([10_000_000_001, 10_000_000_002], dtype), + DecimalArray::from_iter::( + [10_000_000_001, 10_000_000_002], + decimal_add_sub_result_dtype(dtype), + ), &mut ctx ); } @@ -454,7 +512,7 @@ fn test_decimal_empty() { let result = decimal_binary(empty.clone(), empty, Operator::Add).unwrap(); assert_arrays_eq!( result, - DecimalArray::from_iter::([], dtype), + DecimalArray::from_iter::([], decimal_add_sub_result_dtype(dtype)), &mut ctx ); } @@ -469,7 +527,7 @@ fn test_decimal_constant_constant_folds() { let result = decimal_binary(lhs, rhs, Operator::Add).unwrap(); assert_arrays_eq!( result, - DecimalArray::from_iter::([200, 200, 200], dtype), + DecimalArray::from_iter::([200, 200, 200], decimal_add_sub_result_dtype(dtype),), &mut ctx ); }