Add bit-packed widening cast pushdown

encodings/fastlanes/public-api.lock

@@ -34,10 +34,6 @@ pub fn vortex_fastlanes::bitpack_compress::gather_patches(&vortex_array::arrays:
 
 pub mod vortex_fastlanes::bitpack_decompress
 
-pub fn vortex_fastlanes::bitpack_decompress::apply_patches_to_uninit_range<T: vortex_array::dtype::ptype::NativePType>(&mut vortex_array::builders::primitive::UninitRange<'_, T>, &vortex_array::patches::Patches, &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<()>
-
-pub fn vortex_fastlanes::bitpack_decompress::apply_patches_to_uninit_range_fn<T: vortex_array::dtype::ptype::NativePType, F: core::ops::function::Fn(T) -> T>(&mut vortex_array::builders::primitive::UninitRange<'_, T>, &vortex_array::patches::Patches, &mut vortex_array::executor::ExecutionCtx, F) -> vortex_error::VortexResult<()>
-
 pub fn vortex_fastlanes::bitpack_decompress::count_exceptions(u8, &[usize]) -> usize
 
 pub fn vortex_fastlanes::bitpack_decompress::unpack_array(vortex_array::array::view::ArrayView<'_, vortex_fastlanes::BitPacked>, &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::arrays::primitive::vtable::PrimitiveArray>

encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs

@@ -1,6 +1,8 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::mem::MaybeUninit;
+
 use fastlanes::BitPacking;
 use itertools::Itertools;
 use num_traits::AsPrimitive;
@@ -21,6 +23,7 @@ use vortex_error::VortexResult;
 use crate::BitPacked;
 use crate::BitPackedArrayExt;
 use crate::unpack_iter::BitPacked as BitPackedUnpack;
+use crate::unpack_iter::BitUnpackedChunks;
 
 /// Unpacks a bit-packed array into a primitive array.
 pub fn unpack_array(
@@ -42,52 +45,96 @@ pub fn unpack_primitive_array<T: BitPackedUnpack>(
     Ok(builder.finish_into_primitive())
 }
 
+/// Unpack a bit-packed array directly into a same-typed `PrimitiveBuilder`.
+///
+/// This is the fast path for ordinary decompression: full FastLanes chunks are unpacked straight
+/// into the final output buffer, avoiding the scratch chunk and copy needed by mapped decode.
 pub(crate) fn unpack_into_primitive_builder<T: BitPackedUnpack>(
     array: ArrayView<'_, BitPacked>,
-    // TODO(ngates): do we want to use fastlanes alignment for this buffer?
     builder: &mut PrimitiveBuilder<T>,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<()> {
-    // If the array is empty, then we don't need to add anything to the builder.
+    unpack_into_builder_with(
+        array,
+        builder,
+        ctx,
+        |v: T| v,
+        |chunks, output, _| {
+            chunks.decode_into(output);
+        },
+    )
+}
+
+/// Unpack a bit-packed array of physical type `F` into a `PrimitiveBuilder<T>`, applying `map`
+/// to each value during decompression.
+///
+/// Use [`unpack_into_primitive_builder`] for same-type plain decompression. This mapped path is
+/// for widening casts or other element-wise transforms: each 1024-element FastLanes chunk is
+/// unpacked into a cache-resident scratch buffer and written through `map` directly into the `T`
+/// output, so when `F != T` no full-length `F`-typed intermediate is materialized.
+///
+/// The caller must ensure that every valid source value is representable in `T` under `map`; no
+/// per-value bounds check is performed.
+pub(crate) fn unpack_map_into_builder<F, T, M>(
+    array: ArrayView<'_, BitPacked>,
+    builder: &mut PrimitiveBuilder<T>,
+    ctx: &mut ExecutionCtx,
+    map: M,
+) -> VortexResult<()>
+where
+    F: BitPackedUnpack,
+    T: NativePType,
+    M: Fn(F) -> T,
+{
+    unpack_into_builder_with(array, builder, ctx, map, |chunks, output, map| {
+        chunks.decode_map_into(output, map);
+    })
+}
+
+fn unpack_into_builder_with<F, T, M, D>(
+    array: ArrayView<'_, BitPacked>,
+    builder: &mut PrimitiveBuilder<T>,
+    ctx: &mut ExecutionCtx,
+    map: M,
+    decode: D,
+) -> VortexResult<()>
+where
+    F: BitPackedUnpack,
+    T: NativePType,
+    M: Fn(F) -> T,
+    D: FnOnce(&mut BitUnpackedChunks<F>, &mut [MaybeUninit<T>], &M),
+{
     if array.is_empty() {
         return Ok(());
     }
 
-    let mut uninit_range = builder.uninit_range(array.len());
+    let len = array.len();
+    let mut uninit_range = builder.uninit_range(len);
 
-    // SAFETY: We later initialize the the uninitialized range of values with `copy_from_slice`.
+    // SAFETY: We initialize all `len` values below via `decode` and the patch loop.
     unsafe {
-        // Append a dense null Mask.
-        uninit_range.append_mask(array.validity()?.execute_mask(array.as_ref().len(), ctx)?);
+        uninit_range.append_mask(array.validity()?.execute_mask(len, ctx)?);
     }
 
-    // SAFETY: `decode_into` will initialize all values in this range.
-    let uninit_slice = unsafe { uninit_range.slice_uninit_mut(0, array.len()) };
+    // SAFETY: `decode` writes a value to every slot in this range.
+    let uninit_slice = unsafe { uninit_range.slice_uninit_mut(0, len) };
 
-    let mut bit_packed_iter = array.unpacked_chunks()?;
-    bit_packed_iter.decode_into(uninit_slice);
+    let mut chunks = array.unpacked_chunks::<F>()?;
+    decode(&mut chunks, uninit_slice, &map);
 
     if let Some(patches) = array.patches() {
-        apply_patches_to_uninit_range(&mut uninit_range, &patches, ctx)?;
-    };
+        apply_patches_to_uninit_range(&mut uninit_range, &patches, ctx, &map)?;
+    }
 
-    // SAFETY: We have set a correct validity mask via `append_mask` with `array.len()` values and
-    // initialized the same number of values needed via `decode_into`.
+    // SAFETY: A correct validity mask of `len` values was set via `append_mask`, and the same
+    // number of values was initialized via `decode` (and overwritten by patches).
     unsafe {
         uninit_range.finish();
     }
     Ok(())
 }
 
-pub fn apply_patches_to_uninit_range<T: NativePType>(
-    dst: &mut UninitRange<T>,
-    patches: &Patches,
-    ctx: &mut ExecutionCtx,
-) -> VortexResult<()> {
-    apply_patches_to_uninit_range_fn(dst, patches, ctx, |x| x)
-}
-
-pub fn apply_patches_to_uninit_range_fn<T: NativePType, F: Fn(T) -> T>(
+pub(crate) fn apply_patches_to_uninit_range<S: NativePType, T: NativePType, F: Fn(S) -> T>(
     dst: &mut UninitRange<T>,
     patches: &Patches,
     ctx: &mut ExecutionCtx,
@@ -98,7 +145,7 @@ pub fn apply_patches_to_uninit_range_fn<T: NativePType, F: Fn(T) -> T>(
     let indices = patches.indices().clone().execute::<PrimitiveArray>(ctx)?;
     let values = patches.values().clone().execute::<PrimitiveArray>(ctx)?;
     assert!(values.all_valid(ctx)?, "Patch values must be all valid");
-    let values = values.as_slice::<T>();
+    let values = values.as_slice::<S>();
 
     match_each_unsigned_integer_ptype!(indices.ptype(), |P| {
         for (index, &value) in indices.as_slice::<P>().iter().zip_eq(values) {
@@ -335,10 +382,11 @@ mod tests {
         let bitpacked = encode(&empty, 0);
 
         let mut builder = PrimitiveBuilder::<u32>::new(Nullability::NonNullable);
-        unpack_into_primitive_builder(
+        unpack_map_into_builder(
             bitpacked.as_view(),
             &mut builder,
             &mut SESSION.create_execution_ctx(),
+            |v: u32| v,
         )?;
 
         let result = builder.finish_into_primitive();
@@ -363,10 +411,11 @@ mod tests {
 
         // Unpack into a new builder.
         let mut builder = PrimitiveBuilder::<u32>::with_capacity(Nullability::Nullable, 5);
-        unpack_into_primitive_builder(
+        unpack_map_into_builder(
             bitpacked.as_view(),
             &mut builder,
             &mut SESSION.create_execution_ctx(),
+            |v: u32| v,
         )?;
 
         let result = builder.finish_into_primitive();
@@ -400,10 +449,11 @@ mod tests {
 
         // Unpack into a new builder.
         let mut builder = PrimitiveBuilder::<u32>::with_capacity(Nullability::NonNullable, 100);
-        unpack_into_primitive_builder(
+        unpack_map_into_builder(
             bitpacked.as_view(),
             &mut builder,
             &mut SESSION.create_execution_ctx(),
+            |v: u32| v,
         )?;
 
         let result = builder.finish_into_primitive();

encodings/fastlanes/src/bitpacking/array/unpack_iter.rs

@@ -181,50 +181,96 @@ impl<T: PhysicalPType, S: UnpackStrategy<T>> UnpackedChunks<T, S> {
         })
     }
 
-    /// Decode all chunks (initial, full, and trailer) into the output range.
-    /// This consolidates the logic for handling all three chunk types in one place.
+    /// Decode all chunks (initial, full, and trailer) directly into the output range.
     pub fn decode_into(&mut self, output: &mut [MaybeUninit<T>]) {
+        debug_assert_eq!(output.len(), self.len);
         let mut local_idx = 0;
 
-        // Handle initial partial chunk if present
         if let Some(initial) = self.initial() {
             local_idx = initial.len();
 
-            // TODO(connor): use `maybe_uninit_write_slice` feature when it gets stabilized.
-            // https://github.com/rust-lang/rust/issues/79995
+            // TODO(connor): use maybe_uninit_write_slice when it gets stabilized.
             // SAFETY: &[T] and &[MaybeUninit<T>] have the same layout.
             let init_initial: &[MaybeUninit<T>] = unsafe { mem::transmute(initial) };
             output[..local_idx].copy_from_slice(init_initial);
         }
 
-        // Handle full chunks
         local_idx = self.decode_full_chunks_into_at(output, local_idx);
 
-        // Handle trailing partial chunk if present
         if let Some(trailer) = self.trailer() {
-            // TODO(connor): use `maybe_uninit_write_slice` feature when it gets stabilized.
-            // https://github.com/rust-lang/rust/issues/79995
+            // TODO(connor): use maybe_uninit_write_slice when it gets stabilized.
             // SAFETY: &[T] and &[MaybeUninit<T>] have the same layout.
             let init_trailer: &[MaybeUninit<T>] = unsafe { mem::transmute(trailer) };
             output[local_idx..][..init_trailer.len()].copy_from_slice(init_trailer);
+            local_idx += init_trailer.len();
         }
+
+        debug_assert_eq!(local_idx, self.len);
+    }
+
+    /// Decode all chunks (initial, full, and trailer), mapping each unpacked value through f.
+    pub(crate) fn decode_map_into<U>(
+        &mut self,
+        output: &mut [MaybeUninit<U>],
+        mut f: impl FnMut(T) -> U,
+    ) {
+        debug_assert_eq!(output.len(), self.len);
+        let mut local_idx = 0;
+
+        if let Some(initial) = self.initial() {
+            let chunk_len = initial.len();
+            write_map(initial, &mut output[..chunk_len], &mut f);
+            local_idx += chunk_len;
+        }
+
+        if self.num_chunks != 1 {
+            let first_chunk_is_sliced = self.first_chunk_is_sliced();
+            let last_chunk_is_sliced = self.last_chunk_is_sliced();
+            let full_chunks_range =
+                (first_chunk_is_sliced as usize)..(self.num_chunks - last_chunk_is_sliced as usize);
+
+            let packed_slice: &[T::Physical] = buffer_as_slice(&self.packed);
+            let elems_per_chunk = self.elems_per_chunk();
+            for i in full_chunks_range {
+                let chunk = &packed_slice[i * elems_per_chunk..][..elems_per_chunk];
+                unsafe {
+                    let dst: &mut [T::Physical] = mem::transmute(&mut self.buffer[..]);
+                    self.strategy.unpack_chunk(self.bit_width, chunk, dst);
+                    let unpacked: &[T] = mem::transmute(&self.buffer[..]);
+                    write_map(
+                        unpacked,
+                        &mut output[local_idx..local_idx + CHUNK_SIZE],
+                        &mut f,
+                    );
+                }
+                local_idx += CHUNK_SIZE;
+            }
+        }
+
+        if let Some(trailer) = self.trailer() {
+            let chunk_len = trailer.len();
+            write_map(
+                trailer,
+                &mut output[local_idx..local_idx + chunk_len],
+                &mut f,
+            );
+            local_idx += chunk_len;
+        }
+
+        debug_assert_eq!(local_idx, self.len);
     }
 
     /// Unpack full chunks into output range starting at the given index.
-    /// Returns the next local index to write to.
     fn decode_full_chunks_into_at(
         &mut self,
         output: &mut [MaybeUninit<T>],
         start_idx: usize,
     ) -> usize {
-        // If there's only one chunk it has been handled already by `initial` method
         if self.num_chunks == 1 {
-            // Return the start_idx since initial already wrote everything.
             return start_idx;
         }
 
         let first_chunk_is_sliced = self.first_chunk_is_sliced();
-
         let last_chunk_is_sliced = self.last_chunk_is_sliced();
         let full_chunks_range =
             (first_chunk_is_sliced as usize)..(self.num_chunks - last_chunk_is_sliced as usize);
@@ -238,7 +284,7 @@ impl<T: PhysicalPType, S: UnpackStrategy<T>> UnpackedChunks<T, S> {
 
             unsafe {
                 let uninit_dst = &mut output[local_idx..local_idx + CHUNK_SIZE];
-                // SAFETY: &[T] and &[MaybeUninit<T>] have the same layout
+                // SAFETY: &[T] and &[MaybeUninit<T>] have the same layout.
                 let dst: &mut [T::Physical] = mem::transmute(uninit_dst);
                 self.strategy.unpack_chunk(self.bit_width, chunk, dst);
             }
@@ -340,6 +386,12 @@ fn buffer_as_slice<T>(buffer: &ByteBuffer) -> &[T] {
     unsafe { std::slice::from_raw_parts(packed_ptr, packed_len) }
 }
 
+fn write_map<T: Copy, U>(src: &[T], dst: &mut [MaybeUninit<U>], f: &mut impl FnMut(T) -> U) {
+    for (dst, &src) in dst.iter_mut().zip(src.iter()) {
+        dst.write(f(src));
+    }
+}
+
 pub trait BitPacked: PhysicalPType<Physical: BitPacking> {}
 
 impl BitPacked for i8 {}