fix

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>

Author: joseph-isaacs

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

@@ -37,67 +37,31 @@ pub fn unpack_primitive_array<T: BitPackedUnpack>(
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<PrimitiveArray> {
     let mut builder = PrimitiveBuilder::with_capacity(array.dtype().nullability(), array.len());
-    unpack_into_primitive_builder::<T>(array, &mut builder, ctx)?;
+    unpack_map_into_builder::<T, T, _>(array, &mut builder, ctx, |v| v)?;
     assert_eq!(builder.len(), array.len());
     Ok(builder.finish_into_primitive())
 }
 
-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.
-    if array.is_empty() {
-        return Ok(());
-    }
-
-    let mut uninit_range = builder.uninit_range(array.len());
-
-    // SAFETY: We later initialize the the uninitialized range of values with `copy_from_slice`.
-    unsafe {
-        // Append a dense null Mask.
-        uninit_range.append_mask(array.validity()?.execute_mask(array.as_ref().len(), ctx)?);
-    }
-
-    // SAFETY: `decode_into` will initialize all values in this range.
-    let uninit_slice = unsafe { uninit_range.slice_uninit_mut(0, array.len()) };
-
-    let mut bit_packed_iter = array.unpacked_chunks()?;
-    bit_packed_iter.decode_into(uninit_slice);
-
-    if let Some(patches) = array.patches() {
-        apply_patches_to_uninit_range(&mut uninit_range, &patches, ctx)?;
-    };
-
-    // 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`.
-    unsafe {
-        uninit_range.finish();
-    }
-    Ok(())
-}
-
-/// Unpack a bit-packed array of physical type `F` directly into a wider primitive type `T`,
-/// casting each value during decompression.
+/// Unpack a bit-packed array of physical type `F` into a `PrimitiveBuilder<T>`, applying `map`
+/// to each value during decompression.
 ///
-/// This is the "cast pushdown" path: rather than canonicalizing to a full-length `F`-typed
-/// `PrimitiveArray` and then casting it to `T` (two full-length buffers, with the `F` intermediate
-/// written out to RAM), we unpack each 1024-element FastLanes chunk into a small cache-resident
-/// scratch buffer and cast-copy straight into the `T` output. Only the `T` output buffer is
-/// allocated and touched in RAM.
+/// Pass `|v| v` (with `F = T`) for plain decompression, `|v: F| v.as_()` for a widening cast,
+/// or any other element-wise transform. 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 all valid values are representable in `T` (it is intended for widening
-/// casts such as `u16 -> u32`); narrowing or sign-changing casts are not validated here.
-pub(crate) fn unpack_and_cast_into_builder<F, T>(
+/// 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 + AsPrimitive<T>,
+    F: BitPackedUnpack,
     T: NativePType,
+    M: Fn(F) -> T,
 {
     if array.is_empty() {
         return Ok(());
@@ -115,10 +79,10 @@ where
     let uninit_slice = unsafe { uninit_range.slice_uninit_mut(0, len) };
 
     let mut chunks = array.unpacked_chunks::<F>()?;
-    chunks.decode_map_into(uninit_slice, |v: F| v.as_());
+    chunks.decode_map_into(uninit_slice, &map);
 
     if let Some(patches) = array.patches() {
-        apply_patches_to_uninit_range_map(&mut uninit_range, &patches, ctx, |v: F| v.as_())?;
+        apply_patches_to_uninit_range(&mut uninit_range, &patches, ctx, &map)?;
     }
 
     // SAFETY: A correct validity mask of `len` values was set via `append_mask`, and the same
@@ -129,24 +93,7 @@ where
     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, |v: T| v)
-}
-
-pub fn apply_patches_to_uninit_range_fn<T: NativePType, F: Fn(T) -> T>(
-    dst: &mut UninitRange<T>,
-    patches: &Patches,
-    ctx: &mut ExecutionCtx,
-    f: F,
-) -> VortexResult<()> {
-    apply_patches_to_uninit_range_map(dst, patches, ctx, f)
-}
-
-pub(crate) fn apply_patches_to_uninit_range_map<S: NativePType, T: NativePType, F: Fn(S) -> T>(
+pub fn apply_patches_to_uninit_range<S: NativePType, T: NativePType, F: Fn(S) -> T>(
     dst: &mut UninitRange<T>,
     patches: &Patches,
     ctx: &mut ExecutionCtx,
@@ -394,10 +341,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| v,
         )?;
 
         let result = builder.finish_into_primitive();
@@ -422,10 +370,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| v,
         )?;
 
         let result = builder.finish_into_primitive();
@@ -459,10 +408,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| v,
         )?;
 
         let result = builder.finish_into_primitive();

encodings/fastlanes/src/bitpacking/compute/cast.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use num_traits::AsPrimitive;
 use vortex_array::ArrayRef;
 use vortex_array::ArrayView;
 use vortex_array::ExecutionCtx;
@@ -17,7 +18,7 @@ use vortex_error::VortexResult;
 
 use crate::bitpacking::BitPacked;
 use crate::bitpacking::array::BitPackedArrayExt;
-use crate::bitpacking::array::bitpack_decompress::unpack_and_cast_into_builder;
+use crate::bitpacking::array::bitpack_decompress::unpack_map_into_builder;
 
 /// Returns `true` if casting `src` to `tgt` is a widening integer cast for which every value a
 /// bit-packed array can hold is guaranteed to be representable in `tgt` (so no per-value bounds
@@ -101,7 +102,7 @@ impl CastKernel for BitPacked {
         let result = match_each_integer_ptype!(tgt, |T| {
             let mut builder = PrimitiveBuilder::<T>::with_capacity(tgt_nullability, array.len());
             match_each_integer_ptype!(src, |F| {
-                unpack_and_cast_into_builder::<F, T>(array, &mut builder, ctx)?;
+                unpack_map_into_builder::<F, T, _>(array, &mut builder, ctx, |v: F| v.as_())?;
             });
             builder.finish_into_primitive().into_array()
         });

encodings/fastlanes/src/bitpacking/vtable/mod.rs

@@ -44,7 +44,7 @@ use crate::BitPackedArrayExt;
 use crate::BitPackedData;
 use crate::BitPackedDataParts;
 use crate::bitpack_decompress::unpack_array;
-use crate::bitpack_decompress::unpack_into_primitive_builder;
+use crate::bitpack_decompress::unpack_map_into_builder;
 use crate::bitpacking::array::BitPackedSlots;
 use crate::bitpacking::array::BitPackedSlotsView;
 use crate::bitpacking::array::PATCH_SLOTS;
@@ -239,13 +239,14 @@ impl VTable for BitPacked {
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<()> {
         match_each_integer_ptype!(array.dtype().as_ptype(), |T| {
-            unpack_into_primitive_builder::<T>(
+            unpack_map_into_builder::<T, T, _>(
                 array,
                 builder
                     .as_any_mut()
                     .downcast_mut()
                     .vortex_expect("bit packed array must canonicalize into a primitive array"),
                 ctx,
+                |v| v,
             )
         })
     }

encodings/fastlanes/src/for/array/for_decompress.rs

@@ -119,7 +119,7 @@ pub(crate) fn fused_decompress<
     unpacked.decode_into(uninit_slice);
 
     if let Some(patches) = bp.patches() {
-        bitpack_decompress::apply_patches_to_uninit_range_fn(
+        bitpack_decompress::apply_patches_to_uninit_range(
             &mut uninit_range,
             &patches,
             ctx,