diff --git a/vortex-array/src/arrays/primitive/compute/cast.rs b/vortex-array/src/arrays/primitive/compute/cast.rs index c937d68d611..231c5e4f05b 100644 --- a/vortex-array/src/arrays/primitive/compute/cast.rs +++ b/vortex-array/src/arrays/primitive/compute/cast.rs @@ -310,17 +310,40 @@ where // signed carrier that can represent both the source and target. let carrier_type = primitive_decimal_carrier_type(array.ptype()).max(T::DECIMAL_TYPE); match_each_decimal_value_type!(carrier_type, |W| { - let scale_factor = decimal_scale_factor::(decimal_dtype.scale())?; - cast_integer_values_to_decimal_buffer(values, decimal_dtype, valid_values, |value| { - let value = ::from(value)?; - let quotient = value / scale_factor; - let value = (quotient * scale_factor == value).then_some(quotient)?; - let value = ::from(value)?; - decimal_value_fits_precision(value, decimal_dtype).then_some(value) - }) + // `S` and `T` fix the carrier, so only one arm can match. The const condition lets the + // compiler skip codegen of the kernel for the other arms. + if const { is_scaled_down_carrier::() } { + let scale_factor = decimal_scale_factor::(decimal_dtype.scale())?; + cast_integer_values_to_decimal_buffer(values, decimal_dtype, valid_values, |value| { + let value = ::from(value)?; + let quotient = value / scale_factor; + let value = (quotient * scale_factor == value).then_some(quotient)?; + let value = ::from(value)?; + decimal_value_fits_precision(value, decimal_dtype).then_some(value) + }) + } else { + unreachable!( + "{carrier_type} is not the scaled-down carrier for {}", + S::PTYPE + ) + } }) } +/// Whether `W` is the carrier that a scaled-down cast from `S` to `T` selects. +const fn is_scaled_down_carrier() -> bool +where + S: NativePType, + T: NativeDecimalType, + W: NativeDecimalType, +{ + let carrier = primitive_decimal_carrier_type(S::PTYPE) as u8; + let target = T::DECIMAL_TYPE as u8; + let selected = if carrier > target { carrier } else { target }; + + W::DECIMAL_TYPE as u8 == selected +} + fn cast_integer_values_to_decimal_buffer( values: &[S], decimal_dtype: DecimalDType, @@ -350,7 +373,11 @@ fn signed_primitive_decimal_type(ptype: PType) -> Option { } /// The smallest signed decimal physical type that represents every `ptype` value. -fn primitive_decimal_carrier_type(ptype: PType) -> DecimalType { +#[expect( + clippy::panic, + reason = "`unreachable!` with a message is not available in a const fn" +)] +const fn primitive_decimal_carrier_type(ptype: PType) -> DecimalType { match ptype { PType::I8 => DecimalType::I8, PType::U8 | PType::I16 => DecimalType::I16, @@ -358,7 +385,7 @@ fn primitive_decimal_carrier_type(ptype: PType) -> DecimalType { PType::U32 | PType::I64 => DecimalType::I64, PType::U64 => DecimalType::I128, PType::F16 | PType::F32 | PType::F64 => { - unreachable!("floating primitives are rejected before selecting a decimal carrier") + panic!("floating primitives are rejected before selecting a decimal carrier") } } } @@ -685,6 +712,7 @@ mod test { use crate::dtype::PType; use crate::dtype::i256; use crate::expr::stats::Stat; + use crate::match_each_decimal_value_type; use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -1011,6 +1039,39 @@ mod test { Ok(()) } + #[rstest] + fn cast_integer_to_negative_scale_decimal_every_width( + #[values( + PType::U8, PType::U16, PType::U32, PType::U64, PType::I8, PType::I16, PType::I32, + PType::I64 + )] + ptype: PType, + #[values(2, 4, 9, 18, 38, 76)] precision: u8, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let decimal_dtype = DecimalDType::new(precision, -1); + let casted = PrimitiveArray::from_iter([100u8, 30]) + .into_array() + .cast(DType::Primitive(ptype, Nullability::NonNullable))? + .cast(DType::Decimal(decimal_dtype, Nullability::NonNullable))? + .execute::(&mut ctx)?; + + assert_eq!( + casted.values_type(), + DecimalType::smallest_decimal_value_type(&decimal_dtype) + ); + match_each_decimal_value_type!(casted.values_type(), |T| { + let expected = [10i8, 3].map(::from); + let actual = casted + .buffer::() + .iter() + .map(|&value| Some(value)) + .collect::>(); + assert_eq!(actual, expected); + }); + Ok(()) + } + #[test] fn cast_zero_to_large_negative_scale_decimal() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx();