Convert offsets into subtraction in RunEnd and Patches

Signed-off-by: Robert Kruszewski <github@robertk.io>

Author: robert3005

encodings/alp/src/alp/array.rs

@@ -181,7 +181,9 @@ impl VTable for ALP {
                     .map(|dtype| children.get(3, &dtype, usize::try_from(p.chunk_offsets_len())?))
                     .transpose()?;
 
-                Patches::new(len, p.offset()?, indices, values, chunk_offsets)
+                // The indices child already has the offset embedded (via Binary(Sub) expression),
+                // so we pass offset 0 to avoid double-wrapping.
+                Patches::new(len, 0, indices, values, chunk_offsets)
             })
             .transpose()?;
 
@@ -200,10 +202,10 @@ impl VTable for ALP {
         let patches_info = array
             .patches
             .as_ref()
-            .map(|p| (p.array_len(), p.offset(), p.chunk_offsets().is_some()));
+            .map(|p| (p.array_len(), p.chunk_offsets().is_some()));
 
         let expected_children = match &patches_info {
-            Some((_, _, has_chunk_offsets)) => 1 + 2 + if *has_chunk_offsets { 1 } else { 0 },
+            Some((_, has_chunk_offsets)) => 1 + 2 + if *has_chunk_offsets { 1 } else { 0 },
             None => 1,
         };
 
@@ -219,7 +221,7 @@ impl VTable for ALP {
             .next()
             .ok_or_else(|| vortex_err!("Expected encoded child"))?;
 
-        if let Some((array_len, offset, _has_chunk_offsets)) = patches_info {
+        if let Some((array_len, _has_chunk_offsets)) = patches_info {
             let indices = children_iter
                 .next()
                 .ok_or_else(|| vortex_err!("Expected patch indices child"))?;
@@ -228,13 +230,9 @@ impl VTable for ALP {
                 .ok_or_else(|| vortex_err!("Expected patch values child"))?;
             let chunk_offsets = children_iter.next();
 
-            array.patches = Some(Patches::new(
-                array_len,
-                offset,
-                indices,
-                values,
-                chunk_offsets,
-            )?);
+            // The indices child already has the offset embedded (via Binary(Sub) expression),
+            // so we pass offset 0 to avoid double-wrapping.
+            array.patches = Some(Patches::new(array_len, 0, indices, values, chunk_offsets)?);
         }
 
         Ok(())
@@ -799,10 +797,11 @@ mod tests {
         );
 
         // Rebuild the patches WITHOUT chunk_offsets to simulate deserialized patches.
+        let (raw_indices, offset) = original_patches.raw_indices_and_offset();
         let patches_without_chunk_offsets = Patches::new(
             original_patches.array_len(),
-            original_patches.offset(),
-            original_patches.indices().clone(),
+            offset,
+            raw_indices.clone(),
             original_patches.values().clone(),
             None, // NO chunk_offsets - this triggers the bug!
         )

encodings/alp/src/alp/compute/cast.rs

@@ -31,10 +31,11 @@ impl CastReduce for ALP {
                     if p.values().dtype() == dtype {
                         Ok(p.clone())
                     } else {
+                        let (raw_indices, offset) = p.raw_indices_and_offset();
                         Patches::new(
                             p.array_len(),
-                            p.offset(),
-                            p.indices().clone(),
+                            offset,
+                            raw_indices.clone(),
                             p.values().cast(dtype.clone())?,
                             p.chunk_offsets().clone(),
                         )

encodings/alp/src/alp/decompress.rs

@@ -34,7 +34,8 @@ pub fn decompress_into_array(
     {
         let prim_encoded = encoded.execute::<PrimitiveArray>(ctx)?;
         let patches_chunk_offsets = chunk_offsets.clone().execute::<PrimitiveArray>(ctx)?;
-        let patches_indices = patches.indices().clone().execute::<PrimitiveArray>(ctx)?;
+        let (raw_indices, _) = patches.raw_indices_and_offset();
+        let patches_indices = raw_indices.clone().execute::<PrimitiveArray>(ctx)?;
         let patches_values = patches.values().clone().execute::<PrimitiveArray>(ctx)?;
         Ok(decompress_chunked_core(
             prim_encoded,
@@ -67,7 +68,8 @@ pub fn execute_decompress(array: ALPArray, ctx: &mut ExecutionCtx) -> VortexResu
         // TODO(joe): have into parts.
         let encoded = encoded.execute::<PrimitiveArray>(ctx)?;
         let patches_chunk_offsets = chunk_offsets.clone().execute::<PrimitiveArray>(ctx)?;
-        let patches_indices = patches.indices().clone().execute::<PrimitiveArray>(ctx)?;
+        let (raw_indices, _) = patches.raw_indices_and_offset();
+        let patches_indices = raw_indices.clone().execute::<PrimitiveArray>(ctx)?;
         let patches_values = patches.values().clone().execute::<PrimitiveArray>(ctx)?;
         Ok(decompress_chunked_core(
             encoded,

encodings/alp/src/alp_rd/array.rs

@@ -234,12 +234,10 @@ impl VTable for ALPRD {
                 let indices = children.get(2, &p.indices_dtype()?, p.len()?)?;
                 let values = children.get(3, &left_parts_dtype, p.len()?)?;
 
+                // The indices child already has the offset embedded (via Binary(Sub) expression),
+                // so we pass offset 0 to avoid double-wrapping.
                 Patches::new(
-                    len,
-                    p.offset()?,
-                    indices,
-                    values,
-                    // TODO(0ax1): handle chunk offsets
+                    len, 0, indices, values, // TODO(0ax1): handle chunk offsets
                     None,
                 )
             })
@@ -262,10 +260,7 @@ impl VTable for ALPRD {
 
     fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
         // Children: left_parts, right_parts, patches (if present): indices, values
-        let patches_info = array
-            .left_parts_patches
-            .as_ref()
-            .map(|p| (p.array_len(), p.offset()));
+        let patches_info = array.left_parts_patches.as_ref().map(|p| p.array_len());
 
         let expected_children = if patches_info.is_some() { 4 } else { 2 };
 
@@ -284,16 +279,18 @@ impl VTable for ALPRD {
             .next()
             .ok_or_else(|| vortex_err!("Expected right_parts child"))?;
 
-        if let Some((array_len, offset)) = patches_info {
+        if let Some(array_len) = patches_info {
             let indices = children_iter
                 .next()
                 .ok_or_else(|| vortex_err!("Expected patch indices child"))?;
             let values = children_iter
                 .next()
                 .ok_or_else(|| vortex_err!("Expected patch values child"))?;
 
+            // The indices child already has the offset embedded (via Binary(Sub) expression),
+            // so we pass offset 0 to avoid double-wrapping.
             array.left_parts_patches = Some(Patches::new(
-                array_len, offset, indices, values,
+                array_len, 0, indices, values,
                 None, // chunk_offsets not currently supported for ALPRD
             )?);
         }

encodings/alp/src/alp_rd/mod.rs

@@ -311,9 +311,10 @@ pub fn alp_rd_decode<T: ALPRDFloat>(
 
     // Apply any patches
     if let Some(patches) = left_parts_patches {
-        let indices = patches.indices().clone().execute::<PrimitiveArray>(ctx)?;
+        let (raw_indices, offset) = patches.raw_indices_and_offset();
+        let indices = raw_indices.clone().execute::<PrimitiveArray>(ctx)?;
         let patch_values = patches.values().clone().execute::<PrimitiveArray>(ctx)?;
-        alp_rd_apply_patches(&mut values, &indices, &patch_values, patches.offset());
+        alp_rd_apply_patches(&mut values, &indices, &patch_values, offset);
     }
 
     // Shift the left-parts and add in the right-parts.

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

@@ -155,7 +155,8 @@ pub fn apply_patches_to_uninit_range_fn<T: NativePType, F: Fn(T) -> T>(
 ) -> VortexResult<()> {
     assert_eq!(patches.array_len(), dst.len());
 
-    let indices = patches.indices().clone().execute::<PrimitiveArray>(ctx)?;
+    let (raw_indices, offset) = patches.raw_indices_and_offset();
+    let indices = raw_indices.clone().execute::<PrimitiveArray>(ctx)?;
     let values = patches.values().clone().execute::<PrimitiveArray>(ctx)?;
     let validity = values.validity_mask()?;
     let values = values.as_slice::<T>();
@@ -166,7 +167,7 @@ pub fn apply_patches_to_uninit_range_fn<T: NativePType, F: Fn(T) -> T>(
             indices.as_slice::<P>(),
             values,
             validity,
-            patches.offset(),
+            offset,
             f,
         )
     });

encodings/fastlanes/src/bitpacking/compute/cast.rs

@@ -29,10 +29,11 @@ impl CastReduce for BitPacked {
                         .patches()
                         .map(|patches| {
                             let new_values = patches.values().cast(dtype.clone())?;
+                            let (raw_indices, offset) = patches.raw_indices_and_offset();
                             Patches::new(
                                 patches.array_len(),
-                                patches.offset(),
-                                patches.indices().clone(),
+                                offset,
+                                raw_indices.clone(),
                                 new_values,
                                 patches.chunk_offsets().clone(),
                             )

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

@@ -173,12 +173,14 @@ impl VTable for BitPacked {
 
     fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
         // Children: patches (if present): indices, values, chunk_offsets; then validity (if present)
-        let patches_info = array
+        let has_patches = array.patches().is_some();
+        let has_chunk_offsets = array
             .patches()
-            .map(|p| (p.offset(), p.chunk_offsets().is_some()));
+            .map(|p| p.chunk_offsets().is_some())
+            .unwrap_or(false);
 
         let mut child_idx = 0;
-        let patches = if let Some((patch_offset, has_chunk_offsets)) = patches_info {
+        let patches = if has_patches {
             let patch_indices = children
                 .get(child_idx)
                 .ok_or_else(|| vortex_err!("Expected patch_indices child at index {}", child_idx))?
@@ -204,9 +206,11 @@ impl VTable for BitPacked {
                 None
             };
 
+            // The indices child already has the offset embedded (via Binary(Sub) expression),
+            // so we pass offset 0 to avoid double-wrapping.
             Some(Patches::new(
                 array.len(),
-                patch_offset,
+                0,
                 patch_indices,
                 patch_values,
                 patch_chunk_offsets,
@@ -317,7 +321,9 @@ impl VTable for BitPacked {
                     .map(|dtype| children.get(2, &dtype, p.chunk_offsets_len() as usize))
                     .transpose()?;
 
-                Patches::new(len, p.offset()?, indices, values, chunk_offsets)
+                // The indices child already has the offset embedded (via Binary(Sub) expression),
+                // so we pass offset 0 to avoid double-wrapping.
+                Patches::new(len, 0, indices, values, chunk_offsets)
             })
             .transpose()?;
 

encodings/runend/public-api.lock

@@ -124,10 +124,12 @@ pub fn vortex_runend::RunEndArray::into_parts(self) -> vortex_runend::RunEndArra
 
 pub fn vortex_runend::RunEndArray::new(ends: vortex_array::array::ArrayRef, values: vortex_array::array::ArrayRef) -> Self
 
-pub unsafe fn vortex_runend::RunEndArray::new_unchecked(ends: vortex_array::array::ArrayRef, values: vortex_array::array::ArrayRef, offset: usize, length: usize) -> Self
+pub unsafe fn vortex_runend::RunEndArray::new_unchecked(ends: vortex_array::array::ArrayRef, values: vortex_array::array::ArrayRef, length: usize) -> Self
 
 pub fn vortex_runend::RunEndArray::offset(&self) -> usize
 
+pub fn vortex_runend::RunEndArray::raw_ends_and_offset(&self) -> (&vortex_array::array::ArrayRef, usize)
+
 pub fn vortex_runend::RunEndArray::try_new(ends: vortex_array::array::ArrayRef, values: vortex_array::array::ArrayRef) -> vortex_error::VortexResult<Self>
 
 pub fn vortex_runend::RunEndArray::try_new_offset_length(ends: vortex_array::array::ArrayRef, values: vortex_array::array::ArrayRef, offset: usize, length: usize) -> vortex_error::VortexResult<Self>