Merge remote-tracking branch 'origin/develop' into ji/chunked-array-store-child-offset-mem

# Conflicts:
#	vortex-array/src/arrays/chunked/vtable/mod.rs

Author: joseph-isaacs

Cargo.lock

@@ -225,6 +225,14 @@ pub trait ALPFloat: private::Sealed + Float + Display + NativePType {
         encoded.map_each_in_place(move |encoded| Self::decode_single(encoded, exponents))
     }
 
+    fn decode_into(encoded: &[Self::ALPInt], exponents: Exponents, output: &mut [Self]) {
+        assert_eq!(encoded.len(), output.len());
+
+        for i in 0..encoded.len() {
+            output[i] = Self::decode_single(encoded[i], exponents)
+        }
+    }
+
     fn decode_slice_inplace(encoded: &mut [Self::ALPInt], exponents: Exponents) {
         let decoded: &mut [Self] = unsafe { transmute(encoded) };
         decoded.iter_mut().for_each(|v| {

encodings/alp/src/alp/mod.rs

@@ -8,46 +8,27 @@ 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) })
+pub fn unpack_array(array: &BitPackedArray) -> PrimitiveArray {
+    match_each_integer_ptype!(array.ptype(), |P| { unpack_primitive_array::<P>(array) })
 }
 
-pub fn unpack_primitive<T: BitPacked>(array: &BitPackedArray) -> PrimitiveArray {
+pub fn unpack_primitive_array<T: BitPacked>(array: &BitPackedArray) -> PrimitiveArray {
     let mut builder = PrimitiveBuilder::with_capacity(array.dtype().nullability(), array.len());
-    unpack_into::<T>(array, &mut builder);
+    unpack_into_primitive_builder::<T>(array, &mut builder);
     assert_eq!(builder.len(), array.len());
     builder.finish_into_primitive()
 }
 
-pub(crate) fn unpack_into<T: BitPacked>(
+pub(crate) fn unpack_into_primitive_builder<T: BitPacked>(
     array: &BitPackedArray,
     // TODO(ngates): do we want to use fastlanes alignment for this buffer?
     builder: &mut PrimitiveBuilder<T>,
@@ -65,25 +46,28 @@ 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);
+        apply_patches_to_uninit_range(&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();
     }
 }
 
-pub fn apply_patches<T: NativePType>(dst: &mut UninitRange<T>, patches: &Patches) {
-    apply_patches_fn(dst, patches, |x| x)
+pub fn apply_patches_to_uninit_range<T: NativePType>(dst: &mut UninitRange<T>, patches: &Patches) {
+    apply_patches_to_uninit_range_fn(dst, patches, |x| x)
 }
 
-pub fn apply_patches_fn<T: NativePType, F: Fn(T) -> T>(
+pub fn apply_patches_to_uninit_range_fn<T: NativePType, F: Fn(T) -> T>(
     dst: &mut UninitRange<T>,
     patches: &Patches,
     f: F,
@@ -96,7 +80,7 @@ pub fn apply_patches_fn<T: NativePType, F: Fn(T) -> T>(
     let values = values.as_slice::<T>();
 
     match_each_unsigned_integer_ptype!(indices.ptype(), |P| {
-        insert_values_and_validity_at_indices(
+        insert_values_and_validity_at_indices_to_uninit_range(
             dst,
             indices.as_slice::<P>(),
             values,
@@ -107,7 +91,11 @@ pub fn apply_patches_fn<T: NativePType, F: Fn(T) -> T>(
     });
 }
 
-fn insert_values_and_validity_at_indices<T: NativePType, IndexT: IntegerPType, F: Fn(T) -> T>(
+fn insert_values_and_validity_at_indices_to_uninit_range<
+    T: NativePType,
+    IndexT: IntegerPType,
+    F: Fn(T) -> T,
+>(
     dst: &mut UninitRange<T>,
     indices: &[IndexT],
     values: &[T],
@@ -223,23 +211,23 @@ mod tests {
     fn test_all_zeros() {
         let zeros = buffer![0u16, 0, 0, 0].into_array().to_primitive();
         let bitpacked = bitpack_encode(&zeros, 0, None).unwrap();
-        let actual = unpack(&bitpacked);
+        let actual = unpack_array(&bitpacked);
         assert_arrays_eq!(actual, PrimitiveArray::from_iter([0u16, 0, 0, 0]));
     }
 
     #[test]
     fn test_simple_patches() {
         let zeros = buffer![0u16, 1, 0, 1].into_array().to_primitive();
         let bitpacked = bitpack_encode(&zeros, 0, None).unwrap();
-        let actual = unpack(&bitpacked);
+        let actual = unpack_array(&bitpacked);
         assert_arrays_eq!(actual, PrimitiveArray::from_iter([0u16, 1, 0, 1]));
     }
 
     #[test]
     fn test_one_full_chunk() {
         let zeros = BufferMut::from_iter(0u16..1024).into_array().to_primitive();
         let bitpacked = bitpack_encode(&zeros, 10, None).unwrap();
-        let actual = unpack(&bitpacked);
+        let actual = unpack_array(&bitpacked);
         assert_arrays_eq!(actual, PrimitiveArray::from_iter(0u16..1024));
     }
 
@@ -250,7 +238,7 @@ mod tests {
             .to_primitive();
         let bitpacked = bitpack_encode(&zeros, 10, None).unwrap();
         assert!(bitpacked.patches().is_some());
-        let actual = unpack(&bitpacked);
+        let actual = unpack_array(&bitpacked);
         assert_arrays_eq!(
             actual,
             PrimitiveArray::from_iter((5u16..1029).chain(5u16..1029).chain(5u16..1029))
@@ -262,7 +250,7 @@ mod tests {
         let zeros = BufferMut::from_iter(0u16..1025).into_array().to_primitive();
         let bitpacked = bitpack_encode(&zeros, 11, None).unwrap();
         assert!(bitpacked.patches().is_none());
-        let actual = unpack(&bitpacked);
+        let actual = unpack_array(&bitpacked);
         assert_arrays_eq!(actual, PrimitiveArray::from_iter(0u16..1025));
     }
 
@@ -274,7 +262,7 @@ mod tests {
         let bitpacked = bitpack_encode(&zeros, 10, None).unwrap();
         assert_eq!(bitpacked.len(), 1025);
         assert!(bitpacked.patches().is_some());
-        let actual = unpack(&bitpacked);
+        let actual = unpack_array(&bitpacked);
         assert_arrays_eq!(actual, PrimitiveArray::from_iter(512u16..1537));
     }
 
@@ -287,7 +275,7 @@ mod tests {
         assert_eq!(bitpacked.len(), 1025);
         assert!(bitpacked.patches().is_some());
         let bitpacked = bitpacked.slice(1023..1025);
-        let actual = unpack(bitpacked.as_::<BitPackedVTable>());
+        let actual = unpack_array(bitpacked.as_::<BitPackedVTable>());
         assert_arrays_eq!(actual, PrimitiveArray::from_iter([1535u16, 1536]));
     }
 
@@ -300,7 +288,7 @@ mod tests {
         assert_eq!(bitpacked.len(), 2229);
         assert!(bitpacked.patches().is_some());
         let bitpacked = bitpacked.slice(1023..2049);
-        let actual = unpack(bitpacked.as_::<BitPackedVTable>());
+        let actual = unpack_array(bitpacked.as_::<BitPackedVTable>());
         assert_arrays_eq!(
             actual,
             PrimitiveArray::from_iter((1023u16..2049).map(|x| x + 512))
@@ -313,7 +301,7 @@ mod tests {
         let bitpacked = bitpack_encode(&empty, 0, None).unwrap();
 
         let mut builder = PrimitiveBuilder::<u32>::new(Nullability::NonNullable);
-        unpack_into(&bitpacked, &mut builder);
+        unpack_into_primitive_builder(&bitpacked, &mut builder);
 
         let result = builder.finish_into_primitive();
         assert_eq!(
@@ -336,7 +324,7 @@ mod tests {
 
         // Unpack into a new builder.
         let mut builder = PrimitiveBuilder::<u32>::with_capacity(Nullability::Nullable, 5);
-        unpack_into(&bitpacked, &mut builder);
+        unpack_into_primitive_builder(&bitpacked, &mut builder);
 
         let result = builder.finish_into_primitive();
 
@@ -367,7 +355,7 @@ mod tests {
 
         // Unpack into a new builder.
         let mut builder = PrimitiveBuilder::<u32>::with_capacity(Nullability::NonNullable, 100);
-        unpack_into(&bitpacked, &mut builder);
+        unpack_into_primitive_builder(&bitpacked, &mut builder);
 
         let result = builder.finish_into_primitive();
 

encodings/fastlanes/src/bitpacking/array/bitpack_decompress.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/bitpacking/array/unpack_iter.rs

@@ -7,17 +7,17 @@ use vortex_array::vtable::CanonicalVTable;
 use vortex_dtype::match_each_integer_ptype;
 use vortex_error::VortexExpect;
 
-use crate::bitpack_decompress::{unpack, unpack_into};
+use crate::bitpack_decompress::{unpack_array, unpack_into_primitive_builder};
 use crate::{BitPackedArray, BitPackedVTable};
 
 impl CanonicalVTable<BitPackedVTable> for BitPackedVTable {
     fn canonicalize(array: &BitPackedArray) -> Canonical {
-        Canonical::Primitive(unpack(array))
+        Canonical::Primitive(unpack_array(array))
     }
 
     fn append_to_builder(array: &BitPackedArray, builder: &mut dyn ArrayBuilder) {
         match_each_integer_ptype!(array.ptype(), |T| {
-            unpack_into::<T>(
+            unpack_into_primitive_builder::<T>(
                 array,
                 builder
                     .as_any_mut()

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

@@ -124,19 +124,25 @@ 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_));
+        bitpack_decompress::apply_patches_to_uninit_range_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();
     }