wip: chunked iter exe

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

Author: joseph-isaacs

vortex-array/src/arrays/chunked/vtable/canonical.rs

@@ -10,8 +10,10 @@ use crate::Canonical;
 use crate::DynArray;
 use crate::ExecutionCtx;
 use crate::IntoArray;
+use crate::arrays::ListView;
 use crate::arrays::ListViewArray;
 use crate::arrays::PrimitiveArray;
+use crate::arrays::Struct;
 use crate::arrays::StructArray;
 use crate::arrays::chunked::vtable::ChunkedArray;
 use crate::arrays::listview::ListViewRebuildMode;
@@ -41,7 +43,6 @@ pub(super) fn _canonicalize(
                 &owned_chunks,
                 Validity::copy_from_array(&array.clone().into_array())?,
                 struct_dtype,
-                ctx,
             )?;
             Canonical::Struct(struct_array)
         }
@@ -62,24 +63,22 @@ pub(super) fn _canonicalize(
 /// Packs many [`StructArray`]s to instead be a single [`StructArray`], where the [`DynArray`] for each
 /// field is a [`ChunkedArray`].
 ///
-/// The caller guarantees there are at least 2 chunks.
+/// The caller guarantees there are at least 2 chunks, and that all chunks are already
+/// canonicalized to [`StructArray`] by iterative execution.
 fn pack_struct_chunks(
     chunks: &[ArrayRef],
     validity: Validity,
     struct_dtype: &StructFields,
-    ctx: &mut ExecutionCtx,
 ) -> VortexResult<StructArray> {
     let len = chunks.iter().map(|chunk| chunk.len()).sum();
     let mut field_arrays = Vec::new();
 
-    let executed_chunks: Vec<StructArray> = chunks
-        .iter()
-        .map(|c| c.clone().execute::<StructArray>(ctx))
-        .collect::<VortexResult<_>>()?;
-
     for (field_idx, field_dtype) in struct_dtype.fields().enumerate() {
         let mut field_chunks = Vec::with_capacity(chunks.len());
-        for struct_array in &executed_chunks {
+        for chunk in chunks {
+            let struct_array = chunk
+                .as_opt::<Struct>()
+                .vortex_expect("struct chunk pre-canonicalized by iterative execution");
             let field = struct_array.unmasked_field(field_idx).to_array();
             field_chunks.push(field);
         }
@@ -99,7 +98,8 @@ fn pack_struct_chunks(
 ///
 /// We use the existing arrays (chunks) to form a chunked array of `elements` (the child array).
 ///
-/// The caller guarantees there are at least 2 chunks.
+/// The caller guarantees there are at least 2 chunks, and that all chunks are already
+/// canonicalized to [`ListViewArray`] by iterative execution.
 fn swizzle_list_chunks(
     chunks: &[ArrayRef],
     validity: Validity,
@@ -131,7 +131,9 @@ fn swizzle_list_chunks(
     let mut sizes = BufferMut::<u64>::with_capacity(len);
 
     for chunk in chunks {
-        let chunk_array = chunk.clone().execute::<ListViewArray>(ctx)?;
+        let chunk_array = chunk
+            .as_opt::<ListView>()
+            .vortex_expect("list chunk pre-canonicalized by iterative execution");
         // By rebuilding as zero-copy to `List` and trimming all elements (to prevent gaps), we make
         // the final output `ListView` also zero-copyable to `List`.
         let chunk_array = chunk_array.rebuild(ListViewRebuildMode::MakeExact)?;

vortex-array/src/arrays/chunked/vtable/mod.rs

@@ -12,6 +12,7 @@ use vortex_error::vortex_err;
 use vortex_error::vortex_panic;
 use vortex_session::VortexSession;
 
+use crate::AnyCanonical;
 use crate::ArrayRef;
 use crate::Canonical;
 use crate::EmptyMetadata;
@@ -202,6 +203,16 @@ impl VTable for Chunked {
     }
 
     fn execute(array: Arc<Array<Self>>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
+        // Iteratively request execution of each chunk until it reaches canonical form.
+        // This gives the scheduler visibility into child execution and enables
+        // cross-step optimizations between chunk decoding steps.
+        for i in 0..array.nchunks() {
+            if !array.chunk(i).is::<AnyCanonical>() {
+                return Ok(ExecutionResult::execute_slot::<AnyCanonical>(array, i + 1));
+            }
+        }
+
+        // All chunks are now canonical — combine them.
         Ok(ExecutionResult::done(
             _canonicalize(&array, ctx)?.into_array(),
         ))