feat: add an arrays slots macro

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>

Author: joseph-isaacs

Cargo.lock

@@ -1,6 +1,7 @@
 [workspace]
 members = [
     # Core crates in dependency order
+    "vortex-array-macros",
     "vortex-error",
     "vortex-buffer",
     "vortex-mask",
@@ -238,6 +239,7 @@ tempfile = "3"
 termtree = { version = "1.0" }
 test-with = "0.15"
 thiserror = "2.0.3"
+vortex-array-macros = { path = "./vortex-array-macros" }
 tokio = { version = "1.48" }
 tokio-stream = "0.1.17"
 tokio-util = "0.7.17"

Cargo.toml

@@ -296,6 +296,58 @@ pub fn vortex_alp::ALPRDMetadata::clear(&mut self)
 
 pub fn vortex_alp::ALPRDMetadata::encoded_len(&self) -> usize
 
+pub struct vortex_alp::ALPSlots
+
+pub vortex_alp::ALPSlots::encoded: vortex_array::array::erased::ArrayRef
+
+pub vortex_alp::ALPSlots::patch_chunk_offsets: core::option::Option<vortex_array::array::erased::ArrayRef>
+
+pub vortex_alp::ALPSlots::patch_indices: core::option::Option<vortex_array::array::erased::ArrayRef>
+
+pub vortex_alp::ALPSlots::patch_values: core::option::Option<vortex_array::array::erased::ArrayRef>
+
+impl vortex_alp::ALPSlots
+
+pub const vortex_alp::ALPSlots::COUNT: usize
+
+pub const vortex_alp::ALPSlots::ENCODED: usize
+
+pub const vortex_alp::ALPSlots::NAMES: [&'static str; 4]
+
+pub const vortex_alp::ALPSlots::PATCH_CHUNK_OFFSETS: usize
+
+pub const vortex_alp::ALPSlots::PATCH_INDICES: usize
+
+pub const vortex_alp::ALPSlots::PATCH_VALUES: usize
+
+pub fn vortex_alp::ALPSlots::from_slots(slots: alloc::vec::Vec<core::option::Option<vortex_array::array::erased::ArrayRef>>) -> Self
+
+pub fn vortex_alp::ALPSlots::into_slots(self) -> alloc::vec::Vec<core::option::Option<vortex_array::array::erased::ArrayRef>>
+
+pub struct vortex_alp::ALPSlotsView<'a>
+
+pub vortex_alp::ALPSlotsView::encoded: &'a vortex_array::array::erased::ArrayRef
+
+pub vortex_alp::ALPSlotsView::patch_chunk_offsets: core::option::Option<&'a vortex_array::array::erased::ArrayRef>
+
+pub vortex_alp::ALPSlotsView::patch_indices: core::option::Option<&'a vortex_array::array::erased::ArrayRef>
+
+pub vortex_alp::ALPSlotsView::patch_values: core::option::Option<&'a vortex_array::array::erased::ArrayRef>
+
+impl<'a> vortex_alp::ALPSlotsView<'a>
+
+pub fn vortex_alp::ALPSlotsView<'a>::from_slots(slots: &'a [core::option::Option<vortex_array::array::erased::ArrayRef>]) -> Self
+
+impl<'a> core::clone::Clone for vortex_alp::ALPSlotsView<'a>
+
+pub fn vortex_alp::ALPSlotsView<'a>::clone(&self) -> vortex_alp::ALPSlotsView<'a>
+
+impl<'a> core::fmt::Debug for vortex_alp::ALPSlotsView<'a>
+
+pub fn vortex_alp::ALPSlotsView<'a>::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl<'a> core::marker::Copy for vortex_alp::ALPSlotsView<'a>
+
 pub struct vortex_alp::Exponents
 
 pub vortex_alp::Exponents::e: u8
@@ -338,18 +390,14 @@ pub fn vortex_alp::RDEncoder::from_parts(right_bit_width: u8, codes: alloc::vec:
 
 pub fn vortex_alp::RDEncoder::new<T>(sample: &[T]) -> Self where T: vortex_alp::ALPRDFloat + vortex_array::dtype::ptype::NativePType, <T as vortex_alp::ALPRDFloat>::UINT: vortex_array::dtype::ptype::NativePType
 
-pub trait vortex_alp::ALPArrayExt: vortex_array::array::typed::TypedArrayRef<vortex_alp::ALP>
-
-pub fn vortex_alp::ALPArrayExt::encoded(&self) -> &vortex_array::array::erased::ArrayRef
+pub trait vortex_alp::ALPArrayExt: vortex_alp::ALPArraySlotsExt
 
 pub fn vortex_alp::ALPArrayExt::exponents(&self) -> vortex_alp::Exponents
 
 pub fn vortex_alp::ALPArrayExt::patches(&self) -> core::option::Option<vortex_array::patches::Patches>
 
 impl<T: vortex_array::array::typed::TypedArrayRef<vortex_alp::ALP>> vortex_alp::ALPArrayExt for T
 
-pub fn T::encoded(&self) -> &vortex_array::array::erased::ArrayRef
-
 pub fn T::exponents(&self) -> vortex_alp::Exponents
 
 pub fn T::patches(&self) -> core::option::Option<vortex_array::patches::Patches>
@@ -362,6 +410,30 @@ impl vortex_alp::ALPArrayOwnedExt for vortex_array::array::typed::Array<vortex_a
 
 pub fn vortex_array::array::typed::Array<vortex_alp::ALP>::into_parts(self) -> (vortex_array::array::erased::ArrayRef, vortex_alp::Exponents, core::option::Option<vortex_array::patches::Patches>)
 
+pub trait vortex_alp::ALPArraySlotsExt: vortex_array::array::typed::TypedArrayRef<vortex_alp::ALP>
+
+pub fn vortex_alp::ALPArraySlotsExt::encoded(&self) -> &vortex_array::array::erased::ArrayRef
+
+pub fn vortex_alp::ALPArraySlotsExt::patch_chunk_offsets(&self) -> core::option::Option<&vortex_array::array::erased::ArrayRef>
+
+pub fn vortex_alp::ALPArraySlotsExt::patch_indices(&self) -> core::option::Option<&vortex_array::array::erased::ArrayRef>
+
+pub fn vortex_alp::ALPArraySlotsExt::patch_values(&self) -> core::option::Option<&vortex_array::array::erased::ArrayRef>
+
+pub fn vortex_alp::ALPArraySlotsExt::slots_view(&self) -> vortex_alp::ALPSlotsView<'_>
+
+impl<T: vortex_array::array::typed::TypedArrayRef<vortex_alp::ALP>> vortex_alp::ALPArraySlotsExt for T
+
+pub fn T::encoded(&self) -> &vortex_array::array::erased::ArrayRef
+
+pub fn T::patch_chunk_offsets(&self) -> core::option::Option<&vortex_array::array::erased::ArrayRef>
+
+pub fn T::patch_indices(&self) -> core::option::Option<&vortex_array::array::erased::ArrayRef>
+
+pub fn T::patch_values(&self) -> core::option::Option<&vortex_array::array::erased::ArrayRef>
+
+pub fn T::slots_view(&self) -> vortex_alp::ALPSlotsView<'_>
+
 pub trait vortex_alp::ALPFloat: vortex_alp::alp::private::Sealed + num_traits::float::Float + core::fmt::Display + vortex_array::dtype::ptype::NativePType
 
 pub type vortex_alp::ALPFloat::ALPInt: num_traits::int::PrimInt + core::fmt::Display + num_traits::cast::ToPrimitive + core::marker::Copy + vortex_array::dtype::ptype::NativePType + core::convert::Into<vortex_array::scalar::typed_view::primitive::pvalue::PValue>

encodings/alp/public-api.lock

@@ -85,9 +85,7 @@ impl VTable for ALP {
             dtype,
             len,
             data.exponents,
-            slots[ENCODED_SLOT]
-                .as_ref()
-                .vortex_expect("ALPArray encoded slot"),
+            ALPSlotsView::from_slots(slots).encoded,
             patches_from_slots(
                 slots,
                 data.patch_offset,
@@ -170,16 +168,16 @@ impl VTable for ALP {
     }
 
     fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
-        SLOT_NAMES[idx].to_string()
+        ALPSlots::NAMES[idx].to_string()
     }
 
     fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
-        let array = require_child!(array, array.encoded(), ENCODED_SLOT => Primitive);
+        let array = require_child!(array, array.encoded(), ALPSlots::ENCODED => Primitive);
         require_patches!(
             array,
-            PATCH_INDICES_SLOT,
-            PATCH_VALUES_SLOT,
-            PATCH_CHUNK_OFFSETS_SLOT
+            ALPSlots::PATCH_INDICES,
+            ALPSlots::PATCH_VALUES,
+            ALPSlots::PATCH_CHUNK_OFFSETS
         );
 
         Ok(ExecutionResult::done(
@@ -205,21 +203,17 @@ impl VTable for ALP {
     }
 }
 
-/// The ALP-encoded values array.
-pub(super) const ENCODED_SLOT: usize = 0;
-/// The indices of exception values that could not be ALP-encoded.
-pub(super) const PATCH_INDICES_SLOT: usize = 1;
-/// The exception values that could not be ALP-encoded.
-pub(super) const PATCH_VALUES_SLOT: usize = 2;
-/// Chunk offsets for the patch indices/values.
-pub(super) const PATCH_CHUNK_OFFSETS_SLOT: usize = 3;
-pub(super) const NUM_SLOTS: usize = 4;
-pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = [
-    "encoded",
-    "patch_indices",
-    "patch_values",
-    "patch_chunk_offsets",
-];
+#[vortex_array::array_slots(ALP)]
+pub struct ALPSlots {
+    /// The ALP-encoded values array.
+    pub encoded: ArrayRef,
+    /// The indices of exception values that could not be ALP-encoded.
+    pub patch_indices: Option<ArrayRef>,
+    /// The exception values that could not be ALP-encoded.
+    pub patch_values: Option<ArrayRef>,
+    /// Chunk offsets for the patch indices/values.
+    pub patch_chunk_offsets: Option<ArrayRef>,
+}
 
 #[derive(Clone, Debug)]
 pub struct ALPData {
@@ -439,13 +433,7 @@ impl ALPData {
     }
 }
 
-pub trait ALPArrayExt: TypedArrayRef<ALP> {
-    fn encoded(&self) -> &ArrayRef {
-        self.as_ref().slots()[ENCODED_SLOT]
-            .as_ref()
-            .vortex_expect("ALPArray encoded slot")
-    }
-
+pub trait ALPArrayExt: ALPArraySlotsExt {
     fn exponents(&self) -> Exponents {
         self.exponents
     }
@@ -466,7 +454,8 @@ fn patches_from_slots(
     patch_offset_within_chunk: Option<usize>,
     len: usize,
 ) -> Option<Patches> {
-    match (&slots[PATCH_INDICES_SLOT], &slots[PATCH_VALUES_SLOT]) {
+    let view = ALPSlotsView::from_slots(slots);
+    match (view.patch_indices, view.patch_values) {
         (Some(indices), Some(values)) => {
             let patch_offset = patch_offset.vortex_expect("has patch slots but no patch_offset");
             Some(unsafe {
@@ -475,7 +464,7 @@ fn patches_from_slots(
                     patch_offset,
                     indices.clone(),
                     values.clone(),
-                    slots[PATCH_CHUNK_OFFSETS_SLOT].clone(),
+                    view.patch_chunk_offsets.cloned(),
                     patch_offset_within_chunk,
                 )
             })

encodings/alp/src/alp/array.rs

@@ -135,6 +135,7 @@ mod tests {
 
     use super::*;
     use crate::alp::array::ALPArrayExt;
+    use crate::alp::array::ALPArraySlotsExt;
     use crate::decompress_into_array;
 
     #[test]

encodings/alp/src/alp/compress.rs

@@ -20,6 +20,7 @@ use vortex_error::VortexResult;
 use crate::ALP;
 use crate::ALPFloat;
 use crate::alp::array::ALPArrayExt;
+use crate::alp::array::ALPArraySlotsExt;
 use crate::match_each_alp_float_ptype;
 
 impl BetweenReduce for ALP {
@@ -103,6 +104,7 @@ mod tests {
 
     use crate::ALPArray;
     use crate::alp::array::ALPArrayExt;
+    use crate::alp::array::ALPArraySlotsExt;
     use crate::alp::compute::between::between_impl;
     use crate::alp_encode;
 

encodings/alp/src/alp/compute/between.rs

@@ -11,6 +11,7 @@ use vortex_array::scalar_fn::fns::cast::CastReduce;
 use vortex_error::VortexResult;
 
 use crate::ALPArrayExt;
+use crate::ALPArraySlotsExt;
 use crate::alp::ALP;
 
 impl CastReduce for ALP {
@@ -72,6 +73,7 @@ mod tests {
     use vortex_error::VortexResult;
 
     use crate::alp::array::ALPArrayExt;
+    use crate::alp::array::ALPArraySlotsExt;
     use crate::alp_encode;
 
     #[test]

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

@@ -20,6 +20,7 @@ use vortex_error::vortex_err;
 
 use crate::ALP;
 use crate::ALPArrayExt;
+use crate::ALPArraySlotsExt;
 use crate::ALPFloat;
 use crate::match_each_alp_float_ptype;
 

encodings/alp/src/alp/compute/compare.rs

@@ -11,6 +11,7 @@ use vortex_mask::Mask;
 
 use crate::ALP;
 use crate::ALPArrayExt;
+use crate::ALPArraySlotsExt;
 
 impl FilterKernel for ALP {
     fn filter(

encodings/alp/src/alp/compute/filter.rs

@@ -13,6 +13,7 @@ use vortex_error::VortexResult;
 
 use crate::ALP;
 use crate::ALPArrayExt;
+use crate::ALPArraySlotsExt;
 
 impl MaskReduce for ALP {
     fn mask(array: ArrayView<'_, Self>, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
@@ -66,6 +67,7 @@ mod test {
     use vortex_buffer::buffer;
 
     use crate::alp::array::ALPArrayExt;
+    use crate::alp::array::ALPArraySlotsExt;
     use crate::alp_encode;
 
     #[rstest]