replace UninitRange for unpack decode

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

Author: connortsui20

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

@@ -65,15 +65,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,7 +9,6 @@ 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;
 
@@ -159,29 +158,34 @@ 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();
+
+            // SAFETY: &[T] and &[MaybeUninit<T>] have the same layout.
+            let uninit_initial: &[MaybeUninit<T>] = unsafe { mem::transmute(initial) };
+            output[..uninit_initial.len()].copy_from_slice(uninit_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);
+            // SAFETY: &[T] and &[MaybeUninit<T>] have the same layout.
+            let uninit_trailer: &[MaybeUninit<T>] = unsafe { mem::transmute(trailer) };
+            output[local_idx..][..uninit_trailer.len()].copy_from_slice(uninit_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 +208,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();
     }