Decimal Add/Sub kernels

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);
+    }
 }

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::<PrimitiveArray>(bencher, lhs, rhs, operator);
 }
 
+fn bench_decimal(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, operator: Operator) {
+    bench_binary::<DecimalArray>(bencher, lhs, rhs, operator);
+}
+
 fn bench_bool(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, operator: Operator) {
     bench_binary::<BoolArray>(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::<i64, _>((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::<i128, _>(
+        (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))
 }

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::<i64, _>([100, 250], DecimalDType::new(10, 2)).into_array(),
+            DecimalArray::from_iter::<i64, _>([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(),

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(),
+        );
+    }
 }

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<W: NativeDecimalType>(array: &DecimalArray) -> Buffer<W> {
+    if array.values_type() == W::DECIMAL_TYPE {
+        return array.buffer::<W>();
+    }
+    match_each_decimal_value_type!(array.values_type(), |T| {
+        array
+            .buffer::<T>()
+            .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),*) => {{

vortex-array/src/compute/conformance/binary_numeric.rs

@@ -40,11 +40,16 @@ 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;
+use crate::scalar_fn::fns::binary::decimal_add_sub_result_dtype;
 
 fn to_vec_of_scalar(array: &ArrayRef, ctx: &mut ExecutionCtx) -> Vec<Scalar> {
     // Not fast, but obviously correct
@@ -240,12 +245,144 @@ pub fn test_binary_numeric_array(array: &ArrayRef, ctx: &mut ExecutionCtx) {
             PType::F32 => test_binary_numeric_conformance::<f32>(array, ctx),
             PType::F64 => test_binary_numeric_conformance::<f64>(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 = <i256 as NativeDecimalType>::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::<Canonical>(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());
+    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];
+
+    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 expected_results: Vec<Option<Scalar>> = 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())
+                    };
+                    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();
+
+            let result = lhs
+                .binary(rhs, operator.into())
+                .vortex_expect("apply shouldn't fail")
+                .execute::<RecursiveCanonical>(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()
+            {
+                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,
+                );
+            }
+        }
+    }
+}
+
 /// Tests binary numeric operations with edge case scalar values.
 ///
 /// This function tests operations with scalar values: