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