feat: iterative execution for SparseArray

- Add #[array_slots] macro for SparseSlots
- Add SparseParts struct + OwnedExt trait for into_parts(self)
- require_child! indices => Primitive, values => AnyCanonical
- execute_sparse takes SparseParts with resolved patches
- Inner functions (varbin, lists, fixed_size_list) take &Patches + &SparseArray for validity

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: joseph-isaacs

encodings/sparse/src/canonical.rs

@@ -57,66 +57,90 @@ use vortex_error::vortex_panic;
 use crate::ConstantArray;
 use crate::Sparse;
 use crate::SparseArray;
+use crate::SparseParts;
+
+/// Build a temporary [`SparseArray`] from resolved patches (for computing validity).
+fn sparse_array_for_validity(patches: &Patches, fill_value: &Scalar, len: usize) -> SparseArray {
+    // Re-wrap resolved patches (offset=0) into a SparseArray so we can call .validity().
+    Sparse::try_new(
+        patches.indices().clone(),
+        patches.values().clone(),
+        len,
+        fill_value.clone(),
+    )
+    .vortex_expect("rebuilding SparseArray for validity")
+}
+
 pub(super) fn execute_sparse(
-    array: &SparseArray,
+    parts: SparseParts,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<ArrayRef> {
-    if array.patches().num_patches() == 0 {
-        return Ok(ConstantArray::new(array.fill_scalar().clone(), array.len()).into_array());
+    let SparseParts { patches, fill_value, dtype, len } = parts;
+
+    if patches.num_patches() == 0 {
+        return Ok(ConstantArray::new(fill_value, len).into_array());
     }
 
-    Ok(match array.dtype() {
+    // Patches are already resolved (offset subtracted) by SparseParts::resolve_patches().
+    Ok(match &dtype {
         DType::Null => {
-            assert!(array.fill_scalar().is_null());
-            NullArray::new(array.len()).into_array()
+            assert!(fill_value.is_null());
+            NullArray::new(len).into_array()
         }
         DType::Bool(..) => {
-            let resolved_patches = array.resolved_patches()?;
-            execute_sparse_bools(&resolved_patches, array.fill_scalar(), ctx)?
+            execute_sparse_bools(&patches, &fill_value, ctx)?
         }
         DType::Primitive(ptype, ..) => {
-            let resolved_patches = array.resolved_patches()?;
             match_each_native_ptype!(ptype, |P| {
-                execute_sparse_primitives::<P>(&resolved_patches, array.fill_scalar(), ctx)?
+                execute_sparse_primitives::<P>(&patches, &fill_value, ctx)?
             })
         }
         DType::Struct(struct_fields, ..) => execute_sparse_struct(
             struct_fields,
-            array.fill_scalar().as_struct(),
-            array.dtype(),
-            array.patches(),
-            array.len(),
+            fill_value.as_struct(),
+            &dtype,
+            &patches,
+            len,
             ctx,
         )?,
         DType::Decimal(decimal_dtype, nullability) => {
             let canonical_decimal_value_type =
                 DecimalType::smallest_decimal_value_type(decimal_dtype);
-            let fill_value = array.fill_scalar().as_decimal();
+            let fill_decimal = fill_value.as_decimal();
             match_each_decimal_value_type!(canonical_decimal_value_type, |D| {
                 execute_sparse_decimal::<D>(
                     *decimal_dtype,
                     *nullability,
-                    fill_value,
-                    array.patches(),
-                    array.len(),
+                    fill_decimal,
+                    &patches,
+                    len,
                     ctx,
                 )?
             })
         }
         dtype @ DType::Utf8(..) => {
-            let fill_value = array.fill_scalar().as_utf8().value().cloned();
-            let fill_value = fill_value.map(BufferString::into_inner);
-            execute_varbin(array, dtype.clone(), fill_value, ctx)?
+            let fill = fill_value.as_utf8().value().cloned();
+            let fill = fill.map(BufferString::into_inner);
+            let validity_arr = sparse_array_for_validity(&patches, &fill_value, len);
+            execute_varbin(&patches, &validity_arr, dtype.clone(), fill, ctx)?
         }
         dtype @ DType::Binary(..) => {
-            let fill_value = array.fill_scalar().as_binary().value().cloned();
-            execute_varbin(array, dtype.clone(), fill_value, ctx)?
+            let fill = fill_value.as_binary().value().cloned();
+            let validity_arr = sparse_array_for_validity(&patches, &fill_value, len);
+            execute_varbin(&patches, &validity_arr, dtype.clone(), fill, ctx)?
         }
         DType::List(values_dtype, nullability) => {
-            execute_sparse_lists(array, Arc::clone(values_dtype), *nullability, ctx)?
+            let validity_arr = sparse_array_for_validity(&patches, &fill_value, len);
+            execute_sparse_lists(
+                &patches, &validity_arr, &fill_value,
+                Arc::clone(values_dtype), *nullability, ctx,
+            )?
         }
         DType::FixedSizeList(.., nullability) => {
-            execute_sparse_fixed_size_list(array, *nullability, ctx)?
+            let validity_arr = sparse_array_for_validity(&patches, &fill_value, len);
+            execute_sparse_fixed_size_list(
+                &patches, &validity_arr, &fill_value, *nullability, ctx,
+            )?
         }
         DType::Extension(_ext_dtype) => todo!(),
         DType::Variant(_) => vortex_bail!("Sparse canonicalization does not support Variant"),
@@ -128,39 +152,39 @@ pub(super) fn execute_sparse(
     reason = "complexity is from nested match_smallest_offset_type macro"
 )]
 fn execute_sparse_lists(
-    array: &SparseArray,
+    resolved: &Patches,
+    validity_array: &SparseArray,
+    fill_value: &Scalar,
     values_dtype: Arc<DType>,
     nullability: Nullability,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<ArrayRef> {
-    let resolved_patches = array.resolved_patches()?;
-
-    let indices = resolved_patches
+    let indices = resolved
         .indices()
         .clone()
         .execute::<PrimitiveArray>(ctx)?;
-    let values = resolved_patches
+    let values = resolved
         .values()
         .clone()
         .execute::<ListViewArray>(ctx)?;
-    let fill_value = array.fill_scalar().as_list();
+    let fill_list = fill_value.as_list();
 
-    let n_filled = array.len() - resolved_patches.num_patches();
-    let total_canonical_values = values.elements().len() + fill_value.len() * n_filled;
+    let len = validity_array.len();
+    let n_filled = len - resolved.num_patches();
+    let total_canonical_values = values.elements().len() + fill_list.len() * n_filled;
 
-    let validity = {
-        let arr = array.as_array();
-        Validity::from_mask(arr.validity()?.execute_mask(arr.len(), ctx)?, nullability)
-    };
+    let arr = validity_array.as_ref();
+    let validity =
+        Validity::from_mask(arr.validity()?.execute_mask(arr.len(), ctx)?, nullability);
 
     Ok(match_each_integer_ptype!(indices.ptype(), |I| {
         match_smallest_offset_type!(total_canonical_values, |O| {
             execute_sparse_lists_inner::<I, O>(
                 indices.as_slice(),
                 values,
-                fill_value,
+                fill_list,
                 values_dtype,
-                array.len(),
+                len,
                 total_canonical_values,
                 validity,
                 ctx,
@@ -224,32 +248,33 @@ fn execute_sparse_lists_inner<I: IntegerPType, O: IntegerPType>(
 
 /// Canonicalize a sparse [`FixedSizeListArray`] by expanding it into a dense representation.
 fn execute_sparse_fixed_size_list(
-    array: &SparseArray,
+    resolved: &Patches,
+    validity_array: &SparseArray,
+    fill_value: &Scalar,
     nullability: Nullability,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<ArrayRef> {
-    let resolved_patches = array.resolved_patches()?;
-    let indices = resolved_patches
+    let indices = resolved
         .indices()
         .clone()
         .execute::<PrimitiveArray>(ctx)?;
-    let values = resolved_patches
+    let values = resolved
         .values()
         .clone()
         .execute::<FixedSizeListArray>(ctx)?;
-    let fill_value = array.fill_scalar().as_list();
+    let fill_list = fill_value.as_list();
+    let len = validity_array.len();
 
-    let validity = {
-        let arr = array.as_array();
-        Validity::from_mask(arr.validity()?.execute_mask(arr.len(), ctx)?, nullability)
-    };
+    let arr = validity_array.as_ref();
+    let validity =
+        Validity::from_mask(arr.validity()?.execute_mask(arr.len(), ctx)?, nullability);
 
     Ok(match_each_integer_ptype!(indices.ptype(), |I| {
         execute_sparse_fixed_size_list_inner::<I>(
             indices.as_slice(),
             values,
-            fill_value,
-            array.len(),
+            fill_list,
+            len,
             validity,
             ctx,
         )
@@ -496,22 +521,20 @@ fn execute_sparse_decimal<D: NativeDecimalType>(
 }
 
 fn execute_varbin(
-    array: &SparseArray,
+    resolved: &Patches,
+    validity_array: &SparseArray,
     dtype: DType,
     fill_value: Option<ByteBuffer>,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<ArrayRef> {
-    let patches = array.resolved_patches()?;
-    let indices = patches.indices().clone().execute::<PrimitiveArray>(ctx)?;
-    let values = patches.values().clone().execute::<VarBinViewArray>(ctx)?;
-    let validity = {
-        let arr = array.as_array();
-        Validity::from_mask(
-            arr.validity()?.execute_mask(arr.len(), ctx)?,
-            dtype.nullability(),
-        )
-    };
-    let len = array.len();
+    let indices = resolved.indices().clone().execute::<PrimitiveArray>(ctx)?;
+    let values = resolved.values().clone().execute::<VarBinViewArray>(ctx)?;
+    let arr = validity_array.as_ref();
+    let validity = Validity::from_mask(
+        arr.validity()?.execute_mask(arr.len(), ctx)?,
+        dtype.nullability(),
+    );
+    let len = arr.len();
 
     Ok(match_each_integer_ptype!(indices.ptype(), |I| {
         let indices = indices.to_buffer::<I>();

encodings/sparse/src/lib.rs

@@ -14,13 +14,15 @@ use vortex_array::ArrayEq;
 use vortex_array::ArrayHash;
 use vortex_array::ArrayId;
 use vortex_array::ArrayParts;
+use vortex_array::AnyCanonical;
 use vortex_array::ArrayRef;
 use vortex_array::ArrayView;
 use vortex_array::Canonical;
 use vortex_array::ExecutionCtx;
 use vortex_array::ExecutionResult;
 use vortex_array::IntoArray;
 use vortex_array::Precision;
+use vortex_array::arrays::Primitive;
 use vortex_array::arrays::BoolArray;
 use vortex_array::arrays::ConstantArray;
 use vortex_array::arrays::PrimitiveArray;
@@ -64,6 +66,69 @@ mod slice;
 /// A [`Sparse`]-encoded Vortex array.
 pub type SparseArray = Array<Sparse>;
 
+#[vortex_array::array_slots(Sparse)]
+pub struct SparseSlots {
+    pub patch_indices: ArrayRef,
+    pub patch_values: ArrayRef,
+    pub patch_chunk_offsets: Option<ArrayRef>,
+}
+
+/// Concrete parts of a [`SparseArray`] after iterative execution.
+pub(crate) struct SparseParts {
+    pub patches: Patches,
+    pub fill_value: Scalar,
+    pub dtype: DType,
+    pub len: usize,
+}
+
+impl SparseParts {
+    /// Resolve patches by subtracting the offset from indices.
+    pub fn resolve_patches(mut self) -> VortexResult<Self> {
+        if self.patches.offset() != 0 {
+            let offset_scalar =
+                Scalar::from(self.patches.offset()).cast(self.patches.indices().dtype())?;
+            let indices = self.patches.indices().binary(
+                ConstantArray::new(offset_scalar, self.patches.indices().len()).into_array(),
+                Operator::Sub,
+            )?;
+            self.patches = Patches::new(
+                self.patches.array_len(),
+                0,
+                indices,
+                self.patches.values().clone(),
+                None,
+            )?;
+        }
+        Ok(self)
+    }
+}
+
+pub(crate) trait SparseOwnedExt {
+    fn into_parts(self) -> VortexResult<SparseParts>;
+}
+
+impl SparseOwnedExt for Array<Sparse> {
+    fn into_parts(self) -> VortexResult<SparseParts> {
+        let patches = Patches::new(
+            self.len(),
+            self.patches().offset(),
+            self.as_ref().slots()[SparseSlots::PATCH_INDICES]
+                .clone()
+                .vortex_expect("indices"),
+            self.as_ref().slots()[SparseSlots::PATCH_VALUES]
+                .clone()
+                .vortex_expect("values"),
+            self.as_ref().slots()[SparseSlots::PATCH_CHUNK_OFFSETS].clone(),
+        )?;
+        Ok(SparseParts {
+            patches,
+            fill_value: self.fill_scalar().clone(),
+            dtype: self.dtype().clone(),
+            len: self.len(),
+        })
+    }
+}
+
 #[derive(Clone, prost::Message)]
 #[repr(C)]
 pub struct SparseMetadata {
@@ -186,7 +251,7 @@ impl VTable for Sparse {
     }
 
     fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
-        SLOT_NAMES[idx].to_string()
+        SparseSlots::NAMES[idx].to_string()
     }
 
     fn reduce_parent(
@@ -207,13 +272,25 @@ impl VTable for Sparse {
     }
 
     fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
-        execute_sparse(&array, ctx).map(ExecutionResult::done)
+        // Require children to be executed through the scheduler,
+        // enabling cross-step optimization via reduce_parent rules.
+        let array = vortex_array::require_child!(
+            array, array.patch_indices(), SparseSlots::PATCH_INDICES => Primitive
+        );
+        let array = vortex_array::require_child!(
+            array, array.patch_values(), SparseSlots::PATCH_VALUES => AnyCanonical
+        );
+        vortex_array::require_opt_child!(
+            array,
+            array.patch_chunk_offsets(),
+            SparseSlots::PATCH_CHUNK_OFFSETS => Primitive
+        );
+
+        let parts = array.into_parts()?.resolve_patches()?;
+        execute_sparse(parts, ctx).map(ExecutionResult::done)
     }
 }
 
-pub(crate) const NUM_SLOTS: usize = 3;
-pub(crate) const SLOT_NAMES: [&str; NUM_SLOTS] =
-    ["patch_indices", "patch_values", "patch_chunk_offsets"];
 
 #[derive(Clone, Debug)]
 pub struct SparseData {