save

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

Author: a10y

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

@@ -76,6 +76,11 @@ impl PatchedArray {
             "Creating PatchedArray from Patches only supported for primitive arrays"
         );
 
+        vortex_ensure!(
+            patches.num_patches() <= u32::MAX as usize,
+            "PatchedArray does not support > u32::MAX patch values"
+        );
+
         let values_ptype = patches.dtype().as_ptype();
 
         let TransposedPatches {

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

@@ -14,6 +14,7 @@ use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_ensure;
+use vortex_error::vortex_err;
 use vortex_error::vortex_panic;
 use vortex_session::VortexSession;
 
@@ -65,10 +66,29 @@ impl ValidityChild<Patched> for Patched {
 
 #[derive(Clone, prost::Message)]
 pub struct PatchedMetadata {
-    #[prost(uint32, tag = "1")]
-    pub(crate) offset: u32,
+    /// Length of the `inner` child.
+    ///
+    /// This may not match the length of the wrapping PatchedArray, if for example
+    /// in a filter or slice it may be sliced to the nearest chunk boundary.
+    #[prost(uint64, tag = "1")]
+    pub(crate) inner_len: u64,
+
+    /// Offset within the first chunk of `inner` where data begins.
+    ///
+    /// This may become nonzero after slicing.
     #[prost(uint32, tag = "2")]
+    pub(crate) offset: u32,
+
+    /// Number of patches. This is the length of the `indices` and `values` children.
+    #[prost(uint32, tag = "3")]
     pub(crate) n_patches: u32,
+
+    /// Number of lanes the patches get spread over.
+    ///
+    /// By default, this is either 16 or 32 depending on the width of the type, but may change
+    /// in the future, so we save it on write.
+    #[prost(uint32, tag = "4")]
+    pub(crate) n_lanes: u32,
 }
 
 impl VTable for Patched {
@@ -158,8 +178,10 @@ impl VTable for Patched {
     #[allow(clippy::cast_possible_truncation)]
     fn metadata(array: &Self::Array) -> VortexResult<Self::Metadata> {
         Ok(ProstMetadata(PatchedMetadata {
+            inner_len: array.inner.len() as u64,
             offset: array.offset as u32,
             n_patches: array.indices.len() as u32,
+            n_lanes: array.n_lanes as u32,
         }))
     }
 
@@ -243,22 +265,32 @@ impl VTable for Patched {
         buffers: &[BufferHandle],
         children: &dyn ArrayChildren,
     ) -> VortexResult<PatchedArray> {
-        let n_chunks = len.div_ceil(1024);
-        let n_lanes = match_each_native_ptype!(dtype.as_ptype(), |P| { patch_lanes::<P>() });
+        let inner_len = usize::try_from(metadata.inner_len).map_err(|_| {
+            vortex_err!(
+                "PatchedMetadata inner_len overflows usize: {}",
+                metadata.inner_len
+            )
+        })?;
+        let offset = metadata.offset as usize;
+
+        // n_chunks should correspond to the chunk in the `inner`.
+        // After slicing when offset > 0, there may be additional chunks.
+        let n_chunks = (len + offset).div_ceil(1024);
+        let n_lanes = metadata.n_lanes as usize;
 
         let &[lane_offsets] = &buffers else {
             vortex_bail!("invalid buffer count for PatchedArray");
         };
 
-        let inner = children.get(0, dtype, len)?;
+        let inner = children.get(0, dtype, inner_len)?;
         let indices = children.get(1, PType::U16.into(), metadata.n_patches as usize)?;
         let values = children.get(2, dtype, metadata.n_patches as usize)?;
 
         Ok(PatchedArray {
             inner,
             n_chunks,
             n_lanes,
-            offset: metadata.offset as usize,
+            offset,
             len,
             lane_offsets: lane_offsets.clone(),
             indices,