Match Arrow decimal Add/Sub result typing

Signed-off-by: Matt Katz <mhkatz97@gmail.com>

Author: mhk197

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<Scalar> {
     // 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();
 

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::<Cast>());
-        assert!(coerced.child(1).is::<Cast>());
-        assert_eq!(coerced.return_dtype(&scope)?, result_dtype);
-        Ok(())
-    }
-
     #[test]
     fn boolean_operators_no_coercion() -> VortexResult<()> {
         let scope = DType::Struct(

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,

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::<i8, i8>(
+    match DecimalType::smallest_decimal_value_type(&result_decimal_dtype) {
+        DecimalType::I8 => execute_decimal_at_widths::<i8, i8>(
             &lhs,
             &rhs,
             op,
-            decimal_dtype,
+            result_decimal_dtype,
             &result_dtype,
             validity,
             &valid_rows,
         ),
-        (DecimalType::I16, DecimalType::I8) => execute_decimal_at_widths::<i16, i8>(
+        DecimalType::I16 => execute_decimal_at_widths::<i16, i16>(
             &lhs,
             &rhs,
             op,
-            decimal_dtype,
+            result_decimal_dtype,
             &result_dtype,
             validity,
             &valid_rows,
         ),
-        (DecimalType::I16, DecimalType::I16) => execute_decimal_at_widths::<i16, i16>(
+        DecimalType::I32 => execute_decimal_at_widths::<i32, i32>(
             &lhs,
             &rhs,
             op,
-            decimal_dtype,
+            result_decimal_dtype,
             &result_dtype,
             validity,
             &valid_rows,
         ),
-        (DecimalType::I32, DecimalType::I32) => execute_decimal_at_widths::<i32, i32>(
+        DecimalType::I64 => execute_decimal_at_widths::<i64, i64>(
             &lhs,
             &rhs,
             op,
-            decimal_dtype,
+            result_decimal_dtype,
             &result_dtype,
             validity,
             &valid_rows,
         ),
-        (DecimalType::I64, DecimalType::I64) => execute_decimal_at_widths::<i64, i64>(
+        DecimalType::I128 => execute_decimal_at_widths::<i128, i128>(
             &lhs,
             &rhs,
             op,
-            decimal_dtype,
+            result_decimal_dtype,
             &result_dtype,
             validity,
             &valid_rows,
         ),
-        (DecimalType::I128, DecimalType::I128) => execute_decimal_at_widths::<i128, i128>(
+        DecimalType::I256 => execute_decimal_at_widths::<i256, i256>(
             &lhs,
             &rhs,
             op,
-            decimal_dtype,
+            result_decimal_dtype,
             &result_dtype,
             validity,
             &valid_rows,
         ),
-        (DecimalType::I256, DecimalType::I128) => execute_decimal_at_widths::<i256, i128>(
-            &lhs,
-            &rhs,
-            op,
-            decimal_dtype,
-            &result_dtype,
-            validity,
-            &valid_rows,
-        ),
-        (DecimalType::I256, DecimalType::I256) => execute_decimal_at_widths::<i256, i256>(
-            &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 = <i256 as NativeDecimalType>::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::<i8>().is_some() {
-        DecimalType::I8
-    } else if value.cast::<i16>().is_some() {
-        DecimalType::I16
-    } else if value.cast::<i32>().is_some() {
-        DecimalType::I32
-    } else if value.cast::<i64>().is_some() {
-        DecimalType::I64
-    } else if value.cast::<i128>().is_some() {
-        DecimalType::I128
-    } else {
-        DecimalType::I256
-    }
-}
-
 /// Per-execution constants for checked decimal lane operations at working width `W`.
 struct DecimalOpPlan<W> {
     /// Inclusive stored-value bounds implied by the result precision.
@@ -289,7 +243,7 @@ fn execute_decimal_at_widths<W, O>(
     lhs: &DecimalOperand,
     rhs: &DecimalOperand,
     op: NumericOperator,
-    decimal_dtype: DecimalDType,
+    result_decimal_dtype: DecimalDType,
     result_dtype: &DType,
     validity: Validity,
     valid_rows: &Mask,
@@ -303,15 +257,15 @@ where
         NumericOperator::Add => execute_decimal_typed::<W, O, DecimalAdd>(
             lhs,
             rhs,
-            decimal_dtype,
+            result_decimal_dtype,
             result_dtype,
             validity,
             valid_rows,
         ),
         NumericOperator::Sub => execute_decimal_typed::<W, O, DecimalSub>(
             lhs,
             rhs,
-            decimal_dtype,
+            result_decimal_dtype,
             result_dtype,
             validity,
             valid_rows,
@@ -326,7 +280,7 @@ where
 fn execute_decimal_typed<W, O, Op>(
     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::<W>::new(decimal_dtype);
+    let plan = DecimalOpPlan::<W>::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())

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