chore: have on demand validity and patches for array remove slot extraction (#7217)

This PR get rid of holding validity and patches on each ArrayData,
instead we construct this on-demand.

This will also a follow up to have `slots_mut -> &mut
[Option<ArrayRef>]` instead of `with_slots`

---------

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>
Signed-off-by: Will Manning <will@willmanning.io>

Author: joseph-isaacs

encodings/alp/public-api.lock

@@ -118,7 +118,7 @@ pub fn vortex_alp::ALPArray::into_parts(self) -> (vortex_array::array::ArrayRef,
 
 pub fn vortex_alp::ALPArray::new(encoded: vortex_array::array::ArrayRef, exponents: vortex_alp::Exponents, patches: core::option::Option<vortex_array::patches::Patches>) -> Self
 
-pub fn vortex_alp::ALPArray::patches(&self) -> core::option::Option<&vortex_array::patches::Patches>
+pub fn vortex_alp::ALPArray::patches(&self) -> core::option::Option<vortex_array::patches::Patches>
 
 pub fn vortex_alp::ALPArray::ptype(&self) -> vortex_array::dtype::ptype::PType
 
@@ -278,7 +278,7 @@ pub fn vortex_alp::ALPRDArray::left_parts(&self) -> &vortex_array::array::ArrayR
 
 pub fn vortex_alp::ALPRDArray::left_parts_dictionary(&self) -> &vortex_buffer::buffer::Buffer<u16>
 
-pub fn vortex_alp::ALPRDArray::left_parts_patches(&self) -> core::option::Option<&vortex_array::patches::Patches>
+pub fn vortex_alp::ALPRDArray::left_parts_patches(&self) -> core::option::Option<vortex_array::patches::Patches>
 
 pub fn vortex_alp::ALPRDArray::replace_left_parts_patches(&mut self, patches: core::option::Option<vortex_array::patches::Patches>)
 

encodings/alp/src/alp/array.rs

@@ -76,14 +76,14 @@ impl VTable for ALP {
         array.dtype.hash(state);
         array.encoded().array_hash(state, precision);
         array.exponents.hash(state);
-        array.patches.array_hash(state, precision);
+        array.patches().array_hash(state, precision);
     }
 
     fn array_eq(array: &ALPArray, other: &ALPArray, precision: Precision) -> bool {
         array.dtype == other.dtype
             && array.encoded().array_eq(other.encoded(), precision)
             && array.exponents == other.exponents
-            && array.patches.array_eq(&other.patches, precision)
+            && array.patches().array_eq(&other.patches(), precision)
     }
 
     fn nbuffers(_array: &ALPArray) -> usize {
@@ -114,23 +114,12 @@ impl VTable for ALP {
             slots.len()
         );
 
-        // Reconstruct patches from slots + existing metadata
-        array.patches = match (&slots[PATCH_INDICES_SLOT], &slots[PATCH_VALUES_SLOT]) {
-            (Some(indices), Some(values)) => {
-                let old = array
-                    .patches
-                    .as_ref()
-                    .vortex_expect("ALPArray had patch slots but no patches metadata");
-                Some(Patches::new(
-                    old.array_len(),
-                    old.offset(),
-                    indices.clone(),
-                    values.clone(),
-                    slots[PATCH_CHUNK_OFFSETS_SLOT].clone(),
-                )?)
-            }
-            _ => None,
-        };
+        // If patch slots are being cleared, clear the metadata too
+        if slots[PATCH_INDICES_SLOT].is_none() || slots[PATCH_VALUES_SLOT].is_none() {
+            array.patch_offset = None;
+            array.patch_offset_within_chunk = None;
+        }
+
         array.slots = slots;
         Ok(())
     }
@@ -240,7 +229,8 @@ pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = [
 #[derive(Clone, Debug)]
 pub struct ALPArray {
     slots: Vec<Option<ArrayRef>>,
-    patches: Option<Patches>,
+    patch_offset: Option<usize>,
+    patch_offset_within_chunk: Option<usize>,
     dtype: DType,
     exponents: Exponents,
     stats_set: ArrayStats,
@@ -409,12 +399,17 @@ impl ALPArray {
         };
 
         let slots = Self::make_slots(&encoded, &patches);
+        let (patch_offset, patch_offset_within_chunk) = match &patches {
+            Some(p) => (Some(p.offset()), p.offset_within_chunk()),
+            None => (None, None),
+        };
 
         Ok(Self {
             dtype,
             slots,
             exponents,
-            patches,
+            patch_offset,
+            patch_offset_within_chunk,
             stats_set: Default::default(),
         })
     }
@@ -430,12 +425,17 @@ impl ALPArray {
         dtype: DType,
     ) -> Self {
         let slots = Self::make_slots(&encoded, &patches);
+        let (patch_offset, patch_offset_within_chunk) = match &patches {
+            Some(p) => (Some(p.offset()), p.offset_within_chunk()),
+            None => (None, None),
+        };
 
         Self {
             dtype,
             slots,
             exponents,
-            patches,
+            patch_offset,
+            patch_offset_within_chunk,
             stats_set: Default::default(),
         }
     }
@@ -472,17 +472,38 @@ impl ALPArray {
         self.exponents
     }
 
-    pub fn patches(&self) -> Option<&Patches> {
-        self.patches.as_ref()
+    pub fn patches(&self) -> Option<Patches> {
+        match (
+            &self.slots[PATCH_INDICES_SLOT],
+            &self.slots[PATCH_VALUES_SLOT],
+        ) {
+            (Some(indices), Some(values)) => {
+                let patch_offset = self
+                    .patch_offset
+                    .vortex_expect("has patch slots but no patch_offset");
+                Some(unsafe {
+                    Patches::new_unchecked(
+                        self.encoded().len(),
+                        patch_offset,
+                        indices.clone(),
+                        values.clone(),
+                        self.slots[PATCH_CHUNK_OFFSETS_SLOT].clone(),
+                        self.patch_offset_within_chunk,
+                    )
+                })
+            }
+            _ => None,
+        }
     }
 
     /// Consumes the array and returns its parts.
     #[inline]
     pub fn into_parts(mut self) -> (ArrayRef, Exponents, Option<Patches>, DType) {
+        let patches = self.patches();
         let encoded = self.slots[ENCODED_SLOT]
             .take()
             .vortex_expect("ALPArray encoded slot");
-        (encoded, self.exponents, self.patches, self.dtype)
+        (encoded, self.exponents, patches, self.dtype)
     }
 }
 
@@ -506,7 +527,6 @@ mod tests {
     use vortex_array::arrays::PrimitiveArray;
     use vortex_array::assert_arrays_eq;
     use vortex_array::session::ArraySession;
-    use vortex_array::vtable::ValidityHelper;
     use vortex_session::VortexSession;
 
     use super::*;

encodings/alp/src/alp/compress.rs

@@ -8,7 +8,6 @@ use vortex_array::arrays::PrimitiveArray;
 use vortex_array::dtype::PType;
 use vortex_array::patches::Patches;
 use vortex_array::validity::Validity;
-use vortex_array::vtable::ValidityHelper;
 use vortex_buffer::Buffer;
 use vortex_buffer::BufferMut;
 use vortex_error::VortexResult;
@@ -73,7 +72,7 @@ where
     let (exponents, encoded, exceptional_positions, exceptional_values, mut chunk_offsets) =
         T::encode(values_slice, exponents);
 
-    let encoded_array = PrimitiveArray::new(encoded, values.validity().clone()).into_array();
+    let encoded_array = PrimitiveArray::new(encoded, values.validity()).into_array();
 
     let validity = values.validity_mask()?;
     // exceptional_positions may contain exceptions at invalid positions (which contain garbage

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

@@ -29,7 +29,7 @@ impl CastReduce for ALP {
                 .patches()
                 .map(|p| {
                     if p.values().dtype() == dtype {
-                        Ok(p.clone())
+                        Ok(p)
                     } else {
                         Patches::new(
                             p.array_len(),

encodings/alp/src/alp/decompress.rs

@@ -10,7 +10,6 @@ use vortex_array::arrays::primitive::patch_chunk;
 use vortex_array::dtype::DType;
 use vortex_array::match_each_unsigned_integer_ptype;
 use vortex_array::patches::Patches;
-use vortex_array::vtable::ValidityHelper;
 use vortex_buffer::BufferMut;
 use vortex_error::VortexResult;
 
@@ -101,7 +100,7 @@ fn decompress_chunked_core(
     patches: &Patches,
     dtype: DType,
 ) -> PrimitiveArray {
-    let validity = encoded.validity().clone();
+    let validity = encoded.validity();
     let ptype = dtype.as_ptype();
     let array_len = encoded.len();
     let offset_within_chunk = patches.offset_within_chunk().unwrap_or(0);
@@ -151,7 +150,7 @@ fn decompress_unchunked_core(
     dtype: DType,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<PrimitiveArray> {
-    let validity = encoded.validity().clone();
+    let validity = encoded.validity();
     let ptype = dtype.as_ptype();
 
     let decoded = match_each_alp_float_ptype!(ptype, |T| {

encodings/alp/src/alp_rd/array.rs

@@ -543,8 +543,8 @@ impl ALPRDArray {
     }
 
     /// Patches of left-most bits.
-    pub fn left_parts_patches(&self) -> Option<&Patches> {
-        self.left_parts_patches.as_ref()
+    pub fn left_parts_patches(&self) -> Option<Patches> {
+        self.left_parts_patches.clone()
     }
 
     /// The dictionary that maps the codes in `left_parts` into bit patterns.

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

@@ -34,7 +34,7 @@ impl CastReduce for ALPRD {
                     array.left_parts_dictionary().clone(),
                     array.right_parts().clone(),
                     array.right_bit_width(),
-                    array.left_parts_patches().cloned(),
+                    array.left_parts_patches(),
                 )?
                 .into_array(),
             ));

encodings/alp/src/alp_rd/compute/mask.rs

@@ -26,7 +26,7 @@ impl MaskReduce for ALPRD {
                 array.left_parts_dictionary().clone(),
                 array.right_parts().clone(),
                 array.right_bit_width(),
-                array.left_parts_patches().cloned(),
+                array.left_parts_patches(),
             )?
             .into_array(),
         ))

encodings/alp/src/alp_rd/mod.rs

@@ -30,7 +30,6 @@ use vortex_array::arrays::PrimitiveArray;
 use vortex_array::dtype::DType;
 use vortex_array::dtype::NativePType;
 use vortex_array::match_each_integer_ptype;
-use vortex_array::vtable::ValidityHelper;
 use vortex_buffer::Buffer;
 use vortex_buffer::BufferMut;
 use vortex_error::VortexExpect;
@@ -229,7 +228,7 @@ impl RDEncoder {
         }
 
         // Bit-pack down the encoded left-parts array that have been dictionary encoded.
-        let primitive_left = PrimitiveArray::new(left_parts, array.validity().clone());
+        let primitive_left = PrimitiveArray::new(left_parts, array.validity());
         // SAFETY: by construction, all values in left_parts can be packed to left_bit_width.
         let packed_left = unsafe {
             bitpack_encode_unchecked(primitive_left, left_bit_width as _)

encodings/bytebool/public-api.lock

@@ -136,4 +136,4 @@ pub fn vortex_bytebool::ByteBoolArray::into_array(self) -> vortex_array::array::
 
 impl vortex_array::vtable::validity::ValidityHelper for vortex_bytebool::ByteBoolArray
 
-pub fn vortex_bytebool::ByteBoolArray::validity(&self) -> &vortex_array::validity::Validity
+pub fn vortex_bytebool::ByteBoolArray::validity(&self) -> vortex_array::validity::Validity