From c25e1d97f2870c66d0e413ffcf8c6c44b8423614 Mon Sep 17 00:00:00 2001 From: Onur Satici Date: Tue, 1 Sep 2026 15:12:23 +0100 Subject: [PATCH] Support mixed decimal multiplication Derive multiplication result types from both decimal operands and accept signed primitive integers without casting input arrays. Execute directly over canonical integer buffers while preserving checked decimal overflow and null semantics. Signed-off-by: Onur Satici --- .../scalar/typed_view/decimal/arithmetic.rs | 62 ++-- .../src/scalar/typed_view/decimal/mod.rs | 1 + vortex-array/src/scalar_fn/fns/binary/mod.rs | 19 +- .../scalar_fn/fns/binary/numeric/decimal.rs | 308 +++++++++++++----- .../src/scalar_fn/fns/binary/numeric/mod.rs | 105 ++++-- .../src/scalar_fn/fns/binary/numeric/tests.rs | 207 ++++++++++++ 6 files changed, 535 insertions(+), 167 deletions(-) diff --git a/vortex-array/src/scalar/typed_view/decimal/arithmetic.rs b/vortex-array/src/scalar/typed_view/decimal/arithmetic.rs index 2a47c9cf977..9dea76b1a48 100644 --- a/vortex-array/src/scalar/typed_view/decimal/arithmetic.rs +++ b/vortex-array/src/scalar/typed_view/decimal/arithmetic.rs @@ -3,8 +3,9 @@ //! Arrow's decimal arithmetic rules, and their evaluation over [`DecimalValue`]. //! -//! Vortex coerces both operands of a decimal arithmetic expression to a single [`DecimalDType`], -//! so Arrow's general `(p1, s1) op (p2, s2)` formulas collapse to a function of one input type: +//! Add, Sub, and Div currently require both operands to share one [`DecimalDType`]. Decimal scalar +//! arithmetic has the same constraint. For those operations, and for same-dtype Mul, Arrow's +//! general `(p1, s1) op (p2, s2)` formulas collapse to a function of one input type: //! //! | operator | result precision | result scale | //! | -------- | ---------------- | ------------ | @@ -16,6 +17,9 @@ //! stored integers already sits at the result scale and no rounding is needed. Div instead scales //! the dividend up front and truncates toward zero, which is what Arrow does. //! +//! Array multiplication can also use distinct decimal dtypes. Its general `p1 + p2 + 1, +//! s1 + s2` result is derived by [`decimal_multiply_result_dtype`]. +//! //! Div is the one operator whose intermediate can outgrow the widest native width: scaling the //! dividend by `10^result_scale` overflows `i256` once `p + result_scale` passes //! [`MAX_PRECISION`], even for a quotient that would have fit the result precision. That is @@ -59,27 +63,7 @@ pub(crate) fn decimal_numeric_result_dtype( input.scale(), )) } - NumericOperator::Mul => { - // Doubling the scale in i8 would saturate a very negative sum into a legal-looking - // scale, so widen first. The SQL standard rejects a product whose scale cannot be - // represented rather than rounding it away. - let result_scale = >::from(input.scale()) * 2; - let Some(result_scale) = i8::try_from(result_scale) - .ok() - .filter(|scale| *scale <= MAX_SCALE) - else { - vortex_bail!( - "output scale {result_scale} of {input} {op} {input} is outside the \ - representable scale range of {} to {MAX_SCALE}", - i8::MIN - ); - }; - let result_precision = input - .precision() - .saturating_add(input.precision().saturating_add(1)) - .min(MAX_PRECISION); - DecimalDType::try_new(result_precision, result_scale) - } + NumericOperator::Mul => decimal_multiply_result_dtype(input, input), NumericOperator::Div => { // Arrow follows Postgres and MySQL in adding a fixed four fractional digits. Its // precision formula `p1 - s1 + s2 + result_scale` simplifies to `p + result_scale` @@ -98,6 +82,38 @@ pub(crate) fn decimal_numeric_result_dtype( } } +/// Derive the result decimal dtype of multiplying operands with decimal dtypes `lhs` and `rhs`. +/// +/// Precision follows Arrow's `p1 + p2 + 1` rule and saturates at [`MAX_PRECISION`]. Scale is the +/// exact sum `s1 + s2`, so the product of the stored integers needs no rescaling. +/// +/// # Errors +/// +/// Returns an error if the summed scale is outside Vortex's representable scale range. +pub(crate) fn decimal_multiply_result_dtype( + lhs: DecimalDType, + rhs: DecimalDType, +) -> VortexResult { + // Add in i16 so a very negative sum cannot wrap or saturate into a legal-looking i8 scale. + let result_scale = >::from(lhs.scale()) + >::from(rhs.scale()); + let Some(result_scale) = i8::try_from(result_scale) + .ok() + .filter(|scale| *scale <= MAX_SCALE) + else { + vortex_bail!( + "output scale {result_scale} of {lhs} * {rhs} is outside the representable scale \ + range of {} to {MAX_SCALE}", + i8::MIN + ); + }; + let result_precision = lhs + .precision() + .saturating_add(rhs.precision()) + .saturating_add(1) + .min(MAX_PRECISION); + DecimalDType::try_new(result_precision, result_scale) +} + /// Apply `op` to two stored values of `input`, returning the result stored in `result`'s width. /// /// Returns `None` if the result overflows `result`'s precision, or if `op` is a division by zero. diff --git a/vortex-array/src/scalar/typed_view/decimal/mod.rs b/vortex-array/src/scalar/typed_view/decimal/mod.rs index 149fc8aeadf..ba5390abd54 100644 --- a/vortex-array/src/scalar/typed_view/decimal/mod.rs +++ b/vortex-array/src/scalar/typed_view/decimal/mod.rs @@ -7,6 +7,7 @@ mod arithmetic; mod dvalue; mod scalar; +pub(crate) use arithmetic::decimal_multiply_result_dtype; pub(crate) use arithmetic::decimal_numeric_result_dtype; pub(crate) use arithmetic::decimal_numeric_work_dtype; pub use dvalue::DecimalValue; diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index ed31cdf8d46..dd1cd95a383 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -124,24 +124,7 @@ impl ScalarFnVTable for Binary { let rhs = &arg_dtypes[1]; if operator.is_arithmetic() { - 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 - && lhs.eq_ignore_nullability(rhs) - { - let numeric_op = NumericOperator::try_from(*operator)?; - return Ok(DType::Decimal( - numeric_op_result_decimal_dtype(*decimal_dtype, numeric_op)?, - lhs.nullability() | rhs.nullability(), - )); - } - vortex_bail!( - "incompatible types for arithmetic operation: {} {}", - lhs, - rhs - ); + return numeric_return_dtype(lhs, rhs, NumericOperator::try_from(*operator)?); } if operator.is_comparison() diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/decimal.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/decimal.rs index fa9ffb0b5f1..a9221bf2a98 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/decimal.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/decimal.rs @@ -3,11 +3,11 @@ //! Native execution of the arithmetic operators over decimal arrays. //! -//! 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. Mul takes the raw -//! product, which the doubled result scale leaves correctly scaled, and Div rescales the dividend -//! (or the divisor, for a negative result scale) before integer division. Result precision and -//! scale follow Arrow's rules — see [`decimal_numeric_result_dtype`]. +//! Add, Sub, and Div operate on decimal operands sharing one logical [`DecimalDType`]. Mul also +//! accepts decimals with different logical dtypes and signed integer operands. It multiplies their +//! stored integers directly because the result scale is the sum of the operand scales. Div rescales +//! the dividend (or the divisor, for a negative result scale) before integer division. Result +//! precision and scale follow Arrow's rules. //! //! Lanes execute in a working width wide enough that in-precision inputs cannot spuriously //! overflow an intermediate, then narrow to the result's own storage width. Every lane is still @@ -31,6 +31,7 @@ use vortex_mask::Mask; use super::checked::checked_lanes; use crate::ArrayRef; +use crate::Canonical; use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; @@ -43,31 +44,35 @@ use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::DecimalType; use crate::dtype::NativeDecimalType; +use crate::dtype::PType; use crate::match_each_decimal_value_type; use crate::scalar::DecimalValue; use crate::scalar::NumericOperator; use crate::scalar::Scalar; -use crate::scalar::decimal_numeric_result_dtype; use crate::scalar::decimal_numeric_work_dtype; use crate::validity::Validity; -/// Execute a numeric operation between two decimal arrays sharing a decimal dtype. +/// Execute a numeric operation whose result is decimal. pub(super) fn execute_numeric_decimal( lhs: &ArrayRef, rhs: &ArrayRef, op: NumericOperator, + result_decimal_dtype: DecimalDType, ctx: &mut ExecutionCtx, ) -> VortexResult { - let decimal_dtype = lhs - .dtype() - .as_decimal_opt() - .vortex_expect("inputs are both decimals"); - - let result_decimal_dtype = decimal_numeric_result_dtype(*decimal_dtype, op)?; let result_dtype = DType::Decimal( result_decimal_dtype, lhs.dtype().nullability() | rhs.dtype().nullability(), ); + let work_dtype = if op == NumericOperator::Mul { + DecimalDType::new(result_decimal_dtype.precision(), 0) + } else { + let input_dtype = lhs + .dtype() + .as_decimal_opt() + .vortex_expect("non-multiplication decimal operands share a decimal dtype"); + decimal_numeric_work_dtype(*input_dtype, result_decimal_dtype, op) + }; // Fast path for null constant arrays. if is_null_constant(lhs) || is_null_constant(rhs) { @@ -86,7 +91,6 @@ pub(super) fn execute_numeric_decimal( let validity = lhs.validity().and(rhs.validity())?; let valid_rows = validity.execute_mask(len, ctx)?; - let work_dtype = decimal_numeric_work_dtype(*decimal_dtype, result_decimal_dtype, op); match_each_decimal_value_type!(DecimalType::smallest_decimal_value_type(&work_dtype), |W| { let constants = DecimalOpConstants::::new(result_decimal_dtype, op)?; macro_rules! execute_typed { @@ -122,10 +126,10 @@ fn null_result(dtype: &DType, len: usize) -> ArrayRef { ConstantArray::new(Scalar::null(dtype.clone()), len).into_array() } -/// A decimal binary-operator operand: a canonical decimal array or a non-null constant. +/// An integer-valued decimal operator operand, retaining the input's physical buffer. enum DecimalOperand { Array { - values: DecimalArray, + values: Canonical, validity: Validity, }, Constant { @@ -140,7 +144,7 @@ impl DecimalOperand { let columnar = array.clone().execute::(ctx)?; match columnar { - Columnar::Constant(array) => match array.scalar().as_decimal().decimal_value() { + Columnar::Constant(array) => match integer_scalar_value(array.scalar())? { Some(value) => Ok(Some(Self::Constant { value, len: array.len(), @@ -152,9 +156,14 @@ impl DecimalOperand { })), None => Ok(None), }, - Columnar::Canonical(array) => { - let values = array.as_decimal().to_owned(); - let validity = values.validity()?; + Columnar::Canonical(values) => { + let validity = match &values { + Canonical::Decimal(values) => values.validity()?, + Canonical::Primitive(values) if values.ptype().is_signed_int() => { + values.validity()? + } + _ => unreachable!("unsupported decimal operand dtype {}", values.dtype()), + }; Ok(Some(Self::Array { values, validity })) } } @@ -174,6 +183,29 @@ impl DecimalOperand { } } +fn integer_scalar_value(scalar: &Scalar) -> VortexResult> { + match scalar.dtype() { + DType::Decimal(..) => Ok(scalar.as_decimal().decimal_value()), + DType::Primitive(PType::I8, _) => Ok(scalar + .as_primitive() + .try_typed_value::()? + .map(DecimalValue::from)), + DType::Primitive(PType::I16, _) => Ok(scalar + .as_primitive() + .try_typed_value::()? + .map(DecimalValue::from)), + DType::Primitive(PType::I32, _) => Ok(scalar + .as_primitive() + .try_typed_value::()? + .map(DecimalValue::from)), + DType::Primitive(PType::I64, _) => Ok(scalar + .as_primitive() + .try_typed_value::()? + .map(DecimalValue::from)), + dtype => unreachable!("unsupported decimal operand dtype {dtype}"), + } +} + /// Per-execution bounds for checked decimal lane operations at working width `W`. /// /// Native-width checked arithmetic only detects overflow of `W`, whose range may exceed the @@ -318,6 +350,53 @@ impl CheckedDecimalOp for CheckedDecimalDiv { } } +macro_rules! with_decimal_operand_values { + ($operand:expr, | $values:ident | $body:expr) => {{ + match $operand { + DecimalOperand::Array { + values: Canonical::Decimal(values), + .. + } => { + match_each_decimal_value_type!(values.values_type(), |T| { + let buffer = values.buffer::(); + let $values = buffer.as_slice(); + $body + }) + } + DecimalOperand::Array { + values: Canonical::Primitive(values), + .. + } => match values.ptype() { + PType::I8 => { + let buffer = values.to_buffer::(); + let $values = buffer.as_slice(); + $body + } + PType::I16 => { + let buffer = values.to_buffer::(); + let $values = buffer.as_slice(); + $body + } + PType::I32 => { + let buffer = values.to_buffer::(); + let $values = buffer.as_slice(); + $body + } + PType::I64 => { + let buffer = values.to_buffer::(); + let $values = buffer.as_slice(); + $body + } + ptype => unreachable!("unsupported decimal operand ptype {ptype}"), + }, + DecimalOperand::Array { values, .. } => { + unreachable!("unsupported decimal operand dtype {}", values.dtype()) + } + DecimalOperand::Constant { .. } => unreachable!("operand is an array"), + } + }}; +} + fn execute_decimal_typed( lhs: &DecimalOperand, rhs: &DecimalOperand, @@ -334,47 +413,27 @@ where { let len = lhs.len(); - let values = match (lhs, rhs) { - (DecimalOperand::Array { values: lhs, .. }, DecimalOperand::Array { values: rhs, .. }) => { - checked_decimal_arrays::(lhs, rhs, constants, valid_rows) - } - (DecimalOperand::Array { values: lhs, .. }, DecimalOperand::Constant { value, .. }) => { - let rhs = typed_constant::(value); - match_each_decimal_value_type!(lhs.values_type(), |L| { - let lhs = lhs.buffer::(); - checked_lanes(lhs.as_slice(), valid_rows, |lhs| { - Op::apply(::from(lhs)?, rhs, constants) - }) - }) - } - (DecimalOperand::Constant { value, .. }, DecimalOperand::Array { values: rhs, .. }) => { - let lhs = typed_constant::(value); - match_each_decimal_value_type!(rhs.values_type(), |R| { - let rhs = rhs.buffer::(); - checked_lanes(rhs.as_slice(), valid_rows, |rhs| { - Op::apply(lhs, ::from(rhs)?, constants) - }) - }) - } - ( - DecimalOperand::Constant { value: lhs, .. }, - DecimalOperand::Constant { value: rhs, .. }, - ) => { - let lhs = typed_constant::(lhs); - let rhs = typed_constant::(rhs); - let value = Op::apply(lhs, rhs, constants) - .ok_or_else(|| vortex_err!(InvalidArgument: "{}", Op::ERROR))?; - let value = DecimalValue::from(value) - .normalize(result_decimal_dtype) - .vortex_expect("bounds-checked result fits the result precision"); - return Ok(ConstantArray::new( - Scalar::decimal(value, result_decimal_dtype, result_dtype.nullability()), - len, - ) - .into_array()); - } + if let ( + DecimalOperand::Constant { value: lhs, .. }, + DecimalOperand::Constant { value: rhs, .. }, + ) = (lhs, rhs) + { + let lhs = typed_constant::(lhs); + let rhs = typed_constant::(rhs); + let value = Op::apply(lhs, rhs, constants) + .ok_or_else(|| vortex_err!(InvalidArgument: "{}", Op::ERROR))?; + let value = DecimalValue::from(value) + .normalize(result_decimal_dtype) + .vortex_expect("bounds-checked result fits the result precision"); + return Ok(ConstantArray::new( + Scalar::decimal(value, result_decimal_dtype, result_dtype.nullability()), + len, + ) + .into_array()); } - .map_err(|_lane| vortex_err!(InvalidArgument: "{}", Op::ERROR))?; + + let values = execute_decimal_lanes::(lhs, rhs, constants, valid_rows) + .map_err(|_lane| vortex_err!(InvalidArgument: "{}", Op::ERROR))?; Ok(decimal_array_narrowed( values, @@ -383,6 +442,105 @@ where )) } +fn execute_decimal_lanes( + lhs: &DecimalOperand, + rhs: &DecimalOperand, + constants: &DecimalOpConstants, + valid_rows: &Mask, +) -> Result, usize> +where + W: NativeDecimalType + CheckedAdd + CheckedSub + CheckedMul + CheckedDiv + Mul, + Op: CheckedDecimalOp, +{ + match (lhs, rhs) { + (DecimalOperand::Constant { value, .. }, DecimalOperand::Array { .. }) => { + execute_decimal_constant_array::(value, rhs, constants, valid_rows) + } + (DecimalOperand::Array { .. }, DecimalOperand::Constant { value, .. }) => { + execute_decimal_array_constant::(lhs, value, constants, valid_rows) + } + (DecimalOperand::Array { .. }, DecimalOperand::Array { .. }) => { + execute_decimal_array_array::(lhs, rhs, constants, valid_rows) + } + (DecimalOperand::Constant { .. }, DecimalOperand::Constant { .. }) => { + unreachable!("constant operands are handled before lane execution") + } + } +} + +fn execute_decimal_constant_array( + lhs: &DecimalValue, + rhs: &DecimalOperand, + constants: &DecimalOpConstants, + valid_rows: &Mask, +) -> Result, usize> +where + W: NativeDecimalType + CheckedAdd + CheckedSub + CheckedMul + CheckedDiv + Mul, + Op: CheckedDecimalOp, +{ + let lhs = typed_constant::(lhs); + with_decimal_operand_values!(rhs, |rhs| { + checked_lanes(rhs, valid_rows, |rhs| { + Op::apply(lhs, ::from(rhs)?, constants) + }) + }) +} + +fn execute_decimal_array_constant( + lhs: &DecimalOperand, + rhs: &DecimalValue, + constants: &DecimalOpConstants, + valid_rows: &Mask, +) -> Result, usize> +where + W: NativeDecimalType + CheckedAdd + CheckedSub + CheckedMul + CheckedDiv + Mul, + Op: CheckedDecimalOp, +{ + let rhs = typed_constant::(rhs); + with_decimal_operand_values!(lhs, |lhs| { + checked_lanes(lhs, valid_rows, |lhs| { + Op::apply(::from(lhs)?, rhs, constants) + }) + }) +} + +fn execute_decimal_array_array( + lhs: &DecimalOperand, + rhs: &DecimalOperand, + constants: &DecimalOpConstants, + valid_rows: &Mask, +) -> Result, usize> +where + W: NativeDecimalType + CheckedAdd + CheckedSub + CheckedMul + CheckedDiv + Mul, + Op: CheckedDecimalOp, +{ + with_decimal_operand_values!(lhs, |lhs| { + execute_decimal_array_rhs::(lhs, rhs, constants, valid_rows) + }) +} + +fn execute_decimal_array_rhs( + lhs: &[L], + rhs: &DecimalOperand, + constants: &DecimalOpConstants, + valid_rows: &Mask, +) -> Result, usize> +where + W: NativeDecimalType + CheckedAdd + CheckedSub + CheckedMul + CheckedDiv + Mul, + Op: CheckedDecimalOp, + L: NativeDecimalType, +{ + with_decimal_operand_values!(rhs, |rhs| { + checked_lanes(LaneZip::new(lhs, rhs), valid_rows, |(lhs, rhs)| { + Op::apply( + ::from(lhs)?, + ::from(rhs)?, + constants, + ) + }) + }) +} + /// Build the result array, narrowing to the dtype's own storage width when the working width is /// wider than it. Only division picks a working width above the result precision, and only for a /// negative result scale, so this copies in a corner case rather than on the common path. @@ -410,36 +568,6 @@ fn decimal_array_narrowed( }) } -fn checked_decimal_arrays( - lhs: &DecimalArray, - rhs: &DecimalArray, - constants: &DecimalOpConstants, - valid_rows: &Mask, -) -> Result, usize> -where - W: NativeDecimalType + CheckedAdd + CheckedSub + CheckedMul + CheckedDiv + Mul, - Op: CheckedDecimalOp, -{ - debug_assert_eq!(lhs.len(), rhs.len()); - match_each_decimal_value_type!(lhs.values_type(), |L| { - let lhs = lhs.buffer::(); - match_each_decimal_value_type!(rhs.values_type(), |R| { - let rhs = rhs.buffer::(); - checked_lanes( - LaneZip::new(lhs.as_slice(), rhs.as_slice()), - valid_rows, - |(lhs, rhs)| { - Op::apply( - ::from(lhs)?, - ::from(rhs)?, - constants, - ) - }, - ) - }) - }) -} - fn typed_constant(value: &DecimalValue) -> W { value .cast::() diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs index a0c427b142e..0a3700b9cfe 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs @@ -18,6 +18,7 @@ mod row; use decimal::execute_numeric_decimal; use row::execute_numeric_primitive; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use crate::ArrayRef; @@ -25,9 +26,69 @@ use crate::Canonical; use crate::ExecutionCtx; use crate::IntoArray; use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::PType; use crate::scalar::NumericOperator; +use crate::scalar::decimal_multiply_result_dtype; pub(crate) use crate::scalar::decimal_numeric_result_dtype as numeric_op_result_decimal_dtype; +pub(super) fn numeric_return_dtype( + lhs: &DType, + rhs: &DType, + op: NumericOperator, +) -> VortexResult { + let nullability = lhs.nullability() | rhs.nullability(); + + if lhs.is_primitive() && lhs.eq_ignore_nullability(rhs) { + return Ok(lhs.with_nullability(nullability)); + } + + if op == NumericOperator::Mul + && (lhs.is_decimal() || rhs.is_decimal()) + && let (Some(lhs), Some(rhs)) = ( + decimal_multiply_operand_dtype(lhs), + decimal_multiply_operand_dtype(rhs), + ) + { + return Ok(DType::Decimal( + decimal_multiply_result_dtype(lhs, rhs)?, + nullability, + )); + } + + if let (DType::Decimal(lhs_decimal, _), DType::Decimal(rhs_decimal, _)) = (lhs, rhs) + && lhs_decimal == rhs_decimal + { + return Ok(DType::Decimal( + numeric_op_result_decimal_dtype(*lhs_decimal, op)?, + nullability, + )); + } + + vortex_bail!( + "incompatible types for arithmetic operation: {} {}", + lhs, + rhs + ) +} + +pub(super) fn decimal_multiply_operand_dtype(dtype: &DType) -> Option { + match dtype { + DType::Decimal(dtype, _) => Some(*dtype), + DType::Primitive(ptype, _) if ptype.is_signed_int() => { + let precision = match ptype { + PType::I8 => 3, + PType::I16 => 5, + PType::I32 => 10, + PType::I64 => 19, + _ => unreachable!("ptype is a signed integer"), + }; + Some(DecimalDType::new(precision, 0)) + } + _ => None, + } +} + /// Execute a numeric operation between two arrays. pub(crate) fn execute_numeric( lhs: &ArrayRef, @@ -35,20 +96,6 @@ pub(crate) fn execute_numeric( op: NumericOperator, ctx: &mut ExecutionCtx, ) -> VortexResult { - vortex_ensure!( - lhs.dtype().eq_ignore_nullability(rhs.dtype()), - "numeric operator requires matching types, got {} and {}", - lhs.dtype(), - rhs.dtype() - ); - - let dtype = lhs.dtype(); - vortex_ensure!( - matches!(dtype, DType::Primitive(..) | DType::Decimal(..)), - "numeric operator is not supported for dtype {}", - dtype - ); - vortex_ensure!( lhs.len() == rhs.len(), "numeric operator requires equal lengths, got {} and {}", @@ -56,34 +103,20 @@ pub(crate) fn execute_numeric( rhs.len() ); + let result_dtype = numeric_return_dtype(lhs.dtype(), rhs.dtype(), op)?; + if lhs.is_empty() { - return build_empty_result(lhs, rhs, op); + return Ok(Canonical::empty(&result_dtype).into_array()); } - match dtype { + match result_dtype { DType::Primitive(..) => execute_numeric_primitive(lhs, rhs, op, ctx), - DType::Decimal(..) => execute_numeric_decimal(lhs, rhs, op, ctx), - _ => unreachable!("dtype is either Primitive or Decimal"), + DType::Decimal(decimal_dtype, _) => { + execute_numeric_decimal(lhs, rhs, op, decimal_dtype, ctx) + } + _ => unreachable!("numeric result is either Primitive or Decimal"), } } -fn build_empty_result( - lhs: &ArrayRef, - rhs: &ArrayRef, - op: NumericOperator, -) -> VortexResult { - 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, _) => DType::Decimal( - numeric_op_result_decimal_dtype(*decimal_dtype, op)?, - nullability, - ), - _ => unreachable!("dtype is either Primitive or Decimal"), - }; - - Ok(Canonical::empty(&result_dtype).into_array()) -} - #[cfg(test)] mod tests; diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs index 3813c8612b3..e8f8519c59e 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs @@ -420,6 +420,213 @@ fn test_decimal_mixed_storage_widths() -> VortexResult<()> { Ok(()) } +#[test] +fn test_decimal_mul_different_dtypes() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let lhs = + DecimalArray::from_iter::([1_234, -200], DecimalDType::new(7, 2)).into_array(); + let rhs = DecimalArray::from_iter::([20, 30], DecimalDType::new(3, 1)).into_array(); + + let result = decimal_binary(lhs, rhs, Operator::Mul)?; + assert_arrays_eq!( + result, + DecimalArray::from_iter::([24_680, -6_000], DecimalDType::new(11, 3)), + &mut ctx + ); + Ok(()) +} + +#[test] +fn test_decimal_mul_i64_without_decimal_cast() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal = DecimalArray::from_option_iter::( + [Some(125), None, Some(-300)], + DecimalDType::new(7, 2), + ) + .into_array(); + let integer = PrimitiveArray::from_iter([4i64, i64::MAX, -2]).into_array(); + + let result = decimal_binary(decimal, integer, Operator::Mul)?; + assert_arrays_eq!( + result, + DecimalArray::from_option_iter::( + [Some(500), None, Some(600)], + DecimalDType::new(27, 2), + ), + &mut ctx + ); + Ok(()) +} + +#[test] +fn test_signed_integer_mul_decimal_is_commutative() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let integer = PrimitiveArray::from_iter([-2i16, 3]).into_array(); + let decimal = + DecimalArray::from_iter::([125, -200], DecimalDType::new(7, 2)).into_array(); + + let result = decimal_binary(integer, decimal, Operator::Mul)?; + assert_arrays_eq!( + result, + DecimalArray::from_iter::([-250, -600], DecimalDType::new(13, 2)), + &mut ctx + ); + Ok(()) +} + +#[test] +fn test_decimal_mul_i64_stays_in_i128_at_precision_38() -> VortexResult<()> { + let decimal = DecimalArray::from_iter::([10], DecimalDType::new(18, 2)).into_array(); + let integer = PrimitiveArray::from_iter([20i64]).into_array(); + + let result = decimal_binary(decimal, integer, Operator::Mul)? + .execute::(&mut array_session().create_execution_ctx())?; + + assert_eq!(result.decimal_dtype(), DecimalDType::new(38, 2)); + assert_eq!(result.values_type(), DecimalType::I128); + Ok(()) +} + +#[rstest] +#[case::i8( + PrimitiveArray::from_iter([2i8]).into_array(), + DecimalDType::new(11, 2) +)] +#[case::i16( + PrimitiveArray::from_iter([2i16]).into_array(), + DecimalDType::new(13, 2) +)] +#[case::i32( + PrimitiveArray::from_iter([2i32]).into_array(), + DecimalDType::new(18, 2) +)] +#[case::i64( + PrimitiveArray::from_iter([2i64]).into_array(), + DecimalDType::new(27, 2) +)] +fn test_decimal_mul_all_signed_integer_widths( + #[case] integer: ArrayRef, + #[case] result_dtype: DecimalDType, +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal = DecimalArray::from_iter::([125], DecimalDType::new(7, 2)).into_array(); + + let result = decimal_binary(decimal, integer, Operator::Mul)?; + assert_arrays_eq!( + result, + DecimalArray::from_iter::([250], result_dtype), + &mut ctx + ); + Ok(()) +} + +#[test] +fn test_decimal_mul_signed_integer_null_constant() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal = + DecimalArray::from_iter::([125, 250], DecimalDType::new(7, 2)).into_array(); + let null_integer = ConstantArray::new(Option::::None, decimal.len()).into_array(); + + let result = decimal + .binary(null_integer, Operator::Mul)? + .execute::(&mut ctx)?; + assert!(matches!(&result, Columnar::Constant(_))); + assert_arrays_eq!( + result.into_array(), + DecimalArray::from_option_iter::([None, None], DecimalDType::new(27, 2),), + &mut ctx + ); + Ok(()) +} + +#[test] +fn test_decimal_mul_signed_integer_constants() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal = decimal_constant(125i32, DecimalDType::new(7, 2), 2); + let integer = ConstantArray::new(-2i64, 2).into_array(); + + let result = decimal + .binary(integer, Operator::Mul)? + .execute::(&mut ctx)?; + assert!(matches!(&result, Columnar::Constant(_))); + assert_arrays_eq!( + result.into_array(), + DecimalArray::from_iter::([-250, -250], DecimalDType::new(27, 2)), + &mut ctx + ); + Ok(()) +} + +#[test] +fn test_decimal_mul_signed_integer_empty() -> VortexResult<()> { + let decimal = DecimalArray::from_iter::([], DecimalDType::new(7, 2)).into_array(); + let integer = PrimitiveArray::from_iter(Vec::::new()).into_array(); + + let result = decimal_binary(decimal, integer, Operator::Mul)?; + + assert!(result.is_empty()); + assert_eq!( + result.dtype(), + &DType::Decimal(DecimalDType::new(27, 2), Nullability::NonNullable) + ); + Ok(()) +} + +#[test] +fn test_decimal_mul_signed_integer_widens_to_i256() -> VortexResult<()> { + let decimal = DecimalArray::from_iter::([10], DecimalDType::new(19, 0)).into_array(); + let integer = PrimitiveArray::from_iter([20i64]).into_array(); + + let result = decimal_binary(decimal, integer, Operator::Mul)? + .execute::(&mut array_session().create_execution_ctx())?; + + assert_eq!(result.decimal_dtype(), DecimalDType::new(39, 0)); + assert_eq!(result.values_type(), DecimalType::I256); + Ok(()) +} + +#[test] +fn test_decimal_mul_signed_integer_overflow_errors() { + let dtype = DecimalDType::new(76, 0); + let max = ::MAX_BY_PRECISION[76]; + let decimal = DecimalArray::from_iter::([max], dtype).into_array(); + let integer = PrimitiveArray::from_iter([2i64]).into_array(); + + assert!(decimal_binary(decimal, integer, Operator::Mul).is_err()); +} + +#[test] +fn test_decimal_non_mul_and_unsigned_integer_remain_incompatible() { + let decimal = DecimalArray::from_iter::([100], DecimalDType::new(7, 2)).into_array(); + let other_decimal = + DecimalArray::from_iter::([100], DecimalDType::new(8, 2)).into_array(); + + assert!( + decimal + .binary( + PrimitiveArray::from_iter([2i64]).into_array(), + Operator::Add + ) + .is_err() + ); + assert!( + decimal + .binary( + PrimitiveArray::from_iter([2u64]).into_array(), + Operator::Mul + ) + .is_err() + ); + assert!( + other_decimal + .binary( + DecimalArray::from_iter::([100], DecimalDType::new(7, 2)).into_array(), + Operator::Add, + ) + .is_err() + ); +} + #[test] fn test_decimal_nullable_lanes() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx();