diff --git a/encodings/fastlanes/src/bitpacking/compute/cast.rs b/encodings/fastlanes/src/bitpacking/compute/cast.rs index 8203e68a6b7..ff9cdf48336 100644 --- a/encodings/fastlanes/src/bitpacking/compute/cast.rs +++ b/encodings/fastlanes/src/bitpacking/compute/cast.rs @@ -6,16 +6,20 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::builders::PrimitiveBuilder; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; use vortex_array::dtype::PType; -use vortex_array::match_each_integer_ptype; +use vortex_array::match_each_unsigned_integer_ptype; use vortex_array::scalar_fn::fns::cast::CastKernel; use vortex_array::scalar_fn::fns::cast::CastReduce; use vortex_array::validity::Validity; use vortex_error::VortexResult; +use crate::BitPackedArray; use crate::bitpacking::BitPacked; use crate::bitpacking::BitWidthsView; use crate::bitpacking::array::BitPackedArrayExt; @@ -53,6 +57,39 @@ fn build_with_validity( .into_array()) } +/// The array with the unsigned ptype of the same width. +/// +/// The packed values and patches are never negative, so they keep their value. Only the patch +/// values are reinterpreted. The packed buffer is shared. +fn unsigned_view_of( + array: ArrayView<'_, BitPacked>, + bit_width: u8, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let ptype = array.dtype().as_ptype().to_unsigned(); + let patches = array + .patches() + .map(|patches| { + patches.map_values(|values| { + Ok(values + .execute::(ctx)? + .reinterpret_cast(ptype) + .into_array()) + }) + }) + .transpose()?; + + BitPacked::try_new( + array.packed().clone(), + ptype, + array.validity()?, + patches, + bit_width, + array.len(), + array.offset(), + ) +} + impl CastReduce for BitPacked { fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { let BitWidthsView::Global(bit_width) = array.bit_widths() else { @@ -107,16 +144,36 @@ impl CastKernel for BitPacked { .validity()? .cast_nullability(tgt_nullability, array.len(), ctx)?; - let result = match_each_integer_ptype!(tgt, |T| { + // Bit-packing rejects negative values, so sign extension of a packed value is the same as + // zero extension. A signed cast then gives the same bits as the unsigned cast of the same + // widths, and signed and unsigned casts share one set of kernels. + let unsigned_view; + let source = if src.is_signed_int() { + unsigned_view = unsigned_view_of(array, bit_width, ctx)?; + unsigned_view.as_view() + } else { + array + }; + + let result = match_each_unsigned_integer_ptype!(tgt.to_unsigned(), |T| { let mut builder = PrimitiveBuilder::::with_capacity_in( tgt_nullability, array.len(), ctx.allocator(), ); - match_each_integer_ptype!(src, |F| { - unpack_map_into_builder::(array, &mut builder, ctx, |v: F| v.as_())?; + match_each_unsigned_integer_ptype!(src.to_unsigned(), |F| { + // Only widening casts reach here. The const condition lets the compiler skip + // codegen of the other arms. + if const { F::PTYPE.byte_width() < T::PTYPE.byte_width() } { + unpack_map_into_builder::(source, &mut builder, ctx, |v: F| v.as_())?; + } else { + unreachable!("{src} to {tgt} is not a widening integer cast"); + } }); - builder.finish_into_primitive().into_array() + builder + .finish_into_primitive() + .reinterpret_cast(tgt) + .into_array() }); Ok(Some(result)) } @@ -144,6 +201,7 @@ mod tests { use vortex_session::VortexSession; use crate::BitPackedArray; + use crate::BitPackedArrayExt; use crate::BitPackedData; static SESSION: LazyLock = LazyLock::new(|| { @@ -176,6 +234,24 @@ mod tests { ); } + /// Signed sources are cast through an unsigned view, which reinterprets the patch values. + #[rstest] + #[case::i8_to_i16(PType::I16)] + #[case::i8_to_i64(PType::I64)] + fn test_cast_bitpacked_signed_with_patches(#[case] tgt: PType) -> VortexResult<()> { + let values = PrimitiveArray::from_iter( + (0..2000).map(|i| if i % 97 == 0 { 120i8 } else { (i % 7) as i8 }), + ); + let packed = bp(&values.clone().into_array(), 3); + assert!(packed.patches().is_some(), "test setup expects patches"); + + let target = DType::Primitive(tgt, Nullability::NonNullable); + let casted = packed.into_array().cast(target.clone())?; + let expected = values.into_array().cast(target)?; + assert_arrays_eq!(casted, expected, &mut SESSION.create_execution_ctx()); + Ok(()) + } + #[test] fn test_cast_bitpacked_nullable() { let values = PrimitiveArray::from_option_iter([Some(5u16), None, Some(10), Some(15), None]); 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..ce5c54b593c 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,6 +7,7 @@ use vortex_error::VortexResult; use crate::arrays::DecimalArray; use crate::arrays::decimal::DecimalArrayExt; use crate::dtype::BigCast; +use crate::dtype::NativeDecimalType; use crate::match_each_decimal_value_type; #[expect( @@ -27,16 +28,37 @@ pub(super) fn check_decimal_identical( 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) + // `L` and `R` differ and fix `W` as the wider of them, so only one arm can match. + // The const condition lets the compiler skip codegen of the other arms. + if const { is_widest_of_distinct::() } { + 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") + })) + } else { + unreachable!("{widest} is not the wider of two distinct decimal types") + } }) }) }) } + +/// Whether `L` and `R` differ and `W` is the wider of them. +const fn is_widest_of_distinct() -> bool +where + L: NativeDecimalType, + R: NativeDecimalType, + W: NativeDecimalType, +{ + let lhs = L::DECIMAL_TYPE as u8; + let rhs = R::DECIMAL_TYPE as u8; + let widest = if lhs > rhs { lhs } else { rhs }; + + lhs != rhs && W::DECIMAL_TYPE as u8 == widest +} diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/tests.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/tests.rs index 3f894459e1d..b7bcc76259b 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/tests.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_distinct/tests.rs @@ -4,7 +4,9 @@ use std::sync::LazyLock; use rstest::rstest; +use vortex_buffer::Buffer; use vortex_buffer::buffer; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; @@ -22,11 +24,14 @@ use crate::arrays::NullArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::VarBinViewArray; +use crate::dtype::BigCast; use crate::dtype::DType; use crate::dtype::DecimalDType; +use crate::dtype::DecimalType; use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::dtype::PType; +use crate::match_each_decimal_value_type; use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -293,6 +298,46 @@ fn decimal_different_value_types() -> VortexResult<()> { assert_matches_baseline(&a, &b) } +/// Decimal values `[1, 2, last]` stored as `values_type`. +fn decimal_with_storage(values_type: DecimalType, last: i8) -> ArrayRef { + let dtype = DecimalDType::new(2, 0); + match_each_decimal_value_type!(values_type, |T| { + let values = [1i8, 2, last].map(|value| { + ::from(value).vortex_expect("small values fit every decimal type") + }); + DecimalArray::new(Buffer::copy_from(values), dtype, Validity::NonNullable).into_array() + }) +} + +/// Each pair of storage types selects a different widened comparison. +#[rstest] +fn decimal_every_storage_pair( + #[values( + DecimalType::I8, + DecimalType::I16, + DecimalType::I32, + DecimalType::I64, + DecimalType::I128, + DecimalType::I256 + )] + lhs: DecimalType, + #[values( + DecimalType::I8, + DecimalType::I16, + DecimalType::I32, + DecimalType::I64, + DecimalType::I128, + DecimalType::I256 + )] + rhs: DecimalType, +) -> VortexResult<()> { + assert_matches_baseline(&decimal_with_storage(lhs, 3), &decimal_with_storage(rhs, 3))?; + assert_matches_baseline( + &decimal_with_storage(lhs, 3), + &decimal_with_storage(rhs, -3), + ) +} + #[test] fn decimal_nullable() -> VortexResult<()> { let dtype = DecimalDType::new(5, 2); diff --git a/vortex-array/src/arrays/decimal/utils.rs b/vortex-array/src/arrays/decimal/utils.rs index 3c7304e6fda..30b271c2759 100644 --- a/vortex-array/src/arrays/decimal/utils.rs +++ b/vortex-array/src/arrays/decimal/utils.rs @@ -23,11 +23,21 @@ pub(crate) fn widened_buffer(array: &DecimalArray) -> Buff return array.buffer::(); } match_each_decimal_value_type!(array.values_type(), |T| { - array - .buffer::() - .iter() - .map(|v| W::from(*v).vortex_expect("widening decimal cast must succeed")) - .collect() + // Only storage types narrower than `W` reach here. The const condition lets the compiler + // skip codegen of the other arms. + if const { (T::DECIMAL_TYPE as u8) < (W::DECIMAL_TYPE as u8) } { + array + .buffer::() + .iter() + .map(|v| W::from(*v).vortex_expect("widening decimal cast must succeed")) + .collect() + } else { + unreachable!( + "{} is not narrower than {}", + T::DECIMAL_TYPE, + W::DECIMAL_TYPE + ) + } }) } @@ -46,45 +56,51 @@ pub fn converted_buffer( return Ok(widened_buffer(array)); } match_each_decimal_value_type!(array.values_type(), |T| { - let src = array.buffer::(); - match validity { - Mask::AllTrue(_) => { - // Keeping the overflow scan branchless and vectorizable. Only on overflow - // do we rescan for diagnostics. - let any_overflow = src.iter().fold(false, |acc, v| acc | W::from(*v).is_none()); - if any_overflow { - let (i, v) = src - .iter() - .enumerate() - .find(|&(_, v)| W::from(*v).is_none()) - .vortex_expect("overflow scan found an overflowing value"); - vortex_bail!( - "decimal value {v} at index {i} does not fit {}", - W::DECIMAL_TYPE - ); + // Only storage types wider than `W` reach here. The const condition lets the compiler + // skip codegen of the other arms. + if const { (T::DECIMAL_TYPE as u8) > (W::DECIMAL_TYPE as u8) } { + let src = array.buffer::(); + match validity { + Mask::AllTrue(_) => { + // Keeping the overflow scan branchless and vectorizable. Only on overflow + // do we rescan for diagnostics. + let any_overflow = src.iter().fold(false, |acc, v| acc | W::from(*v).is_none()); + if any_overflow { + let (i, v) = src + .iter() + .enumerate() + .find(|&(_, v)| W::from(*v).is_none()) + .vortex_expect("overflow scan found an overflowing value"); + vortex_bail!( + "decimal value {v} at index {i} does not fit {}", + W::DECIMAL_TYPE + ); + } } - } - Mask::AllFalse(_) => return Ok(Buffer::zeroed(src.len())), - Mask::Values(values) => { - let any_overflow = src.iter().fold(false, |acc, v| acc | W::from(*v).is_none()); - if any_overflow { - for (i, v) in src.iter().enumerate() { - if values.value(i) && W::from(*v).is_none() { - vortex_bail!( - "decimal value {v} at index {i} does not fit {}", - W::DECIMAL_TYPE - ); + Mask::AllFalse(_) => return Ok(Buffer::zeroed(src.len())), + Mask::Values(values) => { + let any_overflow = src.iter().fold(false, |acc, v| acc | W::from(*v).is_none()); + if any_overflow { + for (i, v) in src.iter().enumerate() { + if values.value(i) && W::from(*v).is_none() { + vortex_bail!( + "decimal value {v} at index {i} does not fit {}", + W::DECIMAL_TYPE + ); + } } } } } + // The convert pass is infallible: every valid value fits, and out-of-range null-slot + // garbage is mapped to the default value. Callers must still ignore null slots. + Ok(src + .iter() + .map(|v| W::from(*v).unwrap_or_default()) + .collect()) + } else { + unreachable!("{} is not wider than {}", T::DECIMAL_TYPE, W::DECIMAL_TYPE) } - // The convert pass is infallible: every valid value fits, and out-of-range null-slot - // garbage is mapped to the default value. Callers must still ignore null slots. - Ok(src - .iter() - .map(|v| W::from(*v).unwrap_or_default()) - .collect()) }) } @@ -142,3 +158,61 @@ pub fn narrowed_decimal(decimal_array: DecimalArray) -> DecimalArray { } } } + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_buffer::Buffer; + use vortex_error::VortexExpect; + use vortex_error::VortexResult; + use vortex_mask::Mask; + + use super::converted_buffer; + use crate::arrays::DecimalArray; + use crate::dtype::BigCast; + use crate::dtype::DecimalDType; + use crate::dtype::DecimalType; + use crate::match_each_decimal_value_type; + use crate::validity::Validity; + + /// Each pair selects a different widening or narrowing conversion. + #[rstest] + fn converted_buffer_every_storage_pair( + #[values( + DecimalType::I8, + DecimalType::I16, + DecimalType::I32, + DecimalType::I64, + DecimalType::I128, + DecimalType::I256 + )] + storage: DecimalType, + #[values( + DecimalType::I8, + DecimalType::I16, + DecimalType::I32, + DecimalType::I64, + DecimalType::I128, + DecimalType::I256 + )] + target: DecimalType, + ) -> VortexResult<()> { + let values = [1i8, -2, 3]; + let dtype = DecimalDType::new(2, 0); + let array = match_each_decimal_value_type!(storage, |S| { + let values = values.map(|value| { + ::from(value).vortex_expect("small values fit every decimal type") + }); + DecimalArray::new(Buffer::copy_from(values), dtype, Validity::NonNullable) + }); + + match_each_decimal_value_type!(target, |W| { + let converted = converted_buffer::(&array, &Mask::new_true(values.len()))?; + let expected = values.map(|value| { + ::from(value).vortex_expect("small values fit every decimal type") + }); + assert_eq!(converted.as_slice(), expected.as_slice()); + }); + Ok(()) + } +}