feat: masked array iterative execution clean up (#7690)

Use array_slots macro, lazy validity and iterative execution

see #7674

---------

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

Author: joseph-isaacs

.github/runs-on.yml

@@ -1 +1,6 @@
 _extends: .github-private
+images:
+  ubuntu24-full-arm64-pre:
+    platform: "linux"
+    arch: "arm64"
+    ami: "ami-08f97c6362847dc5d"

vortex-array/public-api.lock

@@ -3408,23 +3408,51 @@ impl vortex_array::ArrayHash for vortex_array::arrays::masked::MaskedData
 
 pub fn vortex_array::arrays::masked::MaskedData::array_hash<H: core::hash::Hasher>(&self, _state: &mut H, _precision: vortex_array::Precision)
 
-pub trait vortex_array::arrays::masked::MaskedArrayExt: vortex_array::TypedArrayRef<vortex_array::arrays::Masked>
+pub struct vortex_array::arrays::masked::MaskedSlots
 
-pub fn vortex_array::arrays::masked::MaskedArrayExt::child(&self) -> &vortex_array::ArrayRef
+pub vortex_array::arrays::masked::MaskedSlots::child: vortex_array::ArrayRef
 
-pub fn vortex_array::arrays::masked::MaskedArrayExt::masked_validity(&self) -> vortex_array::validity::Validity
+pub vortex_array::arrays::masked::MaskedSlots::validity: core::option::Option<vortex_array::ArrayRef>
+
+impl vortex_array::arrays::masked::MaskedSlots
+
+pub const vortex_array::arrays::masked::MaskedSlots::CHILD: usize
+
+pub const vortex_array::arrays::masked::MaskedSlots::COUNT: usize
+
+pub const vortex_array::arrays::masked::MaskedSlots::NAMES: [&'static str; 2]
 
-pub fn vortex_array::arrays::masked::MaskedArrayExt::validity_child(&self) -> core::option::Option<&vortex_array::ArrayRef>
+pub const vortex_array::arrays::masked::MaskedSlots::VALIDITY: usize
+
+pub fn vortex_array::arrays::masked::MaskedSlots::from_slots(slots: alloc::vec::Vec<core::option::Option<vortex_array::ArrayRef>>) -> Self
+
+pub fn vortex_array::arrays::masked::MaskedSlots::into_slots(self) -> alloc::vec::Vec<core::option::Option<vortex_array::ArrayRef>>
+
+pub trait vortex_array::arrays::masked::MaskedArrayExt: vortex_array::TypedArrayRef<vortex_array::arrays::Masked> + vortex_array::arrays::masked::MaskedArraySlotsExt
+
+pub fn vortex_array::arrays::masked::MaskedArrayExt::masked_validity(&self) -> vortex_array::validity::Validity
 
 impl<T: vortex_array::TypedArrayRef<vortex_array::arrays::Masked>> vortex_array::arrays::masked::MaskedArrayExt for T
 
+pub fn T::masked_validity(&self) -> vortex_array::validity::Validity
+
+pub trait vortex_array::arrays::masked::MaskedArraySlotsExt: vortex_array::TypedArrayRef<vortex_array::arrays::Masked>
+
+pub fn vortex_array::arrays::masked::MaskedArraySlotsExt::child(&self) -> &vortex_array::ArrayRef
+
+pub fn vortex_array::arrays::masked::MaskedArraySlotsExt::slots_view(&self) -> vortex_array::arrays::masked::array::MaskedSlotsView<'_>
+
+pub fn vortex_array::arrays::masked::MaskedArraySlotsExt::validity(&self) -> core::option::Option<&vortex_array::ArrayRef>
+
+impl<T: vortex_array::TypedArrayRef<vortex_array::arrays::Masked>> vortex_array::arrays::masked::MaskedArraySlotsExt for T
+
 pub fn T::child(&self) -> &vortex_array::ArrayRef
 
-pub fn T::masked_validity(&self) -> vortex_array::validity::Validity
+pub fn T::slots_view(&self) -> vortex_array::arrays::masked::array::MaskedSlotsView<'_>
 
-pub fn T::validity_child(&self) -> core::option::Option<&vortex_array::ArrayRef>
+pub fn T::validity(&self) -> core::option::Option<&vortex_array::ArrayRef>
 
-pub fn vortex_array::arrays::masked::mask_validity_canonical(canonical: vortex_array::Canonical, validity_mask: &vortex_mask::Mask, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::Canonical>
+pub fn vortex_array::arrays::masked::mask_validity_canonical(canonical: vortex_array::Canonical, validity: vortex_array::validity::Validity, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::Canonical>
 
 pub type vortex_array::arrays::masked::MaskedArray = vortex_array::Array<vortex_array::arrays::Masked>
 

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

@@ -4,7 +4,6 @@
 use std::fmt::Display;
 use std::fmt::Formatter;
 
-use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 
@@ -16,15 +15,17 @@ use crate::array::ArrayParts;
 use crate::array::TypedArrayRef;
 use crate::array::child_to_validity;
 use crate::array::validity_to_child;
+use crate::array_slots;
 use crate::arrays::Masked;
 use crate::validity::Validity;
 
-/// The underlying child array being masked.
-pub(super) const CHILD_SLOT: usize = 0;
-/// The validity bitmap defining which elements are non-null.
-pub(super) const VALIDITY_SLOT: usize = 1;
-pub(super) const NUM_SLOTS: usize = 2;
-pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["child", "validity"];
+#[array_slots(Masked)]
+pub struct MaskedSlots {
+    /// The underlying child array being masked.
+    pub child: ArrayRef,
+    /// The validity bitmap defining which elements are non-null.
+    pub validity: Option<ArrayRef>,
+}
 
 #[derive(Clone, Debug)]
 pub struct MaskedData;
@@ -35,22 +36,10 @@ impl Display for MaskedData {
     }
 }
 
-pub trait MaskedArrayExt: TypedArrayRef<Masked> {
-    fn child(&self) -> &ArrayRef {
-        self.as_ref().slots()[CHILD_SLOT]
-            .as_ref()
-            .vortex_expect("validated masked child slot")
-    }
-
-    fn validity_child(&self) -> Option<&ArrayRef> {
-        self.as_ref().slots()[VALIDITY_SLOT].as_ref()
-    }
-
+pub trait MaskedArrayExt: TypedArrayRef<Masked> + MaskedArraySlotsExt {
     fn masked_validity(&self) -> Validity {
-        child_to_validity(
-            &self.as_ref().slots()[VALIDITY_SLOT],
-            self.as_ref().dtype().nullability(),
-        )
+        let x = &self.as_ref().slots()[MaskedSlots::VALIDITY];
+        child_to_validity(x, self.as_ref().dtype().nullability())
     }
 }
 impl<T: TypedArrayRef<Masked>> MaskedArrayExt for T {}

vortex-array/src/arrays/masked/compute/filter.rs

@@ -10,7 +10,7 @@ use crate::array::ArrayView;
 use crate::arrays::Masked;
 use crate::arrays::MaskedArray;
 use crate::arrays::filter::FilterReduce;
-use crate::arrays::masked::MaskedArrayExt;
+use crate::arrays::masked::MaskedArraySlotsExt;
 
 impl FilterReduce for Masked {
     fn filter(array: ArrayView<'_, Masked>, mask: &Mask) -> VortexResult<Option<ArrayRef>> {

vortex-array/src/arrays/masked/compute/mask.rs

@@ -6,7 +6,7 @@ use vortex_error::VortexResult;
 use crate::ArrayRef;
 use crate::array::ArrayView;
 use crate::arrays::Masked;
-use crate::arrays::masked::MaskedArrayExt;
+use crate::arrays::masked::MaskedArraySlotsExt;
 use crate::arrays::scalar_fn::ScalarFnFactoryExt;
 use crate::scalar_fn::EmptyOptions;
 use crate::scalar_fn::fns::mask::Mask as MaskExpr;

vortex-array/src/arrays/masked/compute/slice.rs

@@ -10,7 +10,7 @@ use crate::IntoArray;
 use crate::array::ArrayView;
 use crate::arrays::Masked;
 use crate::arrays::MaskedArray;
-use crate::arrays::masked::MaskedArrayExt;
+use crate::arrays::masked::MaskedArraySlotsExt;
 use crate::arrays::slice::SliceReduce;
 
 impl SliceReduce for Masked {

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

@@ -9,7 +9,7 @@ use crate::array::ArrayView;
 use crate::arrays::Masked;
 use crate::arrays::MaskedArray;
 use crate::arrays::dict::TakeReduce;
-use crate::arrays::masked::MaskedArrayExt;
+use crate::arrays::masked::MaskedArraySlotsExt;
 use crate::builtins::ArrayBuiltins;
 use crate::scalar::Scalar;
 use crate::validity::Validity;