Merge remote-tracking branch 'origin/develop' into ji/remove-param-names-public-lock

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

# Conflicts:
#	vortex-array/public-api.lock

Author: joseph-isaacs

docs/developer-guide/internals/execution.md

@@ -158,13 +158,14 @@ mod tests {
     /// Filter unwrap one layer at a time so RunEnd's FilterKernel can fire.
     #[test]
     fn filter_sliced_run_end_preserves_encoding() -> VortexResult<()> {
+        let mut ctx = LEGACY_SESSION.create_execution_ctx();
+
         // 4 runs of 32 each = 128 rows. Large enough that FilterKernel takes
         // the run-preserving path (true_count >= 25).
         let values: Vec<i32> = [10, 20, 30, 40]
             .iter()
             .flat_map(|&v| std::iter::repeat_n(v, 32))
             .collect();
-        let mut ctx = LEGACY_SESSION.create_execution_ctx();
         let arr = RunEnd::encode(PrimitiveArray::from_iter(values).into_array(), &mut ctx)?;
 
         // Slice off the first 16 rows. Slice(RunEnd), 112 rows, 4 runs.

encodings/runend/src/compute/filter.rs

@@ -23,6 +23,7 @@ use crate::ArrayHash;
 use crate::ArrayView;
 use crate::Canonical;
 use crate::ExecutionCtx;
+use crate::ExecutionResult;
 use crate::IntoArray;
 use crate::LEGACY_SESSION;
 use crate::VTable;
@@ -93,6 +94,12 @@ impl ArrayRef {
         &self.0
     }
 
+    /// Returns a reference to the inner Arc.
+    #[inline(always)]
+    pub(crate) fn inner_mut(&mut self) -> &mut Arc<dyn DynArray> {
+        &mut self.0
+    }
+
     /// Consumes the array reference, returning the owned backing allocation.
     #[inline(always)]
     pub(crate) fn into_inner(self) -> Arc<dyn DynArray> {
@@ -431,10 +438,13 @@ impl ArrayRef {
         self.with_slots(slots)
     }
 
-    /// Take a slot for executor-owned physical rewrites. This has the result that the array may
-    /// either be taken or cloned from the parent.
+    /// Take a slot for executor-owned physical rewrites.
     ///
-    /// The array can be put back with [`put_slot_unchecked`].
+    /// On return the produced parent has the taken slot set to `None`
+    /// callers must put the slot back (typically via [`put_slot_unchecked`]) before the parent is
+    /// returned from the execution loop.
+    ///
+    /// When the `Arc` was shared this allocates a fresh parent.
     ///
     /// # Safety
     /// The caller must put back a slot with the same logical dtype and length before exposing the
@@ -443,18 +453,27 @@ impl ArrayRef {
         mut self,
         slot_idx: usize,
     ) -> VortexResult<(ArrayRef, ArrayRef)> {
-        let child = if let Some(inner) = Arc::get_mut(&mut self.0) {
-            // # Safety: ensured by the caller.
-            unsafe { inner.slots_mut()[slot_idx].take() }
-                .vortex_expect("take_slot_unchecked cannot take an absent slot")
-        } else {
-            self.slots()[slot_idx]
-                .as_ref()
-                .vortex_expect("take_slot_unchecked cannot take an absent slot")
-                .clone()
-        };
+        if let Some(inner) = Arc::get_mut(&mut self.0) {
+            // SAFETY: ensured by the caller.
+            let child = unsafe { inner.slots_mut()[slot_idx].take() }
+                .vortex_expect("take_slot_unchecked cannot take an absent slot");
+            return Ok((self, child));
+        }
+
+        // Arc is shared: clone the child out and build a fresh parent with slot_idx = None,
+        // bypassing encoding-level validation so the absent slot does not panic `V::validate`.
+        let child = self.slots()[slot_idx]
+            .as_ref()
+            .vortex_expect("take_slot_unchecked cannot take an absent slot")
+            .clone();
+
+        let mut new_slots = self.slots().to_vec();
+        new_slots[slot_idx] = None;
 
-        Ok((self, child))
+        // SAFETY: ensured by the caller — the None slot is either put back or driven to completion
+        // via the builder path before the parent escapes the executor.
+        let new_parent = unsafe { self.0.with_slots_unchecked(&self, new_slots) };
+        Ok((new_parent, child))
     }
 
     /// Puts an array into `slot_idx` by either, cloning the inner array if the Arc is not exclusive
@@ -532,12 +551,28 @@ impl ArrayRef {
         self.0.reduce_parent(self, parent, child_idx)
     }
 
-    pub(crate) fn execute_encoding(
+    pub(crate) fn execute_encoding(self, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
+        let inner = Arc::as_ptr(&self.0);
+        // Safety the Arc outline the DynArray function call
+        unsafe { (&*inner).execute(self, ctx) }
+    }
+
+    /// Execute a single encoding step without applying `Done`-result postconditions.
+    ///
+    /// This is for the iterative executor only. It may operate on suspended executor-private
+    /// arrays whose slots temporarily contain `None`, so the executor itself must interpret
+    /// `Done`, enforce any `len`/`dtype` invariants, and transfer statistics.
+    pub(crate) fn execute_encoding_unchecked(
         self,
         ctx: &mut ExecutionCtx,
-    ) -> VortexResult<crate::ExecutionResult> {
-        let inner = Arc::clone(&self.0);
-        inner.execute(self, ctx)
+    ) -> VortexResult<ExecutionResult> {
+        let inner = Arc::as_ptr(&self.0);
+        // Safety the Arc outline the DynArray function call
+        let inner = unsafe { &*inner };
+        // SAFETY: `inner` points at the allocation owned by `self.0`. `self` stays alive for the
+        // duration of the call, so the pointee remains valid. Avoiding an extra `Arc` clone here
+        // preserves uniqueness so execute-time metadata cursors can use `Arc::get_mut`.
+        unsafe { inner.execute_unchecked(self, ctx) }
     }
 
     pub fn execute_parent(

vortex-array/public-api.lock

@@ -56,6 +56,9 @@ pub(crate) trait DynArray: 'static + private::Sealed + Send + Sync + Debug {
     /// Returns the array as a reference to a generic [`Any`] trait object.
     fn as_any(&self) -> &dyn Any;
 
+    /// Returns the array as a mutable reference to a generic [`Any`] trait object.
+    fn as_any_mut(&mut self) -> &mut dyn Any;
+
     /// Converts an owned array allocation into an owned [`Any`] allocation for downcasting.
     fn into_any_arc(self: std::sync::Arc<Self>) -> std::sync::Arc<dyn Any + Send + Sync>;
 
@@ -143,6 +146,24 @@ pub(crate) trait DynArray: 'static + private::Sealed + Send + Sync + Debug {
     /// Returns a new array with the given slots.
     fn with_slots(&self, this: ArrayRef, slots: Vec<Option<ArrayRef>>) -> VortexResult<ArrayRef>;
 
+    /// Returns a new array with the given slots, bypassing encoding-level validation.
+    ///
+    /// Used by the executor to temporarily carry an array that has had one of its child slots
+    /// taken out (leaving `None`) without panicking `V::validate`. The caller must ensure the
+    /// missing slot is filled back in (via `put_slot_unchecked`) or driven to completion by the
+    /// builder path before the array becomes externally observable.
+    ///
+    /// # Safety
+    ///
+    /// The array returned may have slots whose content does not match the encoding's normal
+    /// invariants. Callers must re-establish those invariants before handing the array to
+    /// anything outside the executor.
+    unsafe fn with_slots_unchecked(
+        &self,
+        this: &ArrayRef,
+        slots: Vec<Option<ArrayRef>>,
+    ) -> ArrayRef;
+
     /// Attempt to reduce the array to a simpler representation.
     fn reduce(&self, this: &ArrayRef) -> VortexResult<Option<ArrayRef>>;
 
@@ -155,8 +176,30 @@ pub(crate) trait DynArray: 'static + private::Sealed + Send + Sync + Debug {
     ) -> VortexResult<Option<ArrayRef>>;
 
     /// Execute the array by taking a single encoding-specific execution step.
+    ///
+    /// This is the checked entry point. If the encoding reports
+    /// [`ExecutionStep::Done`](crate::ExecutionStep::Done), implementations must validate that the
+    /// returned array preserves this array's logical `len` and `dtype`, and must transfer this
+    /// array's statistics to the returned array.
     fn execute(&self, this: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult>;
 
+    /// Execute the array by taking a single encoding-specific execution step without applying
+    /// `Done`-result postconditions.
+    ///
+    /// This exists for the iterative executor, which may call into `execute` on suspended
+    /// executor-private arrays whose slots temporarily contain `None`. In that mode the executor
+    /// itself is responsible for deciding when a `Done` result represents a real logical array,
+    /// enforcing any `len`/`dtype` invariants, and transferring statistics.
+    ///
+    /// # Safety
+    /// The `array` returned should have it's `DType` and len checked
+    /// (optionally it should have its stats propagated from `this`).
+    unsafe fn execute_unchecked(
+        &self,
+        this: ArrayRef,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<ExecutionResult>;
+
     /// Attempt to execute the parent of this array.
     fn execute_parent(
         &self,
@@ -203,6 +246,10 @@ impl<V: VTable> DynArray for ArrayInner<V> {
         self
     }
 
+    fn as_any_mut(&mut self) -> &mut dyn Any {
+        self
+    }
+
     fn into_any_arc(self: std::sync::Arc<Self>) -> std::sync::Arc<dyn Any + Send + Sync> {
         self
     }
@@ -387,6 +434,26 @@ impl<V: VTable> DynArray for ArrayInner<V> {
         .into_array())
     }
 
+    unsafe fn with_slots_unchecked(
+        &self,
+        this: &ArrayRef,
+        slots: Vec<Option<ArrayRef>>,
+    ) -> ArrayRef {
+        // SAFETY: we intentionally skip `V::validate` here. Caller guarantees that the resulting
+        // array is either repaired or not externally observed.
+        let inner = unsafe {
+            ArrayInner::<V>::from_data_unchecked(
+                self.vtable.clone(),
+                this.dtype().clone(),
+                self.len,
+                self.data.clone(),
+                slots,
+                self.stats.clone(),
+            )
+        };
+        ArrayRef::from_inner(std::sync::Arc::new(inner))
+    }
+
     fn reduce(&self, this: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
         let view = unsafe { ArrayView::new_unchecked(this, &self.data) };
         let Some(reduced) = V::reduce(view)? else {
@@ -437,12 +504,8 @@ impl<V: VTable> DynArray for ArrayInner<V> {
     fn execute(&self, this: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
         let len = this.len();
         let dtype = this.dtype().clone();
-        let stats = this.statistics().to_owned();
-
-        let typed = Array::<V>::try_from_array_ref(this)
-            .map_err(|_| vortex_err!("Failed to downcast array for execute"))
-            .vortex_expect("Failed to downcast array for execute");
-        let result = V::execute(typed, ctx)?;
+        let stats = this.statistics().to_array_stats();
+        let result = unsafe { self.execute_unchecked(this, ctx)? };
 
         if matches!(result.step(), ExecutionStep::Done) {
             if cfg!(debug_assertions) {
@@ -458,12 +521,26 @@ impl<V: VTable> DynArray for ArrayInner<V> {
                 );
             }
 
-            result.array().statistics().set_iter(stats.into_iter());
+            result
+                .array()
+                .statistics()
+                .set_iter(crate::stats::StatsSet::from(stats).into_iter());
         }
 
         Ok(result)
     }
 
+    unsafe fn execute_unchecked(
+        &self,
+        this: ArrayRef,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<ExecutionResult> {
+        let typed = Array::<V>::try_from_array_ref(this)
+            .map_err(|_| vortex_err!("Failed to downcast array for execute"))
+            .vortex_expect("Failed to downcast array for execute");
+        V::execute(typed, ctx)
+    }
+
     fn execute_parent(
         &self,
         this: &ArrayRef,

vortex-array/src/array/erased.rs

@@ -280,6 +280,14 @@ impl<V: VTable> Array<V> {
         &self.downcast_inner().data
     }
 
+    /// Try to fetch a mut ref to the inner ArrayData.
+    pub fn data_mut(&mut self) -> Option<&mut V::ArrayData> {
+        let m = self.inner.inner_mut();
+        let inner = Arc::get_mut(m)?;
+        let array_inner = inner.as_any_mut().downcast_mut::<ArrayInner<V>>();
+        Some(&mut array_inner?.data)
+    }
+
     /// Returns the full typed array construction parts if this handle owns the allocation.
     pub fn try_into_parts(self) -> Result<ArrayParts<V>, Self> {
         let Self { inner, _phantom } = self;

vortex-array/src/array/mod.rs

@@ -37,6 +37,8 @@ pub(super) const CHUNKS_OFFSET: usize = 1;
 #[derive(Clone, Debug)]
 pub struct ChunkedData {
     pub(super) chunk_offsets: Vec<usize>,
+    /// This is used to find the next child to execute when in executing into a builder.
+    pub(super) next_builder_slot: usize,
 }
 
 impl Display for ChunkedData {
@@ -114,6 +116,13 @@ pub trait ChunkedArrayExt: TypedArrayRef<Chunked> {
 impl<T: TypedArrayRef<Chunked>> ChunkedArrayExt for T {}
 
 impl ChunkedData {
+    pub(super) fn new(chunk_offsets: Vec<usize>) -> Self {
+        Self {
+            chunk_offsets,
+            next_builder_slot: CHUNKS_OFFSET,
+        }
+    }
+
     pub(super) fn compute_chunk_offsets(chunks: &[ArrayRef]) -> Vec<usize> {
         let mut chunk_offsets = Vec::with_capacity(chunks.len() + 1);
         chunk_offsets.push(0);
@@ -159,24 +168,31 @@ impl ChunkedData {
 }
 
 impl Array<Chunked> {
+    pub(super) fn with_next_builder_slot(mut self, next_builder_slot: usize) -> Self {
+        if let Some(data) = self.data_mut() {
+            data.next_builder_slot = next_builder_slot;
+            return self;
+        }
+        // This is the slow path that will be hit at most once per execution since the second one
+        // *MUST* have execlusive access due to this copy.
+        let stats = self.statistics().to_owned();
+        let mut data = self.data().clone();
+        data.next_builder_slot = next_builder_slot;
+        // SAFETY: we only modified next_builder_slot which doesn't affect array invariants.
+        unsafe {
+            Array::from_parts_unchecked(
+                ArrayParts::new(Chunked, self.dtype().clone(), self.len(), data)
+                    .with_slots(self.slots().to_vec()),
+            )
+        }
+        .with_stats_set(stats)
+    }
+
     /// Constructs a new `ChunkedArray`.
     pub fn try_new(chunks: Vec<ArrayRef>, dtype: DType) -> VortexResult<Self> {
         ChunkedData::validate(&chunks, &dtype)?;
-        let len = chunks.iter().map(|chunk| chunk.len()).sum();
-        let chunk_offsets = ChunkedData::compute_chunk_offsets(&chunks);
-        Ok(unsafe {
-            Array::from_parts_unchecked(
-                ArrayParts::new(
-                    Chunked,
-                    dtype,
-                    len,
-                    ChunkedData {
-                        chunk_offsets: chunk_offsets.clone(),
-                    },
-                )
-                .with_slots(ChunkedData::make_slots(&chunk_offsets, &chunks)),
-            )
-        })
+        // SAFETY just validated on previous line.
+        Ok(unsafe { Self::new_unchecked(chunks, dtype) })
     }
 
     pub fn rechunk(&self, target_bytesize: u64, target_rowsize: usize) -> VortexResult<Self> {
@@ -238,15 +254,8 @@ impl Array<Chunked> {
         let chunk_offsets = ChunkedData::compute_chunk_offsets(&chunks);
         unsafe {
             Array::from_parts_unchecked(
-                ArrayParts::new(
-                    Chunked,
-                    dtype,
-                    len,
-                    ChunkedData {
-                        chunk_offsets: chunk_offsets.clone(),
-                    },
-                )
-                .with_slots(ChunkedData::make_slots(&chunk_offsets, &chunks)),
+                ArrayParts::new(Chunked, dtype, len, ChunkedData::new(chunk_offsets.clone()))
+                    .with_slots(ChunkedData::make_slots(&chunk_offsets, &chunks)),
             )
         }
     }