feat: add an arrays slots macro

.github/workflows/ci.yml

@@ -81,7 +81,7 @@ jobs:
           || 'ubuntu-latest' }}
     timeout-minutes: 40
     env:
-      RUST_LOG: "info,uv=debug"
+      RUST_LOG: "info,maturin=warn,uv=debug"
       MATURIN_PEP517_ARGS: "--profile ci"
     steps:
       - uses: runs-on/action@v2

Cargo.toml

@@ -1,6 +1,7 @@
 [workspace]
 members = [
     # Core crates in dependency order
+    "vortex-array-macros",
     "vortex-error",
     "vortex-buffer",
     "vortex-mask",
@@ -241,6 +242,7 @@ thiserror = "2.0.3"
 tokio = { version = "1.48" }
 tokio-stream = "0.1.17"
 tokio-util = "0.7.17"
+vortex-array-macros = { version = "0.1.0", path = "./vortex-array-macros" }
 # Pull these into non-public crates to support DF 58
 tpchgen = { version = "2.0.2", git = "https://github.com/clflushopt/tpchgen-rs.git", rev = "438e9c2dbc25b2fff82c0efc08b3f13b5707874f" }
 tpchgen-arrow = { version = "2.0.2", git = "https://github.com/clflushopt/tpchgen-rs.git", rev = "438e9c2dbc25b2fff82c0efc08b3f13b5707874f" }

encodings/alp/public-api.lock

@@ -296,6 +296,60 @@ 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
+
+pub fn vortex_alp::ALPSlotsView<'a>::to_owned(&self) -> vortex_alp::ALPSlots
+
+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 +392,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 +412,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/src/alp/array.rs

@@ -20,6 +20,7 @@ use vortex_array::ExecutionResult;
 use vortex_array::IntoArray;
 use vortex_array::Precision;
 use vortex_array::TypedArrayRef;
+use vortex_array::array_slots;
 use vortex_array::arrays::Primitive;
 use vortex_array::buffer::BufferHandle;
 use vortex_array::dtype::DType;
@@ -81,15 +82,14 @@ impl VTable for ALP {
         len: usize,
         slots: &[Option<ArrayRef>],
     ) -> VortexResult<()> {
+        let slots = ALPSlotsView::from_slots(slots);
         validate_parts(
             dtype,
             len,
             data.exponents,
-            slots[ENCODED_SLOT]
-                .as_ref()
-                .vortex_expect("ALPArray encoded slot"),
+            slots.encoded,
             patches_from_slots(
-                slots,
+                &slots,
                 data.patch_offset,
                 data.patch_offset_within_chunk,
                 len,
@@ -170,16 +170,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 +205,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",
-];
+#[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,20 +435,14 @@ 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
     }
 
     fn patches(&self) -> Option<Patches> {
         patches_from_slots(
-            self.as_ref().slots(),
+            &self.slots_view(),
             self.patch_offset,
             self.patch_offset_within_chunk,
             self.as_ref().len(),
@@ -461,12 +451,12 @@ pub trait ALPArrayExt: TypedArrayRef<ALP> {
 }
 
 fn patches_from_slots(
-    slots: &[Option<ArrayRef>],
+    slots: &ALPSlotsView,
     patch_offset: Option<usize>,
     patch_offset_within_chunk: Option<usize>,
     len: usize,
 ) -> Option<Patches> {
-    match (&slots[PATCH_INDICES_SLOT], &slots[PATCH_VALUES_SLOT]) {
+    match (slots.patch_indices, slots.patch_values) {
         (Some(indices), Some(values)) => {
             let patch_offset = patch_offset.vortex_expect("has patch slots but no patch_offset");
             Some(unsafe {
@@ -475,7 +465,7 @@ fn patches_from_slots(
                     patch_offset,
                     indices.clone(),
                     values.clone(),
-                    slots[PATCH_CHUNK_OFFSETS_SLOT].clone(),
+                    slots.patch_chunk_offsets.cloned(),
                     patch_offset_within_chunk,
                 )
             })

encodings/alp/src/alp/compress.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/compute/between.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 {

encodings/alp/src/alp/compute/cast.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 {

encodings/alp/src/alp/compute/compare.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/filter.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/mask.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>> {

encodings/alp/src/alp/compute/slice.rs

@@ -12,6 +12,7 @@ use vortex_error::VortexResult;
 
 use crate::ALP;
 use crate::ALPArrayExt;
+use crate::ALPArraySlotsExt;
 
 impl SliceKernel for ALP {
     fn slice(

encodings/alp/src/alp/compute/take.rs

@@ -10,6 +10,7 @@ use vortex_error::VortexResult;
 
 use crate::ALP;
 use crate::ALPArrayExt;
+use crate::ALPArraySlotsExt;
 
 impl TakeExecute for ALP {
     fn take(

encodings/alp/src/alp/ops.rs

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

encodings/fastlanes/benches/compute_between.rs

@@ -7,6 +7,7 @@ use rand::RngExt;
 use rand::rngs::StdRng;
 use vortex_alp::ALP;
 use vortex_alp::ALPArrayExt;
+use vortex_alp::ALPArraySlotsExt;
 use vortex_alp::alp_encode;
 use vortex_array::ArrayRef;
 use vortex_array::IntoArray;