use child for values instead of buffer

Signed-off-by: Andrew Duffy <andrew@a10y.dev>

Author: a10y

vortex-array/src/arrays/patched/array.rs

@@ -12,17 +12,19 @@ use crate::ArrayRef;
 use crate::Canonical;
 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;
 use crate::stats::ArrayStats;
+use crate::validity::Validity;
 
 /// An array that partially "patches" another array with new values.
 ///
@@ -50,14 +52,17 @@ pub struct PatchedArray {
     /// indices within a 1024-element chunk. The PType of these MUST be u16
     pub(super) indices: BufferHandle,
     /// patch values corresponding to the indices. The ptype is specified by `values_ptype`.
-    pub(super) values: BufferHandle,
-    /// PType of the scalars in `values`. Can be any native type.
-    pub(super) values_ptype: PType,
+    pub(super) values: ArrayRef,
 
     pub(super) stats_set: ArrayStats,
 }
 
 impl PatchedArray {
+    /// Create a new `PatchedArray` from a child array and a set of [`Patches`].
+    ///
+    /// # Errors
+    ///
+    /// The `inner` array must be primitive type, and it must have the same DType as the patches.
     pub fn from_array_and_patches(
         inner: ArrayRef,
         patches: &Patches,
@@ -68,6 +73,11 @@ impl PatchedArray {
             "array DType must match patches DType"
         );
 
+        vortex_ensure!(
+            inner.dtype().is_primitive(),
+            "Creating PatchedArray from Patches only supported for primitive arrays"
+        );
+
         let values_ptype = patches.dtype().as_ptype();
 
         let TransposedPatches {
@@ -80,27 +90,32 @@ impl PatchedArray {
 
         let len = inner.len();
 
+        let values = PrimitiveArray::from_buffer_handle(
+            BufferHandle::new_host(values),
+            values_ptype,
+            Validity::NonNullable,
+        )
+        .into_array();
+
         Ok(Self {
             inner,
             n_chunks,
             n_lanes,
-            values_ptype,
             offset: 0,
             len,
             lane_offsets: BufferHandle::new_host(lane_offsets),
             indices: BufferHandle::new_host(indices),
-            values: BufferHandle::new_host(values),
+            values,
             stats_set: ArrayStats::default(),
         })
     }
 
     /// Get an accessor, which allows ranged access to patches by chunk/lane.
-    pub fn accessor<V: NativePType>(&self) -> PatchAccessor<'_, V> {
+    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>(),
-            values: self.values.as_host().reinterpret::<V>(),
         }
     }
 
@@ -133,7 +148,6 @@ impl PatchedArray {
             len,
             indices,
             values,
-            values_ptype: self.values_ptype,
             lane_offsets: sliced_lane_offsets,
             stats_set: ArrayStats::default(),
         })

vortex-array/src/arrays/patched/compute/compare.rs

@@ -12,6 +12,7 @@ use crate::IntoArray;
 use crate::arrays::BoolArray;
 use crate::arrays::ConstantArray;
 use crate::arrays::Patched;
+use crate::arrays::PrimitiveArray;
 use crate::arrays::bool::BoolArrayParts;
 use crate::arrays::patched::patch_lanes;
 use crate::arrays::primitive::NativeValue;
@@ -28,6 +29,12 @@ impl CompareKernel for Patched {
         operator: CompareOperator,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<Option<ArrayRef>> {
+        // We only accelerate comparisons for primitives
+        if !lhs.dtype().is_primitive() {
+            return Ok(None);
+        }
+
+        // We only accelerate comparisons against constants
         let Some(constant) = rhs.as_constant() else {
             return Ok(None);
         };
@@ -87,9 +94,10 @@ impl CompareKernel for Patched {
 
         let lane_offsets = lhs.lane_offsets.as_host().reinterpret::<u32>();
         let indices = lhs.indices.as_host().reinterpret::<u16>();
+        let values = lhs.values.clone().execute::<PrimitiveArray>(ctx)?;
 
-        match_each_native_ptype!(lhs.values_ptype, |V| {
-            let values = lhs.values.as_host().reinterpret::<V>();
+        match_each_native_ptype!(values.ptype(), |V| {
+            let values = values.as_slice::<V>();
             let constant = constant
                 .as_primitive()
                 .as_::<V>()

vortex-array/src/arrays/patched/compute/take.rs

@@ -24,6 +24,11 @@ impl TakeExecute for Patched {
         indices: &ArrayRef,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<Option<ArrayRef>> {
+        // Only pushdown take when we have primitive types.
+        if !array.dtype().is_primitive() {
+            return Ok(None);
+        }
+
         // Perform take on the inner array, including the placeholders.
         let inner = array
             .inner
@@ -41,6 +46,7 @@ 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 mut output = Buffer::<V>::from_byte_buffer(buffer.unwrap_host()).into_mut();
                 take_map(
                     output.as_mut(),
@@ -51,7 +57,7 @@ impl TakeExecute for Patched {
                     array.n_lanes,
                     array.lane_offsets.as_host().reinterpret::<u32>(),
                     array.indices.as_host().reinterpret::<u16>(),
-                    array.values.as_host().reinterpret::<V>(),
+                    values.as_slice::<V>(),
                 );
 
                 // SAFETY: output and validity still have same length after take_map returns.

vortex-array/src/arrays/patched/mod.rs

@@ -32,23 +32,34 @@ const fn patch_lanes<V: Sized>() -> usize {
     if size_of::<V>() < 8 { 32 } else { 16 }
 }
 
-pub struct PatchAccessor<'a, V> {
+pub struct PatchAccessor<'a> {
     n_lanes: usize,
     lane_offsets: &'a [u32],
     indices: &'a [u16],
-    values: &'a [V],
 }
 
-impl<'a, V: Sized> PatchAccessor<'a, V> {
-    /// Access the patches for a particular lane
-    pub fn access(&'a self, chunk: usize, lane: usize) -> LanePatches<'a, V> {
+pub struct PatchOffset {
+    /// Global offset into the list of patches. These are some of the
+    pub index: usize,
+    /// This is the value stored in the `indices` buffer, which encodes the offset of the `index`-th
+    /// patch
+    pub chunk_offset: 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 component
+    /// is the set of values instead here...I think?
+    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;
 
-        LanePatches {
-            indices: &self.indices[start..stop],
-            values: &self.values[start..stop],
-        }
+        std::iter::zip(start..stop, self.indices[start..stop].iter().copied())
     }
 }
 

vortex-array/src/arrays/patched/vtable/mod.rs

@@ -86,7 +86,6 @@ impl VTable for Patched {
 
     fn array_hash<H: Hasher>(array: &Self::Array, state: &mut H, precision: Precision) {
         array.inner.array_hash(state, precision);
-        array.values_ptype.hash(state);
         array.n_chunks.hash(state);
         array.n_lanes.hash(state);
         array.lane_offsets.array_hash(state, precision);
@@ -97,7 +96,6 @@ impl VTable for Patched {
     fn array_eq(array: &Self::Array, other: &Self::Array, precision: Precision) -> bool {
         array.n_chunks == other.n_chunks
             && array.n_lanes == other.n_lanes
-            && array.values_ptype == other.values_ptype
             && array.inner.array_eq(&other.inner, precision)
             && array.lane_offsets.array_eq(&other.lane_offsets, precision)
             && array.indices.array_eq(&other.indices, precision)
@@ -112,7 +110,6 @@ impl VTable for Patched {
         match idx {
             0 => array.lane_offsets.clone(),
             1 => array.indices.clone(),
-            2 => array.values.clone(),
             _ => vortex_panic!("invalid buffer index for PatchedArray: {idx}"),
         }
     }
@@ -121,28 +118,27 @@ impl VTable for Patched {
         match idx {
             0 => Some("lane_offsets".to_string()),
             1 => Some("patch_indices".to_string()),
-            2 => Some("patch_values".to_string()),
             _ => vortex_panic!("invalid buffer index for PatchedArray: {idx}"),
         }
     }
 
     fn nchildren(_array: &Self::Array) -> usize {
-        1
+        2
     }
 
     fn child(array: &Self::Array, idx: usize) -> ArrayRef {
-        if idx == 0 {
-            array.inner.clone()
-        } else {
-            vortex_panic!("invalid child index for PatchedArray: {idx}");
+        match idx {
+            0 => array.inner.clone(),
+            1 => array.values.clone(),
+            _ => vortex_panic!("invalid buffer index for PatchedArray: {idx}"),
         }
     }
 
     fn child_name(_array: &Self::Array, idx: usize) -> String {
-        if idx == 0 {
-            "inner".to_string()
-        } else {
-            vortex_panic!("invalid child index for PatchedArray: {idx}");
+        match idx {
+            0 => "inner".to_string(),
+            1 => "patch_values".to_string(),
+            _ => vortex_panic!("invalid buffer index for PatchedArray: {idx}"),
         }
     }
 
@@ -181,10 +177,14 @@ impl VTable for Patched {
 
         let n_lanes = match_each_native_ptype!(dtype.as_ptype(), |P| { patch_lanes::<P>() });
 
-        let &[lane_offsets, indices, values] = &buffers else {
+        let &[lane_offsets, indices] = &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)?;
+
         Ok(PatchedArray {
             inner,
             n_chunks,
@@ -193,19 +193,19 @@ impl VTable for Patched {
             len,
             lane_offsets: lane_offsets.clone(),
             indices: indices.clone(),
-            values: values.clone(),
-            values_ptype: dtype.as_ptype(),
+            values,
             stats_set: ArrayStats::default(),
         })
     }
 
     fn with_children(array: &mut Self::Array, mut children: Vec<ArrayRef>) -> VortexResult<()> {
         vortex_ensure!(
-            children.len() == 1,
-            "PatchedArray must have exactly 1 child"
+            children.len() == 2,
+            "PatchedArray must have exactly 2 children"
         );
 
         array.inner = children.remove(0);
+        array.values = children.remove(0);
 
         Ok(())
     }
@@ -226,23 +226,25 @@ 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)?;
+
+        // TODO(aduffy): add support for non-primitive PatchedArray patches application.
 
-        let patched_values = match_each_native_ptype!(array.values_ptype, |V| {
+        let patched_values = match_each_native_ptype!(values.ptype(), |V| {
             let mut output = Buffer::<V>::from_byte_buffer(buffer.unwrap_host()).into_mut();
-            let values: Buffer<V> = Buffer::from_byte_buffer(array.values.clone().unwrap_host());
 
             let offset = array.offset;
             let len = array.len;
 
-            apply::<V>(
+            apply_patches_primitive::<V>(
                 &mut output,
                 offset,
                 len,
                 array.n_chunks,
                 array.n_lanes,
                 &lane_offsets,
                 &indices,
-                &values,
+                values.as_slice::<V>(),
             );
 
             // The output will always be aligned to a chunk boundary, we apply the offset/len
@@ -276,7 +278,7 @@ impl VTable for Patched {
 
 /// Apply patches on top of the existing value types.
 #[allow(clippy::too_many_arguments)]
-fn apply<V: NativePType>(
+fn apply_patches_primitive<V: NativePType>(
     output: &mut [V],
     offset: usize,
     len: usize,

vortex-array/src/arrays/patched/vtable/operations.rs

@@ -7,6 +7,7 @@ use crate::DynArray;
 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::scalar::Scalar;
 use crate::vtable::OperationsVTable;
@@ -17,21 +18,19 @@ impl OperationsVTable<Patched> for Patched {
         let chunk = index / 1024;
         #[allow(clippy::cast_possible_truncation)]
         let chunk_index = (index % 1024) as u16;
-        match_each_native_ptype!(array.values_ptype, |V| {
-            let lane = index % patch_lanes::<V>();
-            let accessor = array.accessor::<V>();
-            let patches = accessor.access(chunk, lane);
-            // NOTE: we do linear scan as lane has <= 32 patches, binary search would likely
-            //  be slower.
-            for (patch_index, patch_value) in patches.iter() {
-                if patch_index == chunk_index {
-                    return Ok(Scalar::primitive(
-                        patch_value,
-                        array.inner.dtype().nullability(),
-                    ));
-                }
+
+        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();
+
+        // 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) {
+            if patch_index == chunk_index {
+                return array.values.scalar_at(index);
             }
-        });
+        }
 
         // Otherwise, access the underlying value.
         array.inner.scalar_at(index)