Chore: Use MaybeUninit instead of UninitRange for bit unpacking (#5427)

This will make porting over to vectors easier.

---------

Signed-off-by: Connor Tsui <[email protected]>

Author: connortsui20

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

@@ -8,33 +8,14 @@ use vortex_array::arrays::PrimitiveArray;
 use vortex_array::builders::{ArrayBuilder, PrimitiveBuilder, UninitRange};
 use vortex_array::patches::Patches;
 use vortex_dtype::{
-    IntegerPType, NativePType, PhysicalPType, match_each_integer_ptype,
-    match_each_unsigned_integer_ptype,
+    IntegerPType, NativePType, match_each_integer_ptype, match_each_unsigned_integer_ptype,
 };
 use vortex_error::VortexExpect;
 use vortex_mask::Mask;
 use vortex_scalar::Scalar;
 
 use crate::BitPackedArray;
-use crate::unpack_iter::{BitPacked, UnpackStrategy};
-
-/// BitPacking strategy - uses plain bitpacking without reference value
-pub struct BitPackingStrategy;
-
-impl<T: PhysicalPType<Physical: BitPacking>> UnpackStrategy<T> for BitPackingStrategy {
-    #[inline(always)]
-    unsafe fn unpack_chunk(
-        &self,
-        bit_width: usize,
-        chunk: &[T::Physical],
-        dst: &mut [T::Physical],
-    ) {
-        // SAFETY: Caller must ensure [`BitPacking::unchecked_unpack`] safety requirements hold.
-        unsafe {
-            BitPacking::unchecked_unpack(bit_width, chunk, dst);
-        }
-    }
-}
+use crate::unpack_iter::BitPacked;
 
 pub fn unpack(array: &BitPackedArray) -> PrimitiveArray {
     match_each_integer_ptype!(array.ptype(), |P| { unpack_primitive::<P>(array) })
@@ -65,15 +46,18 @@ pub(crate) fn unpack_into<T: BitPacked>(
         uninit_range.append_mask(array.validity_mask());
     }
 
+    // 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(&mut uninit_range);
+    bit_packed_iter.decode_into(uninit_slice);
 
     if let Some(patches) = array.patches() {
         apply_patches(&mut uninit_range, patches);
     };
 
     // SAFETY: We have set a correct validity mask via `append_mask` with `array.len()` values and
-    // initialized the same number of values needed via calls to `copy_from_slice`.
+    // initialized the same number of values needed via `decode_into`.
     unsafe {
         uninit_range.finish();
     }

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

@@ -9,12 +9,10 @@ use lending_iterator::gat;
 use lending_iterator::prelude::Item;
 #[gat(Item)]
 use lending_iterator::prelude::LendingIterator;
-use vortex_array::builders::UninitRange;
 use vortex_buffer::ByteBuffer;
 use vortex_dtype::PhysicalPType;
 
 use crate::BitPackedArray;
-use crate::bitpacking::bitpack_decompress::BitPackingStrategy;
 
 const CHUNK_SIZE: usize = 1024;
 
@@ -28,6 +26,24 @@ pub trait UnpackStrategy<T: PhysicalPType> {
     unsafe fn unpack_chunk(&self, bit_width: usize, chunk: &[T::Physical], dst: &mut [T::Physical]);
 }
 
+/// BitPacking strategy - uses plain bitpacking without reference value
+pub struct BitPackingStrategy;
+
+impl<T: PhysicalPType<Physical: BitPacking>> UnpackStrategy<T> for BitPackingStrategy {
+    #[inline(always)]
+    unsafe fn unpack_chunk(
+        &self,
+        bit_width: usize,
+        chunk: &[T::Physical],
+        dst: &mut [T::Physical],
+    ) {
+        // SAFETY: Caller must ensure [`BitPacking::unchecked_unpack`] safety requirements hold.
+        unsafe {
+            BitPacking::unchecked_unpack(bit_width, chunk, dst);
+        }
+    }
+}
+
 /// Accessor to unpacked chunks of bitpacked arrays
 ///
 /// The usual pattern of usage should follow
@@ -159,29 +175,38 @@ 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.
-    pub fn decode_into(&mut self, output: &mut UninitRange<T>) {
+    pub fn decode_into(&mut self, output: &mut [MaybeUninit<T>]) {
         let mut local_idx = 0;
 
         // Handle initial partial chunk if present
         if let Some(initial) = self.initial() {
-            output.copy_from_slice(0, 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
+            // 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() {
-            output.copy_from_slice(local_idx, trailer);
+            // TODO(connor): use `maybe_uninit_write_slice` feature when it gets stabilized.
+            // https://github.com/rust-lang/rust/issues/79995
+            // 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);
         }
     }
 
     /// 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 UninitRange<T>,
+        output: &mut [MaybeUninit<T>],
         start_idx: usize,
     ) -> usize {
         // If there's only one chunk it has been handled already by `initial` method
@@ -204,8 +229,7 @@ impl<T: PhysicalPType, S: UnpackStrategy<T>> UnpackedChunks<T, S> {
             let chunk = &packed_slice[i * elems_per_chunk..][..elems_per_chunk];
 
             unsafe {
-                // SAFETY: We're about to initialize CHUNK_SIZE elements at local_idx.
-                let uninit_dst = output.slice_uninit_mut(local_idx, CHUNK_SIZE);
+                let uninit_dst = &mut output[local_idx..local_idx + CHUNK_SIZE];
                 // 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);

encodings/fastlanes/src/for/compress.rs

@@ -124,19 +124,23 @@ fn fused_decompress<T: PhysicalPType<Physical = T> + UnsignedPType + FoR + Wrapp
 
     let mut builder = PrimitiveBuilder::<T>::with_capacity(for_.dtype().nullability(), bp.len());
     let mut uninit_range = builder.uninit_range(bp.len());
-
-    // Decode all chunks (initial, full, and trailer) in one call
-    unpacked.decode_into(&mut uninit_range);
-
     unsafe {
         // Append a dense null Mask.
         uninit_range.append_mask(bp.validity_mask());
     }
 
+    // SAFETY: `decode_into` will initialize all values in this range.
+    let uninit_slice = unsafe { uninit_range.slice_uninit_mut(0, bp.len()) };
+
+    // Decode all chunks (initial, full, and trailer) in one call
+    unpacked.decode_into(uninit_slice);
+
     if let Some(patches) = bp.patches() {
         bitpack_decompress::apply_patches_fn(&mut uninit_range, patches, |v| v.wrapping_add(&ref_));
     };
 
+    // 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();
     }