From 4eab00aae65e7eda5f022c504f578afc55dff3b1 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 5 Oct 2026 18:13:48 +0100 Subject: [PATCH] Add bounded decimal dispatch macros and migrate kernels Signed-off-by: Joe Isaacs --- .../fns/all_non_distinct/decimal.rs | 29 ++-- .../src/aggregate_fn/fns/sum/decimal.rs | 22 ++- .../src/arrays/decimal/compute/cast.rs | 11 +- vortex-array/src/dtype/decimal/macros.rs | 136 ++++++++++++++++++ .../src/dtype/decimal/macros_tests.rs | 112 +++++++++++++++ 5 files changed, 273 insertions(+), 37 deletions(-) create mode 100644 vortex-array/src/dtype/decimal/macros_tests.rs diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/decimal.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/decimal.rs index 285d42cdfaa..eae102ed413 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_distinct/decimal.rs @@ -7,12 +7,9 @@ use vortex_error::VortexResult; use crate::arrays::DecimalArray; use crate::arrays::decimal::DecimalArrayExt; use crate::dtype::BigCast; +use crate::match_decimal_binary; use crate::match_each_decimal_value_type; -#[expect( - clippy::cognitive_complexity, - reason = "decimal widening depends on both source value types and the chosen widest type" -)] pub(super) fn check_decimal_identical( lhs: &DecimalArray, rhs: &DecimalArray, @@ -24,19 +21,15 @@ pub(super) fn check_decimal_identical( } let widest = lhs.values_type().max(rhs.values_type()); - match_each_decimal_value_type!(lhs.values_type(), |L| { - match_each_decimal_value_type!(rhs.values_type(), |R| { - match_each_decimal_value_type!(widest, |W| { - Ok(lhs - .buffer::() - .iter() - .zip(rhs.buffer::().iter()) - .all(|(lhs, rhs)| { - ::from(*lhs).vortex_expect("decimal widening should succeed") - == ::from(*rhs) - .vortex_expect("decimal widening should succeed") - })) - }) - }) + match_decimal_binary!(lhs.values_type(), rhs.values_type(), widest, |L, R, W| { + Ok(lhs + .buffer::() + .iter() + .zip(rhs.buffer::().iter()) + .all(|(lhs, rhs)| { + ::from(*lhs).vortex_expect("decimal widening should succeed") + == ::from(*rhs) + .vortex_expect("decimal widening should succeed") + })) }) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs index e30bd1f1dfc..84a65a0f4ba 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs @@ -19,7 +19,7 @@ use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::DecimalType; use crate::dtype::NativeDecimalType; -use crate::match_each_decimal_value_type; +use crate::match_decimal_unary; use crate::scalar::DecimalValue; /// Accumulate a decimal array into the sum state. @@ -47,17 +47,15 @@ pub(crate) fn accumulate_decimal( .vortex_expect("decimal sum result dtype"); let values_type = DecimalType::smallest_decimal_value_type(dtype); - match_each_decimal_value_type!(d.values_type(), |T| { - match_each_decimal_value_type!(values_type, |I| { - let initial: I = value - .cast() - .vortex_expect("cannot fail to cast initial value"); - match sum_decimal_value(initial, d.buffer::(), validity, *dtype) { - Some(v) => *value = v, - None => return Ok(true), - } - Ok(false) - }) + match_decimal_unary!(d.values_type(), values_type, |T, I| { + let initial: I = value + .cast() + .vortex_expect("cannot fail to cast initial value"); + match sum_decimal_value(initial, d.buffer::(), validity, *dtype) { + Some(v) => *value = v, + None => return Ok(true), + } + Ok(false) }) } diff --git a/vortex-array/src/arrays/decimal/compute/cast.rs b/vortex-array/src/arrays/decimal/compute/cast.rs index b7e94fbbf66..4e48cb95548 100644 --- a/vortex-array/src/arrays/decimal/compute/cast.rs +++ b/vortex-array/src/arrays/decimal/compute/cast.rs @@ -33,6 +33,7 @@ use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::ToI256; use crate::dtype::i256; +use crate::match_decimal_unary; use crate::match_each_decimal_value_type; use crate::match_each_integer_ptype; use crate::scalar::DecimalToIntegerCast; @@ -461,13 +462,9 @@ pub fn upcast_decimal_values( } let validity = array.validity()?; - // Use match_each_decimal_value_type to dispatch based on source and target types - match_each_decimal_value_type!(from_values_type, |F| { - let from_buffer = array.buffer::(); - match_each_decimal_value_type!(to_values_type, |T| { - let to_buffer = upcast_decimal_buffer::(from_buffer); - Ok(DecimalArray::new(to_buffer, decimal_dtype, validity)) - }) + match_decimal_unary!(from_values_type, to_values_type, |F, T| { + let to_buffer = upcast_decimal_buffer::(array.buffer::()); + Ok(DecimalArray::new(to_buffer, decimal_dtype, validity)) }) } diff --git a/vortex-array/src/dtype/decimal/macros.rs b/vortex-array/src/dtype/decimal/macros.rs index 39629ac285e..5622492b6ed 100644 --- a/vortex-array/src/dtype/decimal/macros.rs +++ b/vortex-array/src/dtype/decimal/macros.rs @@ -58,3 +58,139 @@ macro_rules! match_each_decimal_value_type { } }}; } + +/// Dispatches a decimal input and output type, instantiating only `In <= Out`. +/// This produces 21 type combinations instead of 36. +/// +/// The caller must choose an output type at least as wide as the input storage type. +/// A decimal array's storage is no wider than +/// [`DecimalType::smallest_decimal_value_type`](crate::dtype::DecimalType::smallest_decimal_value_type) +/// for its precision, so using that type (or a wider working type) satisfies this contract. +/// Narrowing casts must use unrestricted dispatch instead. Invalid combinations panic. +/// Each type expression is evaluated once. +/// +/// ``` +/// use num_traits::AsPrimitive; +/// use vortex_array::dtype::DecimalType; +/// use vortex_array::match_decimal_unary; +/// +/// let value = match_decimal_unary!(DecimalType::I8, DecimalType::I64, |In, Out| { +/// let input: In = 42i8.as_(); +/// >::as_(input).to_string() +/// }); +/// assert_eq!(value, "42"); +/// ``` +#[macro_export] +macro_rules! match_decimal_unary { + ($input:expr, $output:expr, | $in:ident, $out:ident | $body:block) => {{ + let (input, output) = ($input, $output); + match output { + $crate::dtype::DecimalType::I8 => { + type $out = i8; + $crate::match_decimal_unary!(@input input, |$in| $body; I8 => i8) + } + $crate::dtype::DecimalType::I16 => { + type $out = i16; + $crate::match_decimal_unary!(@input input, |$in| $body; I8 => i8, I16 => i16) + } + $crate::dtype::DecimalType::I32 => { + type $out = i32; + $crate::match_decimal_unary!(@input input, |$in| $body; I8 => i8, I16 => i16, I32 => i32) + } + $crate::dtype::DecimalType::I64 => { + type $out = i64; + $crate::match_decimal_unary!(@input input, |$in| $body; I8 => i8, I16 => i16, I32 => i32, I64 => i64) + } + $crate::dtype::DecimalType::I128 => { + type $out = i128; + $crate::match_decimal_unary!(@input input, |$in| $body; I8 => i8, I16 => i16, I32 => i32, I64 => i64, I128 => i128) + } + $crate::dtype::DecimalType::I256 => { + type $out = $crate::dtype::i256; + $crate::match_decimal_unary!(@input input, |$in| $body; I8 => i8, I16 => i16, I32 => i32, I64 => i64, I128 => i128, I256 => $crate::dtype::i256) + } + } + }}; + (@input $input:expr, | $in:ident | $body:block; $( $variant:ident => $native:ty ),+) => {{ + #[allow(unreachable_patterns)] + match $input { + $( + $crate::dtype::DecimalType::$variant => { + type $in = $native; + $body + } + )+ + _ => panic!("decimal input storage type must be no wider than the output type"), + } + }}; +} + +/// Dispatches two decimal input types and an output type, instantiating only combinations +/// where both inputs are no wider than the output (`Left <= Out` and `Right <= Out`). +/// +/// This produces 91 type combinations instead of the 216 from unrestricted nested dispatch. +/// See [`match_decimal_unary!`] for the storage contract. Invalid combinations panic. +/// Each type expression is evaluated once, from left to right. +/// +/// ``` +/// use num_traits::AsPrimitive; +/// use vortex_array::dtype::DecimalType; +/// use vortex_array::match_decimal_binary; +/// +/// let value = match_decimal_binary!( +/// DecimalType::I8, DecimalType::I16, DecimalType::I64, |Left, Right, Out| { +/// let lhs: Left = 1i8.as_(); +/// let rhs: Right = 2i8.as_(); +/// (>::as_(lhs) + >::as_(rhs)).to_string() +/// } +/// ); +/// assert_eq!(value, "3"); +/// ``` +#[macro_export] +macro_rules! match_decimal_binary { + ($left:expr, $right:expr, $output:expr, | $lhs:ident, $rhs:ident, $out:ident | $body:block) => {{ + let (left, right, output) = ($left, $right, $output); + match output { + $crate::dtype::DecimalType::I8 => { + type $out = i8; + $crate::match_decimal_unary!(@input left, |$lhs| { + $crate::match_decimal_unary!(@input right, |$rhs| $body; I8 => i8) + }; I8 => i8) + } + $crate::dtype::DecimalType::I16 => { + type $out = i16; + $crate::match_decimal_unary!(@input left, |$lhs| { + $crate::match_decimal_unary!(@input right, |$rhs| $body; I8 => i8, I16 => i16) + }; I8 => i8, I16 => i16) + } + $crate::dtype::DecimalType::I32 => { + type $out = i32; + $crate::match_decimal_unary!(@input left, |$lhs| { + $crate::match_decimal_unary!(@input right, |$rhs| $body; I8 => i8, I16 => i16, I32 => i32) + }; I8 => i8, I16 => i16, I32 => i32) + } + $crate::dtype::DecimalType::I64 => { + type $out = i64; + $crate::match_decimal_unary!(@input left, |$lhs| { + $crate::match_decimal_unary!(@input right, |$rhs| $body; I8 => i8, I16 => i16, I32 => i32, I64 => i64) + }; I8 => i8, I16 => i16, I32 => i32, I64 => i64) + } + $crate::dtype::DecimalType::I128 => { + type $out = i128; + $crate::match_decimal_unary!(@input left, |$lhs| { + $crate::match_decimal_unary!(@input right, |$rhs| $body; I8 => i8, I16 => i16, I32 => i32, I64 => i64, I128 => i128) + }; I8 => i8, I16 => i16, I32 => i32, I64 => i64, I128 => i128) + } + $crate::dtype::DecimalType::I256 => { + type $out = $crate::dtype::i256; + $crate::match_decimal_unary!(@input left, |$lhs| { + $crate::match_decimal_unary!(@input right, |$rhs| $body; I8 => i8, I16 => i16, I32 => i32, I64 => i64, I128 => i128, I256 => $crate::dtype::i256) + }; I8 => i8, I16 => i16, I32 => i32, I64 => i64, I128 => i128, I256 => $crate::dtype::i256) + } + } + }}; +} + +#[cfg(test)] +#[path = "macros_tests.rs"] +mod tests; diff --git a/vortex-array/src/dtype/decimal/macros_tests.rs b/vortex-array/src/dtype/decimal/macros_tests.rs new file mode 100644 index 00000000000..68cfe26979a --- /dev/null +++ b/vortex-array/src/dtype/decimal/macros_tests.rs @@ -0,0 +1,112 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::cell::Cell; +use std::panic::catch_unwind; + +use num_traits::AsPrimitive; + +use crate::dtype::DecimalType; +use crate::dtype::NativeDecimalType; +use crate::dtype::i256; + +const TYPES: [DecimalType; 6] = [ + DecimalType::I8, + DecimalType::I16, + DecimalType::I32, + DecimalType::I64, + DecimalType::I128, + DecimalType::I256, +]; + +// Only widening pairs implement this trait. Using it in the macro bodies makes instantiating +// even one narrowing combination a compile error, independently of the runtime test inputs. +trait Widen { + fn widen(self) -> Out; +} + +macro_rules! impl_widen { + ($out:ty; $( $input:ty ),+) => { + $( + impl Widen<$out> for $input { + fn widen(self) -> $out { + self.as_() + } + } + )+ + }; +} + +impl_widen!(i8; i8); +impl_widen!(i16; i8, i16); +impl_widen!(i32; i8, i16, i32); +impl_widen!(i64; i8, i16, i32, i64); +impl_widen!(i128; i8, i16, i32, i64, i128); +impl_widen!(i256; i8, i16, i32, i64, i128, i256); + +#[test] +fn unary_dispatches_only_widening_pairs() { + for input in TYPES { + for output in TYPES { + let result = catch_unwind(|| { + match_decimal_unary!(input, output, |In, Out| { + let value: In = (-42i8).as_(); + let widened: Out = value.widen(); + assert_eq!(In::DECIMAL_TYPE, input); + assert_eq!(Out::DECIMAL_TYPE, output); + widened.to_string() + }) + }); + if input <= output { + assert_eq!(result.unwrap(), "-42"); + } else { + assert!(result.is_err()); + } + } + } +} + +#[test] +fn binary_dispatches_only_widening_triples() { + for left in TYPES { + for right in TYPES { + for output in TYPES { + let result = catch_unwind(|| { + match_decimal_binary!(left, right, output, |Left, Right, Out| { + let lhs: Left = (-42i8).as_(); + let rhs: Right = 43i8.as_(); + let lhs: Out = lhs.widen(); + let rhs: Out = rhs.widen(); + assert_eq!(Left::DECIMAL_TYPE, left); + assert_eq!(Right::DECIMAL_TYPE, right); + assert_eq!(Out::DECIMAL_TYPE, output); + (lhs + rhs).to_string() + }) + }); + if left <= output && right <= output { + assert_eq!(result.unwrap(), "1"); + } else { + assert!(result.is_err()); + } + } + } + } +} + +#[test] +fn dispatch_evaluates_type_expressions_once_in_order() { + let calls = Cell::new(0); + let next = |expected| { + assert_eq!(calls.get(), expected); + calls.set(expected + 1); + DecimalType::I8 + }; + match_decimal_unary!(next(0), next(1), |In, Out| { + assert_eq!(In::DECIMAL_TYPE, Out::DECIMAL_TYPE); + }); + match_decimal_binary!(next(2), next(3), next(4), |Left, Right, Out| { + assert_eq!(Left::DECIMAL_TYPE, Out::DECIMAL_TYPE); + assert_eq!(Right::DECIMAL_TYPE, Out::DECIMAL_TYPE); + }); + assert_eq!(calls.get(), 5); +}