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..314e6e5e60e 100644 --- a/vortex-array/benches/binary_ops.rs +++ b/vortex-array/benches/binary_ops.rs @@ -18,8 +18,13 @@ 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::dtype::Nullability; +use vortex_array::scalar::DecimalValue; +use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_session::VortexSession; @@ -136,6 +141,47 @@ 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 mul_decimal_i64_nonnull(bencher: Bencher) { + // Multiplication at precision 18 widens the working type to i128. + let lhs = decimal_i64_nonnull(1).into_array(); + let rhs = decimal_i64_nonnull(17).into_array(); + + bench_decimal(bencher, lhs, rhs, Operator::Mul); +} + +#[divan::bench] +fn div_decimal_i128_constant(bencher: Bencher) { + let lhs = decimal_i128_nullable(1_000_000, 7).into_array(); + let rhs = ConstantArray::new( + Scalar::decimal( + DecimalValue::I128(37_00), + DecimalDType::new(38, 2), + Nullability::NonNullable, + ), + LEN, + ) + .into_array(); + + bench_decimal(bencher, lhs, rhs, Operator::Div); +} + #[divan::bench] fn eq_i64_constant(bencher: Bencher) { let lhs = primitive_nonnull(0).into_array(); @@ -172,6 +218,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 +247,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/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/compute/conformance/binary_numeric.rs b/vortex-array/src/compute/conformance/binary_numeric.rs index 1fe409bf648..2226b30573e 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,131 @@ 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()); + + // Skip division by zero. + let operators = if value.is_zero() { + vec![ + NumericOperator::Add, + NumericOperator::Sub, + NumericOperator::Mul, + ] + } else { + vec![ + NumericOperator::Add, + NumericOperator::Sub, + NumericOperator::Mul, + NumericOperator::Div, + ] + }; + + 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: 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..43eba8f014d 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -118,7 +118,7 @@ impl ScalarFnVTable for Binary { let rhs = &arg_dtypes[1]; if operator.is_arithmetic() { - if lhs.is_primitive() && lhs.eq_ignore_nullability(rhs) { + if (lhs.is_primitive() || lhs.is_decimal()) && 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..b7bc165002f --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/decimal.rs @@ -0,0 +1,422 @@ +// 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, truncating toward zero. Add and +//! Sub apply directly to the unscaled stored integers; Mul rescales the raw product from twice +//! the scale back to the scale; Div pre-scales the dividend (or, for negative scales, the +//! divisor). +//! +//! Lanes execute in a working width chosen so that in-precision inputs cannot spuriously +//! overflow an intermediate value (capped at 256 bits, where an overflowing intermediate is +//! reported as an error even if the rescaled result would fit). An operation that overflows the +//! result precision or divides by zero on a valid lane is an error; invalid lanes never error. + +use num_traits::CheckedAdd; +use num_traits::CheckedDiv; +use num_traits::CheckedMul; +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::BigCast; +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, op); + 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 => execute_decimal_typed::( + &lhs, + &rhs, + decimal_dtype, + &result_dtype, + validity, + &valid_rows, + ), + NumericOperator::Div => execute_decimal_typed::( + &lhs, + &rhs, + decimal_dtype, + &result_dtype, + validity, + &valid_rows, + ), + } + }) +} + +/// 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 for executing `op`: 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, + op: NumericOperator, +) -> DecimalType { + let precision = dtype.precision() as u16; + let scale_digits = dtype.scale().unsigned_abs() as u16; + let needed_digits = match op { + // The sum of two p-digit values has at most p + 1 digits. + NumericOperator::Add | NumericOperator::Sub => precision + 1, + // The raw product has at most 2p digits; a negative scale then multiplies it by 10^|s|. + NumericOperator::Mul => 2 * precision + if dtype.scale() < 0 { scale_digits } else { 0 }, + // The pre-scaled dividend (or divisor) has at most p + |s| digits. + NumericOperator::Div => precision + scale_digits, + }; + let needed_digits = u8::try_from(needed_digits) + .unwrap_or(::MAX_PRECISION) + .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, + /// `10^|scale|` in `W`, or `None` if unrepresentable (only for scales below -76, where any + /// nonzero rescale overflows anyway). + scale_factor: Option, + /// The shared operand/result scale. + scale: i8, +} + +impl DecimalOpPlan { + fn new(dtype: DecimalDType) -> Self { + let precision = dtype.precision() as usize; + let scale = dtype.scale(); + let scale_factor = i256::from_i128(10) + .checked_pow(scale.unsigned_abs() as u32) + .and_then(::from); + Self { + prec_min: W::MIN_BY_PRECISION[precision], + prec_max: W::MAX_BY_PRECISION[precision], + scale_factor, + scale, + } + } + + /// 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 + CheckedMul + CheckedDiv; +} + +struct DecimalAdd; + +struct DecimalSub; + +struct DecimalMul; + +struct DecimalDiv; + +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 + CheckedMul + CheckedDiv, + { + 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 + CheckedMul + CheckedDiv, + { + plan.in_precision(lhs.checked_sub(&rhs)?) + } +} + +impl CheckedDecimalOp for DecimalMul { + const ERROR: &'static str = "decimal overflow in checked mul"; + + #[inline(always)] + fn checked(lhs: W, rhs: W, plan: &DecimalOpPlan) -> Option + where + W: NativeDecimalType + CheckedAdd + CheckedSub + CheckedMul + CheckedDiv, + { + let product = lhs.checked_mul(&rhs)?; + // The raw product has twice the operand scale; bring it back, truncating toward zero. + let rescaled = if product == W::default() || plan.scale == 0 { + product + } else if plan.scale > 0 { + product.checked_div(&plan.scale_factor?)? + } else { + product.checked_mul(&plan.scale_factor?)? + }; + plan.in_precision(rescaled) + } +} + +impl CheckedDecimalOp for DecimalDiv { + const ERROR: &'static str = "decimal division by zero or overflow in checked div"; + + #[inline(always)] + fn checked(lhs: W, rhs: W, plan: &DecimalOpPlan) -> Option + where + W: NativeDecimalType + CheckedAdd + CheckedSub + CheckedMul + CheckedDiv, + { + if rhs == W::default() { + return None; + } + if lhs == W::default() { + return Some(lhs); + } + // Pre-scaling the dividend (or, for negative scales, the divisor) by 10^|s| yields a + // quotient at the operand scale, truncating toward zero. + let quotient = if plan.scale >= 0 { + lhs.checked_mul(&plan.scale_factor?)?.checked_div(&rhs)? + } else { + lhs.checked_div(&rhs.checked_mul(&plan.scale_factor?)?)? + }; + plan.in_precision(quotient) + } +} + +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 + CheckedMul + CheckedDiv, + 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 new file mode 100644 index 00000000000..b0b7d754e81 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs @@ -0,0 +1,204 @@ +// 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 decimal; +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_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. +pub(crate) fn execute_numeric( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: NumericOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + 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. +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 65% 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..327f078bba0 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,21 +109,14 @@ 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, 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 { @@ -214,7 +212,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 +226,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 +250,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 +266,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 +474,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 +763,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..a760d47c46a --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs @@ -0,0 +1,559 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +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::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::dtype::i256; +use crate::scalar::DecimalValue; +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 + ); +} + +// -- 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])] +#[case::mul(Operator::Mul, [50i64, 200], [450i64, 2000])] // 9.00 * 0.50, 10.00 * 2.00 +#[case::div(Operator::Div, [300i64, 400], [300i64, 250])] // 9.00 / 3.00, 10.00 / 4.00 +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()); +} + +#[rstest] +#[case::rescales([50i64], [10i64], [5i64])] // 0.50 * 0.10 = 0.05 +#[case::truncates([15i64], [15i64], [2i64])] // 0.15 * 0.15 = 0.0225 -> 0.02 +#[case::truncates_toward_zero([-15i64], [15i64], [-2i64])] +fn test_decimal_mul_fixed_point( + #[case] lhs: [i64; 1], + #[case] rhs: [i64; 1], + #[case] expected: [i64; 1], +) { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DecimalDType::new(10, 2); + let lhs = DecimalArray::from_iter::(lhs, dtype).into_array(); + let rhs = DecimalArray::from_iter::(rhs, dtype).into_array(); + + let result = decimal_binary(lhs, rhs, Operator::Mul).unwrap(); + assert_arrays_eq!( + result, + DecimalArray::from_iter::(expected, dtype), + &mut ctx + ); +} + +/// The raw product may overflow the operand storage width while the rescaled result fits: +/// 3000.00 * 3000.00 = 9,000,000.00 needs a wider intermediate than the i32 storage. +#[test] +fn test_decimal_mul_wider_than_operand_storage() { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DecimalDType::new(10, 2); + let values = DecimalArray::from_iter::([300_000], dtype).into_array(); + + let result = decimal_binary(values.clone(), values, Operator::Mul).unwrap(); + assert_arrays_eq!( + result, + DecimalArray::from_iter::([900_000_000], dtype), + &mut ctx + ); +} + +#[rstest] +#[case::rescales([1000i64], [10i64], [10000i64])] // 10.00 / 0.10 = 100.00 +#[case::truncates([1000i64], [300i64], [333i64])] // 10.00 / 3.00 = 3.33... +#[case::truncates_toward_zero([-1000i64], [300i64], [-333i64])] +fn test_decimal_div_fixed_point( + #[case] lhs: [i64; 1], + #[case] rhs: [i64; 1], + #[case] expected: [i64; 1], +) { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DecimalDType::new(10, 2); + let lhs = DecimalArray::from_iter::(lhs, dtype).into_array(); + let rhs = DecimalArray::from_iter::(rhs, dtype).into_array(); + + let result = decimal_binary(lhs, rhs, Operator::Div).unwrap(); + assert_arrays_eq!( + result, + DecimalArray::from_iter::(expected, dtype), + &mut ctx + ); +} + +#[test] +fn test_decimal_divide_by_zero_errors() { + let dtype = DecimalDType::new(10, 2); + let lhs = DecimalArray::from_iter::([100], dtype).into_array(); + let rhs = decimal_constant(0i64, dtype, 1); + + assert!(decimal_binary(lhs, rhs, Operator::Div).is_err()); +} + +#[test] +fn test_decimal_divide_by_zero_on_null_lane_ignored() { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DecimalDType::new(10, 2); + let lhs = DecimalArray::from_option_iter::([None, Some(1000)], dtype).into_array(); + let rhs = DecimalArray::from_iter::([0, 500], dtype).into_array(); + + let result = decimal_binary(lhs, rhs, Operator::Div).unwrap(); + assert_arrays_eq!( + result, + DecimalArray::from_option_iter::([None, Some(200)], dtype), + &mut ctx + ); +} + +#[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 + ); +} + +/// p=38 multiplication requires a 256-bit intermediate product. +#[test] +fn test_decimal_mul_i256_working_width() { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DecimalDType::new(38, 2); + let big = 10_i128.pow(19); + let values = DecimalArray::from_iter::([big], dtype).into_array(); + + // (10^17).00 * (10^17).00 = 10^34.00, stored as 10^36. + let result = decimal_binary(values.clone(), values, Operator::Mul).unwrap(); + assert_arrays_eq!( + result, + DecimalArray::from_iter::([10_i128.pow(36)], dtype), + &mut ctx + ); +} + +/// Near the 76-digit cap, a raw product that overflows i256 errors even if the rescaled result +/// would fit. This is a documented limitation of the 256-bit intermediate. +#[test] +fn test_decimal_mul_i256_intermediate_overflow_errors() { + let dtype = DecimalDType::new(76, 76); + let big = i256::from_i128(10).checked_pow(75).unwrap(); + let values = DecimalArray::from_iter::([big], dtype).into_array(); + + assert!(decimal_binary(values.clone(), values, Operator::Mul).is_err()); +} + +#[rstest] +#[case::mul(Operator::Mul, [5i64], [3i64], [1500i64])] // 500 * 300 = 150,000, stored 1500 +#[case::div(Operator::Div, [600i64], [3i64], [2i64])] // 60,000 / 300 = 200, stored 2 +#[case::div_truncates(Operator::Div, [5000i64], [3i64], [16i64])] // 500,000 / 300 = 1666.67 -> 1600 +fn test_decimal_negative_scale( + #[case] op: Operator, + #[case] lhs: [i64; 1], + #[case] rhs: [i64; 1], + #[case] expected: [i64; 1], +) { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DecimalDType::new(10, -2); + let lhs = DecimalArray::from_iter::(lhs, 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_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::Mul).unwrap(); + assert_arrays_eq!( + result, + DecimalArray::from_iter::([75, 75, 75], dtype), + &mut ctx + ); +}