fix: patch_chunk offset calc (#5222)

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

Author: 0ax1

encodings/alp/src/alp/compress.rs

@@ -4,6 +4,7 @@
 use std::mem::transmute;
 
 use itertools::Itertools;
+use num_traits::AsPrimitive;
 use vortex_array::arrays::{PrimitiveArray, patch_chunk};
 use vortex_array::patches::Patches;
 use vortex_array::validity::Validity;
@@ -169,21 +170,27 @@ pub fn decompress_chunked(
     let exponents = array.exponents();
     let patches_offset = patches.offset();
 
-    // 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.
+    // 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);
 
     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) };
@@ -194,6 +201,8 @@ pub fn decompress_chunked(
                         patches_offset,
                         patches_chunk_offsets,
                         chunk_idx,
+                        base_offset.as_(),
+                        offset_within_chunk,
                     );
                 }
 
@@ -214,8 +223,9 @@ fn decompress_unchunked(array: ALPArray) -> PrimitiveArray {
     let exponents = array.exponents();
     let ptype = array.dtype().as_ptype();
 
-    // 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.
+    // 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 decoded = match_each_alp_float_ptype!(ptype, |T| {
@@ -503,6 +513,27 @@ mod tests {
         assert!(encoded.patches().is_some());
     }
 
+    #[test]
+    fn test_slice_across_chunks_with_patches_roundtrip() {
+        let mut values = vec![1.0f64; 2048];
+        values[100] = PI;
+        values[200] = E;
+        values[600] = 42.42;
+        values[800] = 42.42;
+        values[1000] = 42.42;
+        values[1023] = 42.42;
+
+        let original = PrimitiveArray::new(Buffer::from(values), Validity::NonNullable);
+        let encoded = alp_encode(&original, None).unwrap();
+
+        let sliced_alp = encoded.slice(1023..1025);
+        let decoded = sliced_alp.to_primitive();
+
+        let expected_slice = original.slice(1023..1025).to_primitive();
+        assert_eq!(expected_slice.as_slice::<f64>(), decoded.as_slice::<f64>());
+        assert!(encoded.patches().is_some());
+    }
+
     #[test]
     fn test_slice_half_chunk_nullable_roundtrip() {
         let values = (0..1024)

vortex-array/src/arrays/primitive/array/patch.rs

@@ -67,22 +67,29 @@ impl PrimitiveArray {
 /// * `patches_offset` - Offset to subtract from patch indices
 /// * `chunk_offsets_slice` - Slice containing offsets for each chunk
 /// * `chunk_idx` - Index of the chunk to patch
+/// * `base_offset` - Base offset from the first chunk
+/// * `offset_within_chunk` - Offset within chunk for sliced patches
 #[inline]
+#[allow(clippy::too_many_arguments)]
 pub fn patch_chunk<T, I, C>(
     decoded_values: &mut [T],
     patches_indices: &[I],
     patches_values: &[T],
     patches_offset: usize,
     chunk_offsets_slice: &[C],
     chunk_idx: usize,
+    base_offset: usize,
+    offset_within_chunk: usize,
 ) where
     T: NativePType,
     I: UnsignedPType,
     C: UnsignedPType,
 {
-    let patches_start_idx = chunk_offsets_slice[chunk_idx].as_();
+    // Use the same logic as patches slice implementation for calculating patch ranges.
+    let patches_start_idx =
+        (chunk_offsets_slice[chunk_idx].as_() - base_offset).saturating_sub(offset_within_chunk);
     let patches_end_idx = if chunk_idx + 1 < chunk_offsets_slice.len() {
-        chunk_offsets_slice[chunk_idx + 1].as_()
+        chunk_offsets_slice[chunk_idx + 1].as_() - base_offset - offset_within_chunk
     } else {
         patches_indices.len()
     };

vortex-array/src/patches.rs

@@ -263,6 +263,11 @@ impl Patches {
             .vortex_expect("chunk offset must be usize")
     }
 
+    #[inline]
+    pub fn offset_within_chunk(&self) -> Option<usize> {
+        self.offset_within_chunk
+    }
+
     #[inline]
     pub fn indices_ptype(&self) -> PType {
         PType::try_from(self.indices.dtype()).vortex_expect("primitive indices")