comments

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

Author: mhk197

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

@@ -49,7 +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;
+use crate::scalar_fn::fns::binary::numeric_op_result_decimal_dtype;
 
 fn to_vec_of_scalar(array: &ArrayRef, ctx: &mut ExecutionCtx) -> Vec<Scalar> {
     // Not fast, but obviously correct
@@ -300,13 +300,14 @@ 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];
 
     for operator in operators {
+        let result_decimal_dtype = numeric_op_result_decimal_dtype(decimal_dtype, operator)
+            .vortex_expect("decimal Add/Sub must have a result dtype");
+        let result_dtype = DType::Decimal(result_decimal_dtype, array.dtype().nullability());
         let rhs_const = ConstantArray::new(scalar.clone(), array.len()).into_array();
 
         for lhs_is_array in [true, false] {

vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs

@@ -5,6 +5,9 @@
 
 use std::fmt;
 
+use vortex_error::VortexError;
+use vortex_error::vortex_err;
+
 use crate::scalar_fn::fns::operators::Operator;
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -38,3 +41,17 @@ impl From<NumericOperator> for Operator {
         }
     }
 }
+
+impl TryFrom<Operator> for NumericOperator {
+    type Error = VortexError;
+
+    fn try_from(op: Operator) -> Result<Self, Self::Error> {
+        match op {
+            Operator::Add => Ok(NumericOperator::Add),
+            Operator::Sub => Ok(NumericOperator::Sub),
+            Operator::Mul => Ok(NumericOperator::Mul),
+            Operator::Div => Ok(NumericOperator::Div),
+            _ => Err(vortex_err!(InvalidArgument: "{op} is not a numeric operator")),
+        }
+    }
+}

vortex-array/src/scalar_fn/fns/binary/mod.rs

@@ -17,8 +17,6 @@ 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;
@@ -50,14 +48,6 @@ 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;
 
@@ -131,12 +121,13 @@ impl ScalarFnVTable for Binary {
             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)
             {
+                let numeric_op = NumericOperator::try_from(*operator)?;
                 return Ok(DType::Decimal(
-                    decimal_add_sub_result_dtype(*decimal_dtype),
+                    numeric_op_result_decimal_dtype(*decimal_dtype, numeric_op)?,
                     lhs.nullability() | rhs.nullability(),
                 ));
             }