wip

Signed-off-by: Alexander Droste <[email protected]>

Author: 0ax1

encodings/alp/src/alp/array.rs

@@ -360,6 +360,12 @@ impl ALPArray {
     pub fn patches(&self) -> Option<&Patches> {
         self.patches.as_ref()
     }
+
+    /// Consumes the array and returns its parts.
+    #[inline]
+    pub fn into_parts(self) -> (ArrayRef, Exponents, Option<Patches>, DType) {
+        (self.encoded, self.exponents, self.patches, self.dtype)
+    }
 }
 
 impl ValidityChild<ALPVTable> for ALPVTable {

encodings/alp/src/alp/decompress.rs

@@ -4,15 +4,15 @@
 use std::mem::transmute;
 
 use num_traits::AsPrimitive;
+use vortex_array::ArrayRef;
 use vortex_array::ToCanonical;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::patch_chunk;
 use vortex_array::patches::Patches;
-use vortex_array::validity::Validity;
 use vortex_array::vtable::ValidityHelper;
 use vortex_buffer::BufferMut;
+use vortex_dtype::DType;
 use vortex_dtype::match_each_unsigned_integer_ptype;
-use vortex_mask::MaskMut;
 use vortex_vector::primitive::PVectorMut;
 
 use crate::ALPArray;
@@ -31,40 +31,26 @@ pub(crate) fn decompress_to_pvector<T: ALPFloat>(array: ALPArray) -> PVectorMut<
         return PVectorMut::with_capacity(0);
     }
 
-    let (buffer, validity) = decompress_to_components::<T>(array);
-    // SAFETY: buffer and validity have same length from decompress_to_components
-    unsafe { PVectorMut::new_unchecked(buffer, validity) }
-}
-
-/// Decompresses an ALP-encoded array to a buffer and mask.
-///
-/// Returns `(BufferMut<T>, MaskMut)` with matching lengths.
-pub(crate) fn decompress_to_components<T: ALPFloat>(array: ALPArray) -> (BufferMut<T>, MaskMut) {
-    let encoded = array.encoded().to_primitive();
-    let validity = encoded.validity().clone();
-    let exponents = array.exponents();
-    let patches = array.patches().cloned();
-
-    // We need to drop ALPArray here in case converting encoded buffer into
-    // primitive didn't create a copy. In that case both alp_encoded and array
-    // will hold a reference to the buffer we want to mutate.
-    drop(array);
-
-    let encoded_buffer = encoded.into_buffer_mut();
-
-    // Extract chunk offsets if patches exist
-    let chunk_offsets = patches
-        .as_ref()
-        .and_then(|p| p.chunk_offsets().as_ref())
-        .map(|c| c.to_primitive());
-
-    decompress_core::<T>(
-        encoded_buffer,
-        exponents,
-        validity,
-        patches.as_ref(),
-        chunk_offsets.as_ref(),
-    )
+    let (encoded, exponents, patches, dtype) = array.into_parts();
+    let decompressed = if let Some(ref patches) = patches
+        && let Some(chunk_offsets) = patches.chunk_offsets()
+    {
+        decompress_chunked(
+            encoded,
+            exponents,
+            patches,
+            &chunk_offsets.as_ref().to_primitive(),
+            dtype,
+        )
+    } else {
+        decompress_unchunked(encoded, exponents, patches, dtype)
+    };
+
+    let validity = decompressed.validity().clone();
+    let buffer = decompressed.into_buffer_mut::<T>();
+    let validity_mask = validity.to_mask(buffer.len()).into_mut();
+    // SAFETY: buffer and validity_mask have same length.
+    unsafe { PVectorMut::new_unchecked(buffer, validity_mask) }
 }
 
 /// Decompresses an ALP-encoded array.
@@ -77,34 +63,19 @@ pub(crate) fn decompress_to_components<T: ALPFloat>(array: ALPArray) -> (BufferM
 ///
 /// A `PrimitiveArray` containing the decompressed floating-point values with all patches applied.
 pub fn decompress(array: ALPArray) -> PrimitiveArray {
-    let dtype = array.dtype().clone();
-
-    match_each_alp_float_ptype!(dtype.as_ptype(), |T| {
-        let (buffer, validity_mask) = decompress_to_components::<T>(array);
-        PrimitiveArray::new(
-            buffer.freeze(),
-            Validity::from_mask(validity_mask.freeze(), dtype.nullability()),
-        )
-    })
-}
-
-/// Core decompression logic that all public APIs use.
-/// Returns raw components that can be assembled into different array types.
-fn decompress_core<T: ALPFloat>(
-    encoded_buffer: BufferMut<T::ALPInt>,
-    exponents: Exponents,
-    validity: Validity,
-    patches: Option<&Patches>,
-    patches_chunk_offsets: Option<&PrimitiveArray>,
-) -> (BufferMut<T>, MaskMut) {
-    // Only use chunked decompression if we have both patches AND chunk offsets
-    if let Some(patches) = patches
-        && let Some(chunk_offsets) = patches_chunk_offsets
+    let (encoded, exponents, patches, dtype) = array.into_parts();
+    if let Some(ref patches) = patches
+        && let Some(chunk_offsets) = patches.chunk_offsets()
     {
-        decompress_chunked_core::<T>(encoded_buffer, exponents, validity, patches, chunk_offsets)
+        decompress_chunked(
+            encoded,
+            exponents,
+            patches,
+            &chunk_offsets.as_ref().to_primitive(),
+            dtype,
+        )
     } else {
-        // Handle both: no patches, or patches without chunk offsets
-        decompress_unchunked_core::<T>(encoded_buffer, exponents, validity, patches)
+        decompress_unchunked(encoded, exponents, patches, dtype)
     }
 }
 
@@ -114,103 +85,113 @@ fn decompress_core<T: ALPFloat>(
 ///
 /// # Arguments
 ///
-/// * `encoded_buffer` - The encoded buffer containing ALP-compressed data
-/// * `exponents` - The ALP exponents used for encoding
-/// * `validity` - The validity mask for the array
+/// * `array` - The ALP-encoded array to decompress
 /// * `patches` - The patches containing exceptional values and their positions
 /// * `patches_chunk_offsets` - Offsets into the patches array for each chunk
 ///
 /// # Returns
 ///
-/// A tuple containing the decompressed buffer and validity mask
+/// A `PrimitiveArray` containing the decompressed values with all patches applied.
 #[expect(
     clippy::cognitive_complexity,
     reason = "complexity is from nested match_each_* macros"
 )]
-fn decompress_chunked_core<T: ALPFloat>(
-    mut encoded_buffer: BufferMut<T::ALPInt>,
+fn decompress_chunked(
+    encoded_ref: ArrayRef,
     exponents: Exponents,
-    validity: Validity,
     patches: &Patches,
     patches_chunk_offsets: &PrimitiveArray,
-) -> (BufferMut<T>, MaskMut) {
-    let array_len = encoded_buffer.len();
-    let validity_mask = validity.to_mask(array_len).into_mut();
-    let patches_offset = patches.offset();
+    dtype: DType,
+) -> PrimitiveArray {
+    let encoded = encoded_ref.to_primitive();
+
+    let validity = encoded.validity().clone();
 
     let patches_indices = patches.indices().as_ref().to_primitive();
     let patches_values = patches.values().as_ref().to_primitive();
-    let patches_values_slice = patches_values.as_slice::<T>();
-
-    match_each_unsigned_integer_ptype!(patches_chunk_offsets.ptype(), |C| {
-        let patches_chunk_offsets_slice = patches_chunk_offsets.as_slice::<C>();
-        // Base offset is the first chunk offset
-        let base_offset = if !patches_chunk_offsets_slice.is_empty() {
-            AsPrimitive::<usize>::as_(patches_chunk_offsets_slice[0])
-        } else {
-            0
-        };
-
-        match_each_unsigned_integer_ptype!(patches_indices.ptype(), |I| {
-            let patches_indices_slice = patches_indices.as_slice::<I>();
-
-            // Process in chunks of 1024 for cache efficiency
-            for (chunk_idx, chunk_start) in (0..array_len).step_by(1024).enumerate() {
-                let chunk_end = (chunk_start + 1024).min(array_len);
-                let chunk_slice = &mut encoded_buffer.as_mut_slice()[chunk_start..chunk_end];
-
-                T::decode_slice_inplace(chunk_slice, exponents);
-
-                // SAFETY: After decoding, T::ALPInt values are valid T values
-                let decoded_chunk: &mut [T] = unsafe { transmute(chunk_slice) };
-
-                // Apply patches for this chunk using the original patch_chunk function
-                let offset_within_chunk = patches.offset_within_chunk().unwrap_or(0);
-                patch_chunk(
-                    decoded_chunk,
-                    patches_indices_slice,
-                    patches_values_slice,
-                    patches_offset,
-                    patches_chunk_offsets_slice,
-                    chunk_idx,
-                    base_offset,
-                    offset_within_chunk,
-                );
-            }
-        })
-    });
+    let ptype = dtype.as_ptype();
+    let array_len = encoded_ref.len();
+    let patches_offset = patches.offset();
 
-    // SAFETY: After decoding, T::ALPInt buffer contains valid T values
-    let decoded_buffer = unsafe { transmute::<BufferMut<T::ALPInt>, BufferMut<T>>(encoded_buffer) };
-    (decoded_buffer, validity_mask)
+    // We need to drop ALPArray here in case converting encoded buffer into
+    // primitive didn't create a copy. In that case both alp_encoded and array
+    // will hold a reference to the buffer we want to mutate.
+    drop(encoded_ref);
+
+    match_each_alp_float_ptype!(ptype, |T| {
+        let patches_values = patches_values.as_slice::<T>();
+        let mut alp_buffer = encoded.into_buffer_mut();
+        match_each_unsigned_integer_ptype!(patches_chunk_offsets.ptype(), |C| {
+            let patches_chunk_offsets = patches_chunk_offsets.as_slice::<C>();
+            // There always is at least one chunk offset.
+            let base_offset = patches_chunk_offsets[0];
+            let offset_within_chunk = patches.offset_within_chunk().unwrap_or(0);
+
+            match_each_unsigned_integer_ptype!(patches_indices.ptype(), |I| {
+                let patches_indices = patches_indices.as_slice::<I>();
+
+                for (chunk_idx, chunk_start) in (0..array_len).step_by(1024).enumerate() {
+                    let chunk_end = (chunk_start + 1024).min(array_len);
+                    let chunk_slice = &mut alp_buffer.as_mut_slice()[chunk_start..chunk_end];
+
+                    <T>::decode_slice_inplace(chunk_slice, exponents);
+
+                    let decoded_chunk: &mut [T] = unsafe { transmute(chunk_slice) };
+                    patch_chunk(
+                        decoded_chunk,
+                        patches_indices,
+                        patches_values,
+                        patches_offset,
+                        patches_chunk_offsets,
+                        chunk_idx,
+                        base_offset.as_(),
+                        offset_within_chunk,
+                    );
+                }
+
+                let decoded_buffer: BufferMut<T> = unsafe { transmute(alp_buffer) };
+                PrimitiveArray::new::<T>(decoded_buffer.freeze(), validity)
+            })
+        })
+    })
 }
 
 /// Decompresses an ALP-encoded array without chunk offsets.
 ///
-/// This function decodes the complete buffer at once and applies any patches at the end in one go.
-fn decompress_unchunked_core<T: ALPFloat>(
-    encoded_buffer: BufferMut<T::ALPInt>,
+/// This function decodes the complete array at once and then applies any patches after.
+fn decompress_unchunked(
+    encoded_ref: ArrayRef,
     exponents: Exponents,
-    validity: Validity,
-    patches: Option<&Patches>,
-) -> (BufferMut<T>, MaskMut) {
-    // Decode entire buffer.
-    let mut decoded_buffer = T::decode_buffer(encoded_buffer, exponents);
-    let mut validity_mask = validity.to_mask(decoded_buffer.len()).into_mut();
-
-    // Apply patches if present
+    patches: Option<Patches>,
+    dtype: DType,
+) -> PrimitiveArray {
+    let encoded = encoded_ref.to_primitive();
+
+    // We need to drop ALPArray here in case converting encoded buffer into
+    // primitive didn't create a copy. In that case both alp_encoded and array
+    // will hold a reference to the buffer we want to mutate.
+    drop(encoded_ref);
+
+    let validity = encoded.validity().clone();
+    let ptype = dtype.as_ptype();
+
+    let decoded = match_each_alp_float_ptype!(ptype, |T| {
+        PrimitiveArray::new::<T>(
+            <T>::decode_buffer(encoded.into_buffer_mut(), exponents),
+            validity,
+        )
+    });
+
     if let Some(patches) = patches {
-        // SAFETY: Patches invariants guarantee indices are valid.
-        unsafe {
-            patches.apply_to_buffer(decoded_buffer.as_mut_slice(), &mut validity_mask);
-        }
+        decoded.patch(&patches)
+    } else {
+        decoded
     }
-
-    (decoded_buffer, validity_mask)
 }
 
 #[cfg(test)]
 mod tests {
+    use vortex_array::validity::Validity;
     use vortex_buffer::buffer;
     use vortex_vector::VectorMutOps;
 
@@ -355,19 +336,4 @@ mod tests {
         let expected_f64 = decompress(f64_encoded);
         assert_eq!(expected_f64.as_slice::<f64>(), f64_vector.as_ref());
     }
-
-    #[test]
-    fn test_components_api() {
-        let array = PrimitiveArray::from_iter([1.5f32, 2.5, 3.5]);
-        let encoded = alp_encode(&array, None).unwrap();
-
-        let (buffer, validity) = decompress_to_components::<f32>(encoded);
-        assert_eq!(buffer.len(), 3);
-        assert_eq!(validity.len(), 3);
-
-        // All should be valid
-        for i in 0..3 {
-            assert!(validity.value(i));
-        }
-    }
 }