From fd2091f9867f23d8e6d212cf5079474c3431a397 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 26 Oct 2026 14:25:46 +0000 Subject: [PATCH] Unpack bit-packed arrays directly for narrowing casts that fit BitPacked's CastKernel only took its direct unpack path for widening integer casts. A patch-free bit-packed array holds only values in 0..2^bit_width, so a narrowing cast whose target holds every such value cannot overflow. Take the same unpack path for those casts too, skipping the full-length intermediate buffer and the generic cast's bounds-check scan. Split out of vortex-data/vortex#10761. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01SQocaYaDvJHfmorKXnyvJC --- .../fastlanes/src/bitpacking/compute/cast.rs | 52 +++++++++++++++++-- 1 file changed, 48 insertions(+), 4 deletions(-) diff --git a/encodings/fastlanes/src/bitpacking/compute/cast.rs b/encodings/fastlanes/src/bitpacking/compute/cast.rs index 8203e68a6b7..53f0e8d8481 100644 --- a/encodings/fastlanes/src/bitpacking/compute/cast.rs +++ b/encodings/fastlanes/src/bitpacking/compute/cast.rs @@ -32,6 +32,14 @@ fn is_widening_int_cast(src: PType, tgt: PType) -> bool { && (src.is_unsigned_int() || tgt.is_signed_int()) } +/// Returns `true` if every value a patch-free array packed at `bit_width` can hold fits `tgt`. +/// Packed values are non-negative and below `2^bit_width`. +fn packed_values_fit(bit_width: u8, src: PType, tgt: PType) -> bool { + src.is_int() + && tgt.is_int() + && usize::from(bit_width) + usize::from(tgt.is_signed_int()) <= tgt.bit_width() +} + fn build_with_validity( array: ArrayView<'_, BitPacked>, bit_width: u8, @@ -89,15 +97,17 @@ impl CastKernel for BitPacked { return build_with_validity(array, bit_width, dtype, new_validity).map(Some); } - // Widening integer cast: unpack each FastLanes chunk into a cache-resident scratch buffer - // and cast-copy straight into the wide output, avoiding a full-length intermediate buffer - // and the generic cast kernel's bounds-check scan (unnecessary when widening). + // Widening integer cast, or a narrowing one whose target holds every packed value: unpack + // each FastLanes chunk into a cache-resident scratch buffer and cast-copy straight into the + // output, avoiding a full-length intermediate buffer and the generic cast kernel's + // bounds-check scan (unnecessary when every value fits). let DType::Primitive(tgt, tgt_nullability) = dtype else { return Ok(None); }; let (tgt, tgt_nullability) = (*tgt, *tgt_nullability); let src = array.dtype().as_ptype(); - if !is_widening_int_cast(src, tgt) { + let fits = || array.patches().is_none() && packed_values_fit(bit_width, src, tgt); + if !is_widening_int_cast(src, tgt) && !fits() { return Ok(None); } @@ -144,6 +154,7 @@ mod tests { use vortex_session::VortexSession; use crate::BitPackedArray; + use crate::BitPackedArrayExt; use crate::BitPackedData; static SESSION: LazyLock = LazyLock::new(|| { @@ -266,6 +277,39 @@ mod tests { Ok(()) } + /// Narrowing casts of patch-free arrays whose bit width fits the target unpack directly and + /// match a plain primitive cast, including sliced and nullable inputs. + #[rstest] + #[case(PType::I64, PType::U8)] + #[case(PType::I64, PType::I8)] + #[case(PType::U64, PType::U16)] + #[case(PType::I32, PType::I16)] + fn test_cast_bitpacked_narrowing_fits( + #[case] src: PType, + #[case] tgt: PType, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let source = match_each_integer_ptype!(src, |S| { + PrimitiveArray::from_option_iter((0..2051usize).map(|i| { + (i % 11 != 5).then(|| { + ::from_usize(i % 128) + .expect("test values fit every integer ptype") + }) + })) + }) + .into_array(); + let packed = bp(&source, 7); + assert!(packed.patches().is_none()); + let target = DType::Primitive(tgt, Nullability::Nullable); + + let casted = packed.clone().into_array().cast(target.clone())?; + assert_arrays_eq!(casted, source.cast(target.clone())?, &mut ctx); + + let sliced = packed.into_array().slice(500..1700)?.cast(target.clone())?; + assert_arrays_eq!(sliced, source.slice(500..1700)?.cast(target)?, &mut ctx); + Ok(()) + } + #[rstest] #[case(bp(&buffer![0u8, 10, 20, 30, 40, 50, 60, 63].into_array(), 6))] #[case(bp(&buffer![0u16, 100, 200, 300, 400, 500].into_array(), 9))]