indices as child instead of buffer

a10y · a10y · commit df8dbf595052 · 2026-03-26T15:49:17.000-04:00
Signed-off-by: Andrew Duffy &lt;andrew@a10y.dev&gt;
diff --git a/vortex-array/src/arrays/patched/array.rs b/vortex-array/src/arrays/patched/array.rs
@@ -14,12 +14,12 @@ use crate::DynArray;
 use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::arrays::PrimitiveArray;
-use crate::arrays::patched::PatchAccessor;
 use crate::arrays::patched::TransposedPatches;
 use crate::arrays::patched::patch_lanes;
 use crate::buffer::BufferHandle;
 use crate::dtype::IntegerPType;
 use crate::dtype::NativePType;
+use crate::dtype::PType;
 use crate::match_each_native_ptype;
 use crate::match_each_unsigned_integer_ptype;
 use crate::patches::Patches;
@@ -48,7 +48,7 @@ pub struct PatchedArray {
     /// lane offsets. The PType of these MUST be u32
     pub(super) lane_offsets: BufferHandle,
     /// indices within a 1024-element chunk. The PType of these MUST be u16
-    pub(super) indices: BufferHandle,
+    pub(super) indices: ArrayRef,
     /// patch values corresponding to the indices. The ptype is specified by `values_ptype`.
     pub(super) values: ArrayRef,
 
@@ -86,6 +86,12 @@ impl PatchedArray {
             values,
         } = transpose_patches(patches, ctx)?;
 
+        let indices = PrimitiveArray::from_buffer_handle(
+            BufferHandle::new_host(indices),
+            PType::U16,
+            Validity::NonNullable,
+        )
+        .into_array();
         let values = PrimitiveArray::from_buffer_handle(
             BufferHandle::new_host(values),
             values_ptype,
@@ -102,19 +108,24 @@ impl PatchedArray {
             offset: 0,
             len,
             lane_offsets: BufferHandle::new_host(lane_offsets),
-            indices: BufferHandle::new_host(indices),
+            indices,
             values,
             stats_set: ArrayStats::default(),
         })
     }
+}
 
-    /// Get an accessor, which allows ranged access to patches by chunk/lane.
-    pub fn accessor(&self) -> PatchAccessor<'_> {
-        PatchAccessor {
-            n_lanes: self.n_lanes,
-            lane_offsets: self.lane_offsets.as_host().reinterpret::<u32>(),
-            indices: self.indices.as_host().reinterpret::<u16>(),
-        }
+impl PatchedArray {
+    pub(crate) fn seek_to_lane(&self, chunk: usize, lane: usize) -> Range<usize> {
+        assert!(chunk < self.n_chunks);
+        assert!(lane < self.n_lanes);
+
+        let lane_offsets = self.lane_offsets.as_host().reinterpret::<u32>();
+
+        let start = lane_offsets[chunk * self.n_lanes + lane] as usize;
+        let stop = lane_offsets[chunk * self.n_lanes + lane + 1] as usize;
+
+        start..stop
     }
 
     /// Slice the array to just the patches and inner values that are within the chunk range.
diff --git a/vortex-array/src/arrays/patched/compute/compare.rs b/vortex-array/src/arrays/patched/compute/compare.rs
@@ -93,10 +93,11 @@ impl CompareKernel for Patched {
         }
 
         let lane_offsets = lhs.lane_offsets.as_host().reinterpret::<u32>();
-        let indices = lhs.indices.as_host().reinterpret::<u16>();
+        let indices = lhs.indices.clone().execute::<PrimitiveArray>(ctx)?;
         let values = lhs.values.clone().execute::<PrimitiveArray>(ctx)?;
 
         match_each_native_ptype!(values.ptype(), |V| {
+            let indices = indices.as_slice::<u16>();
             let values = values.as_slice::<V>();
             let constant = constant
                 .as_primitive()
diff --git a/vortex-array/src/arrays/patched/compute/take.rs b/vortex-array/src/arrays/patched/compute/take.rs
@@ -46,7 +46,8 @@ impl TakeExecute for Patched {
         match_each_unsigned_integer_ptype!(indices_ptype, |I| {
             match_each_native_ptype!(ptype, |V| {
                 let indices = indices.clone().execute::<PrimitiveArray>(ctx)?;
-                let values = array.values.clone().execute::<PrimitiveArray>(ctx)?;
+                let patch_indices = array.indices.clone().execute::<PrimitiveArray>(ctx)?;
+                let patch_values = array.values.clone().execute::<PrimitiveArray>(ctx)?;
                 let mut output = Buffer::<V>::from_byte_buffer(buffer.unwrap_host()).into_mut();
                 take_map(
                     output.as_mut(),
@@ -56,8 +57,8 @@ impl TakeExecute for Patched {
                     array.n_chunks,
                     array.n_lanes,
                     array.lane_offsets.as_host().reinterpret::<u32>(),
-                    array.indices.as_host().reinterpret::<u16>(),
-                    values.as_slice::<V>(),
+                    patch_indices.as_slice::<u16>(),
+                    patch_values.as_slice::<V>(),
                 );
 
                 // SAFETY: output and validity still have same length after take_map returns.
diff --git a/vortex-array/src/arrays/patched/mod.rs b/vortex-array/src/arrays/patched/mod.rs
@@ -31,27 +31,3 @@ const fn patch_lanes<V: Sized>() -> usize {
     // from shared to global memory.
     if size_of::<V>() < 8 { 32 } else { 16 }
 }
-
-/// A cached accessor to the patch values.
-pub struct PatchAccessor<'a> {
-    n_lanes: usize,
-    lane_offsets: &'a [u32],
-    indices: &'a [u16],
-}
-
-impl<'a> PatchAccessor<'a> {
-    /// Get an iterator over indices and values offsets.
-    ///
-    /// The first component is the index into the `indices` and `values`, and the second is
-    /// the datum from `indices[index]` already prefetched.
-    pub fn offsets_iter(
-        &self,
-        chunk: usize,
-        lane: usize,
-    ) -> impl Iterator<Item = (usize, u16)> + '_ {
-        let start = self.lane_offsets[chunk * self.n_lanes + lane] as usize;
-        let stop = self.lane_offsets[chunk * self.n_lanes + lane + 1] as usize;
-
-        std::iter::zip(start..stop, self.indices[start..stop].iter().copied())
-    }
-}
diff --git a/vortex-array/src/arrays/patched/vtable/mod.rs b/vortex-array/src/arrays/patched/vtable/mod.rs
@@ -39,6 +39,7 @@ use crate::builders::ArrayBuilder;
 use crate::builders::PrimitiveBuilder;
 use crate::dtype::DType;
 use crate::dtype::NativePType;
+use crate::dtype::PType;
 use crate::match_each_native_ptype;
 use crate::serde::ArrayChildren;
 use crate::stats::ArrayStats;
@@ -64,6 +65,8 @@ impl ValidityChild<Patched> for Patched {
 pub struct PatchedMetadata {
     #[prost(uint32, tag = "1")]
     pub(crate) offset: u32,
+    #[prost(uint32, tag = "2")]
+    pub(crate) n_patches: u32,
 }
 
 impl VTable for Patched {
@@ -111,21 +114,19 @@ impl VTable for Patched {
     }
 
     fn nbuffers(_array: &Self::Array) -> usize {
-        3
+        1
     }
 
     fn buffer(array: &Self::Array, idx: usize) -> BufferHandle {
         match idx {
             0 => array.lane_offsets.clone(),
-            1 => array.indices.clone(),
             _ => vortex_panic!("invalid buffer index for PatchedArray: {idx}"),
         }
     }
 
     fn buffer_name(_array: &Self::Array, idx: usize) -> Option<String> {
         match idx {
             0 => Some("lane_offsets".to_string()),
-            1 => Some("patch_indices".to_string()),
             _ => vortex_panic!("invalid buffer index for PatchedArray: {idx}"),
         }
     }
@@ -137,15 +138,17 @@ impl VTable for Patched {
     fn child(array: &Self::Array, idx: usize) -> ArrayRef {
         match idx {
             0 => array.inner.clone(),
-            1 => array.values.clone(),
+            1 => array.indices.clone(),
+            2 => array.values.clone(),
             _ => vortex_panic!("invalid buffer index for PatchedArray: {idx}"),
         }
     }
 
     fn child_name(_array: &Self::Array, idx: usize) -> String {
         match idx {
             0 => "inner".to_string(),
-            1 => "patch_values".to_string(),
+            1 => "patch_indices".to_string(),
+            2 => "patch_values".to_string(),
             _ => vortex_panic!("invalid buffer index for PatchedArray: {idx}"),
         }
     }
@@ -154,6 +157,7 @@ impl VTable for Patched {
     fn metadata(array: &Self::Array) -> VortexResult<Self::Metadata> {
         Ok(ProstMetadata(PatchedMetadata {
             offset: array.offset as u32,
+            n_patches: array.indices.len() as u32,
         }))
     }
 
@@ -198,7 +202,7 @@ impl VTable for Patched {
         let offset = array.offset;
         let lane_offsets: Buffer<u32> =
             Buffer::from_byte_buffer(array.lane_offsets.clone().unwrap_host());
-        let indices: Buffer<u16> = Buffer::from_byte_buffer(array.indices.clone().unwrap_host());
+        let indices = array.indices.clone().execute::<PrimitiveArray>(ctx)?;
         let values = array.values.clone().execute::<PrimitiveArray>(ctx)?;
 
         match_each_native_ptype!(ptype, |V| {
@@ -219,7 +223,7 @@ impl VTable for Patched {
                 array.n_chunks,
                 array.n_lanes,
                 &lane_offsets,
-                &indices,
+                indices.as_slice::<u16>(),
                 values.as_slice::<V>(),
             );
         });
@@ -234,19 +238,16 @@ impl VTable for Patched {
         buffers: &[BufferHandle],
         children: &dyn ArrayChildren,
     ) -> VortexResult<PatchedArray> {
-        let inner = children.get(0, dtype, len)?;
-
         let n_chunks = len.div_ceil(1024);
-
         let n_lanes = match_each_native_ptype!(dtype.as_ptype(), |P| { patch_lanes::<P>() });
 
-        let &[lane_offsets, indices] = &buffers else {
+        let &[lane_offsets] = &buffers else {
             vortex_bail!("invalid buffer count for PatchedArray");
         };
 
-        // values and indices should have same len.
-        let expected_len = indices.as_host().reinterpret::<u16>().len();
-        let values = children.get(1, dtype, expected_len)?;
+        let inner = children.get(0, dtype, len)?;
+        let indices = children.get(1, PType::U16.into(), metadata.n_patches as usize)?;
+        let values = children.get(1, dtype, metadata.n_patches as usize)?;
 
         Ok(PatchedArray {
             inner,
@@ -255,7 +256,7 @@ impl VTable for Patched {
             offset: metadata.offset as usize,
             len,
             lane_offsets: lane_offsets.clone(),
-            indices: indices.clone(),
+            indices,
             values,
             stats_set: ArrayStats::default(),
         })
@@ -288,10 +289,10 @@ impl VTable for Patched {
 
         let lane_offsets: Buffer<u32> =
             Buffer::from_byte_buffer(array.lane_offsets.clone().unwrap_host());
-        let indices: Buffer<u16> = Buffer::from_byte_buffer(array.indices.clone().unwrap_host());
-        let values = array.values.clone().execute::<PrimitiveArray>(ctx)?;
+        let indices = array.indices.clone().execute::<PrimitiveArray>(ctx)?;
 
-        // TODO(aduffy): add support for non-primitive PatchedArray patches application.
+        // TODO(aduffy): add support for non-primitive PatchedArray patches application (?)
+        let values = array.values.clone().execute::<PrimitiveArray>(ctx)?;
 
         let patched_values = match_each_native_ptype!(values.ptype(), |V| {
             let mut output = Buffer::<V>::from_byte_buffer(buffer.unwrap_host()).into_mut();
@@ -306,7 +307,7 @@ impl VTable for Patched {
                 array.n_chunks,
                 array.n_lanes,
                 &lane_offsets,
-                &indices,
+                indices.as_slice::<u16>(),
                 values.as_slice::<V>(),
             );
 
diff --git a/vortex-array/src/arrays/patched/vtable/operations.rs b/vortex-array/src/arrays/patched/vtable/operations.rs
@@ -4,16 +4,23 @@
 use vortex_error::VortexResult;
 
 use crate::DynArray;
+use crate::ExecutionCtx;
+use crate::arrays::PrimitiveArray;
 use crate::arrays::patched::Patched;
 use crate::arrays::patched::PatchedArray;
 use crate::arrays::patched::patch_lanes;
 use crate::dtype::PType;
 use crate::match_each_native_ptype;
+use crate::optimizer::ArrayOptimizer;
 use crate::scalar::Scalar;
 use crate::vtable::OperationsVTable;
 
 impl OperationsVTable<Patched> for Patched {
-    fn scalar_at(array: &PatchedArray, index: usize) -> VortexResult<Scalar> {
+    fn scalar_at(
+        array: &PatchedArray,
+        index: usize,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Scalar> {
         // First check the patches
         let chunk = index / 1024;
         #[allow(clippy::cast_possible_truncation)]
@@ -22,11 +29,20 @@ impl OperationsVTable<Patched> for Patched {
         let values_ptype = PType::try_from(array.dtype())?;
 
         let lane = match_each_native_ptype!(values_ptype, |V| { index % patch_lanes::<V>() });
-        let accessor = array.accessor();
+
+        let range = array.seek_to_lane(chunk, lane);
+
+        // Get the range of indices corresponding to the lane, potentially decoding them to avoid
+        // the overhead of repeated scalar_at calls.
+        let patch_indices = array
+            .indices
+            .slice(range.clone())?
+            .optimize()?
+            .execute::<PrimitiveArray>(ctx)?;
 
         // NOTE: we do linear scan as lane has <= 32 patches, binary search would likely
         //  be slower.
-        for (index, patch_index) in accessor.offsets_iter(chunk, lane) {
+        for (&patch_index, index) in std::iter::zip(patch_indices.as_slice::<u16>(), range) {
             if patch_index == chunk_index {
                 return array.values.scalar_at(index)?.cast(array.dtype());
             }
