diff --git a/vortex-array/benches/chunked_dict_builder.rs b/vortex-array/benches/chunked_dict_builder.rs index b42514fcf19..bb47720705d 100644 --- a/vortex-array/benches/chunked_dict_builder.rs +++ b/vortex-array/benches/chunked_dict_builder.rs @@ -59,6 +59,6 @@ fn chunked_dict_primitive_into_canonical( let chunk = gen_dict_primitive_chunks::(len, unique_values, chunk_count); bencher - .with_inputs(|| chunk.clone()) - .bench_values(|chunk| chunk.execute::(&mut SESSION.create_execution_ctx())) + .with_inputs(|| (gen_dict_primitive_chunks::(len, unique_values, chunk_count), SESSION.create_execution_ctx())) + .bench_values(|(chunk, mut ctx)| chunk.execute::(&mut ctx)) } diff --git a/vortex-array/public-api.lock b/vortex-array/public-api.lock index 84b3fb9513d..16f7122e63d 100644 --- a/vortex-array/public-api.lock +++ b/vortex-array/public-api.lock @@ -3182,7 +3182,7 @@ pub fn vortex_array::arrays::patched::Patched::child_name(array: vortex_array::A pub fn vortex_array::arrays::patched::Patched::deserialize(&self, dtype: &vortex_array::dtype::DType, len: usize, metadata: &[u8], _buffers: &[vortex_array::buffer::BufferHandle], children: &dyn vortex_array::serde::ArrayChildren, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult> -pub fn vortex_array::arrays::patched::Patched::execute(array: vortex_array::Array, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult +pub fn vortex_array::arrays::patched::Patched::execute(array: vortex_array::Array, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult pub fn vortex_array::arrays::patched::Patched::execute_parent(array: vortex_array::ArrayView<'_, Self>, parent: &vortex_array::ArrayRef, child_idx: usize, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult> @@ -3922,7 +3922,7 @@ pub fn vortex_array::arrays::slice::Slice::child_name(array: vortex_array::Array pub fn vortex_array::arrays::slice::Slice::deserialize(&self, _dtype: &vortex_array::dtype::DType, _len: usize, _metadata: &[u8], _buffers: &[vortex_array::buffer::BufferHandle], _children: &dyn vortex_array::serde::ArrayChildren, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult> -pub fn vortex_array::arrays::slice::Slice::execute(array: vortex_array::Array, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult +pub fn vortex_array::arrays::slice::Slice::execute(array: vortex_array::Array, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult pub fn vortex_array::arrays::slice::Slice::execute_parent(array: vortex_array::ArrayView<'_, Self>, parent: &vortex_array::ArrayRef, child_idx: usize, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult> @@ -6022,7 +6022,7 @@ pub fn vortex_array::arrays::patched::Patched::child_name(array: vortex_array::A pub fn vortex_array::arrays::patched::Patched::deserialize(&self, dtype: &vortex_array::dtype::DType, len: usize, metadata: &[u8], _buffers: &[vortex_array::buffer::BufferHandle], children: &dyn vortex_array::serde::ArrayChildren, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult> -pub fn vortex_array::arrays::patched::Patched::execute(array: vortex_array::Array, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult +pub fn vortex_array::arrays::patched::Patched::execute(array: vortex_array::Array, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult pub fn vortex_array::arrays::patched::Patched::execute_parent(array: vortex_array::ArrayView<'_, Self>, parent: &vortex_array::ArrayRef, child_idx: usize, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult> @@ -6320,7 +6320,7 @@ pub fn vortex_array::arrays::slice::Slice::child_name(array: vortex_array::Array pub fn vortex_array::arrays::slice::Slice::deserialize(&self, _dtype: &vortex_array::dtype::DType, _len: usize, _metadata: &[u8], _buffers: &[vortex_array::buffer::BufferHandle], _children: &dyn vortex_array::serde::ArrayChildren, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult> -pub fn vortex_array::arrays::slice::Slice::execute(array: vortex_array::Array, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult +pub fn vortex_array::arrays::slice::Slice::execute(array: vortex_array::Array, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult pub fn vortex_array::arrays::slice::Slice::execute_parent(array: vortex_array::ArrayView<'_, Self>, parent: &vortex_array::ArrayRef, child_idx: usize, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult> @@ -19676,7 +19676,7 @@ pub fn vortex_array::arrays::patched::Patched::child_name(array: vortex_array::A pub fn vortex_array::arrays::patched::Patched::deserialize(&self, dtype: &vortex_array::dtype::DType, len: usize, metadata: &[u8], _buffers: &[vortex_array::buffer::BufferHandle], children: &dyn vortex_array::serde::ArrayChildren, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult> -pub fn vortex_array::arrays::patched::Patched::execute(array: vortex_array::Array, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult +pub fn vortex_array::arrays::patched::Patched::execute(array: vortex_array::Array, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult pub fn vortex_array::arrays::patched::Patched::execute_parent(array: vortex_array::ArrayView<'_, Self>, parent: &vortex_array::ArrayRef, child_idx: usize, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult> @@ -19760,7 +19760,7 @@ pub fn vortex_array::arrays::slice::Slice::child_name(array: vortex_array::Array pub fn vortex_array::arrays::slice::Slice::deserialize(&self, _dtype: &vortex_array::dtype::DType, _len: usize, _metadata: &[u8], _buffers: &[vortex_array::buffer::BufferHandle], _children: &dyn vortex_array::serde::ArrayChildren, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult> -pub fn vortex_array::arrays::slice::Slice::execute(array: vortex_array::Array, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult +pub fn vortex_array::arrays::slice::Slice::execute(array: vortex_array::Array, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult pub fn vortex_array::arrays::slice::Slice::execute_parent(array: vortex_array::ArrayView<'_, Self>, parent: &vortex_array::ArrayRef, child_idx: usize, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult> @@ -20692,7 +20692,7 @@ pub fn vortex_array::arrays::patched::Patched::child_name(array: vortex_array::A pub fn vortex_array::arrays::patched::Patched::deserialize(&self, dtype: &vortex_array::dtype::DType, len: usize, metadata: &[u8], _buffers: &[vortex_array::buffer::BufferHandle], children: &dyn vortex_array::serde::ArrayChildren, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult> -pub fn vortex_array::arrays::patched::Patched::execute(array: vortex_array::Array, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult +pub fn vortex_array::arrays::patched::Patched::execute(array: vortex_array::Array, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult pub fn vortex_array::arrays::patched::Patched::execute_parent(array: vortex_array::ArrayView<'_, Self>, parent: &vortex_array::ArrayRef, child_idx: usize, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult> @@ -20776,7 +20776,7 @@ pub fn vortex_array::arrays::slice::Slice::child_name(array: vortex_array::Array pub fn vortex_array::arrays::slice::Slice::deserialize(&self, _dtype: &vortex_array::dtype::DType, _len: usize, _metadata: &[u8], _buffers: &[vortex_array::buffer::BufferHandle], _children: &dyn vortex_array::serde::ArrayChildren, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult> -pub fn vortex_array::arrays::slice::Slice::execute(array: vortex_array::Array, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult +pub fn vortex_array::arrays::slice::Slice::execute(array: vortex_array::Array, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult pub fn vortex_array::arrays::slice::Slice::execute_parent(array: vortex_array::ArrayView<'_, Self>, parent: &vortex_array::ArrayRef, child_idx: usize, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult> @@ -23398,7 +23398,7 @@ pub fn vortex_array::arrays::patched::Patched::child_name(array: vortex_array::A pub fn vortex_array::arrays::patched::Patched::deserialize(&self, dtype: &vortex_array::dtype::DType, len: usize, metadata: &[u8], _buffers: &[vortex_array::buffer::BufferHandle], children: &dyn vortex_array::serde::ArrayChildren, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult> -pub fn vortex_array::arrays::patched::Patched::execute(array: vortex_array::Array, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult +pub fn vortex_array::arrays::patched::Patched::execute(array: vortex_array::Array, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult pub fn vortex_array::arrays::patched::Patched::execute_parent(array: vortex_array::ArrayView<'_, Self>, parent: &vortex_array::ArrayRef, child_idx: usize, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult> @@ -23482,7 +23482,7 @@ pub fn vortex_array::arrays::slice::Slice::child_name(array: vortex_array::Array pub fn vortex_array::arrays::slice::Slice::deserialize(&self, _dtype: &vortex_array::dtype::DType, _len: usize, _metadata: &[u8], _buffers: &[vortex_array::buffer::BufferHandle], _children: &dyn vortex_array::serde::ArrayChildren, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult> -pub fn vortex_array::arrays::slice::Slice::execute(array: vortex_array::Array, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult +pub fn vortex_array::arrays::slice::Slice::execute(array: vortex_array::Array, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult pub fn vortex_array::arrays::slice::Slice::execute_parent(array: vortex_array::ArrayView<'_, Self>, parent: &vortex_array::ArrayRef, child_idx: usize, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult> @@ -24662,7 +24662,7 @@ pub fn vortex_array::arrays::patched::Patched::child_name(array: vortex_array::A pub fn vortex_array::arrays::patched::Patched::deserialize(&self, dtype: &vortex_array::dtype::DType, len: usize, metadata: &[u8], _buffers: &[vortex_array::buffer::BufferHandle], children: &dyn vortex_array::serde::ArrayChildren, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult> -pub fn vortex_array::arrays::patched::Patched::execute(array: vortex_array::Array, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult +pub fn vortex_array::arrays::patched::Patched::execute(array: vortex_array::Array, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult pub fn vortex_array::arrays::patched::Patched::execute_parent(array: vortex_array::ArrayView<'_, Self>, parent: &vortex_array::ArrayRef, child_idx: usize, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult> @@ -24746,7 +24746,7 @@ pub fn vortex_array::arrays::slice::Slice::child_name(array: vortex_array::Array pub fn vortex_array::arrays::slice::Slice::deserialize(&self, _dtype: &vortex_array::dtype::DType, _len: usize, _metadata: &[u8], _buffers: &[vortex_array::buffer::BufferHandle], _children: &dyn vortex_array::serde::ArrayChildren, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult> -pub fn vortex_array::arrays::slice::Slice::execute(array: vortex_array::Array, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult +pub fn vortex_array::arrays::slice::Slice::execute(array: vortex_array::Array, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult pub fn vortex_array::arrays::slice::Slice::execute_parent(array: vortex_array::ArrayView<'_, Self>, parent: &vortex_array::ArrayRef, child_idx: usize, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult> diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index cb5d2e1bb86..a44e60d67b1 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -395,6 +395,29 @@ impl ArrayRef { self.is::() } + /// Returns a new array with the slot at `slot_idx` replaced by `replacement`. + /// + /// This is only valid for physical rewrites: the replacement must have the same logical + /// `DType` and `len` as the existing slot. + /// + /// Takes ownership to allow in-place mutation when the refcount is 1. + pub unsafe fn with_slot_unchecked( + mut self, + slot_idx: usize, + replacement: ArrayRef, + ) -> ArrayRef { + // Fast path: if we have the only reference, mutate in place. + if let Some(inner) = Arc::get_mut(&mut self.0) { + inner.replace_slot(slot_idx, replacement); + return self; + } + + // Slow path: clone the slots and rebuild. + let mut slots = self.slots().to_vec(); + slots[slot_idx] = Some(replacement); + self.with_slots(slots).vortex_expect("cannot fail") + } + /// Returns a new array with the slot at `slot_idx` replaced by `replacement`. /// /// This is only valid for physical rewrites: the replacement must have the same logical diff --git a/vortex-array/src/array/mod.rs b/vortex-array/src/array/mod.rs index f6ed110e172..97abb620994 100644 --- a/vortex-array/src/array/mod.rs +++ b/vortex-array/src/array/mod.rs @@ -141,6 +141,9 @@ pub(crate) trait DynArray: 'static + private::Sealed + Send + Sync + Debug { /// Compares two arrays of the same concrete type for equality. fn dyn_array_eq(&self, other: &ArrayRef, precision: crate::Precision) -> bool; + /// Replace a single slot in-place. Only called when the Arc refcount is 1. + fn replace_slot(&mut self, slot_idx: usize, replacement: ArrayRef); + /// Returns a new array with the given slots. fn with_slots(&self, this: ArrayRef, slots: Vec>) -> VortexResult; @@ -378,6 +381,10 @@ impl DynArray for ArrayInner { }) } + fn replace_slot(&mut self, slot_idx: usize, replacement: ArrayRef) { + self.slots[slot_idx] = Some(replacement); + } + fn with_slots(&self, this: ArrayRef, slots: Vec>) -> VortexResult { let data = self.data.clone(); let stats = this.statistics().to_owned(); diff --git a/vortex-array/src/array/typed.rs b/vortex-array/src/array/typed.rs index 0d8239af589..05ed863be24 100644 --- a/vortex-array/src/array/typed.rs +++ b/vortex-array/src/array/typed.rs @@ -91,7 +91,15 @@ impl ArrayInner { pub fn try_new(new: ArrayParts) -> VortexResult { new.vtable .validate(&new.data, &new.dtype, new.len, &new.slots)?; - Ok(unsafe { + Ok(unsafe { Self::new_unchecked(new) }) + } + + /// Create from [`ArrayParts`] without validation. + /// + /// # Safety + /// Caller must ensure dtype and len match the data. + pub unsafe fn new_unchecked(new: ArrayParts) -> Self { + unsafe { Self::from_data_unchecked( new.vtable, new.dtype, @@ -100,7 +108,7 @@ impl ArrayInner { new.slots, ArrayStats::default(), ) - }) + } } /// Create without validation. @@ -209,16 +217,7 @@ impl Array { /// Caller must ensure the provided parts are logically consistent. #[doc(hidden)] pub unsafe fn from_parts_unchecked(new: ArrayParts) -> Self { - let inner = ArrayRef::from_inner(Arc::new(unsafe { - ArrayInner::::from_data_unchecked( - new.vtable, - new.dtype, - new.len, - new.data, - new.slots, - ArrayStats::default(), - ) - })); + let inner = ArrayRef::from_inner(Arc::new(unsafe { ArrayInner::new_unchecked(new) })); Self { inner, _phantom: PhantomData, diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index 73a9b1c856a..70f3bb4b3cd 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -12,8 +12,10 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; +use crate::arrays::ListView; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; +use crate::arrays::Struct; use crate::arrays::StructArray; use crate::arrays::chunked::ChunkedArrayExt; use crate::arrays::listview::ListViewArrayExt; @@ -45,7 +47,6 @@ pub(super) fn _canonicalize( &owned_chunks, Validity::copy_from_array(array.array())?, struct_dtype, - ctx, )?; Canonical::Struct(struct_array) } @@ -66,24 +67,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 { let len = chunks.iter().map(|chunk| chunk.len()).sum(); let mut field_arrays = Vec::new(); - let executed_chunks: Vec = chunks - .iter() - .map(|c| c.clone().execute::(ctx)) - .collect::>()?; - 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::() + .vortex_expect("struct chunk pre-canonicalized by iterative execution"); let field = struct_array.unmasked_field(field_idx).clone(); field_chunks.push(field); } @@ -103,7 +102,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, @@ -135,7 +135,11 @@ fn swizzle_list_chunks( let mut sizes = BufferMut::::with_capacity(len); for chunk in chunks { - let chunk_array = chunk.clone().execute::(ctx)?; + let chunk_array = chunk + .clone() + .try_downcast::() + .ok() + .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)?; diff --git a/vortex-array/src/arrays/chunked/vtable/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index b152311e6bb..5d2d7d35734 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/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::ArrayEq; use crate::ArrayHash; use crate::ArrayRef; @@ -237,6 +238,19 @@ impl VTable for Chunked { } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + // 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::() { + return Ok(ExecutionResult::execute_slot::( + array, + i + CHUNKS_OFFSET, + )); + } + } + + // All chunks are now canonical — combine them. Ok(ExecutionResult::done(_canonicalize(array.as_view(), ctx)?)) } diff --git a/vortex-array/src/arrays/filter/vtable.rs b/vortex-array/src/arrays/filter/vtable.rs index 75258f542e5..3237a680d8e 100644 --- a/vortex-array/src/arrays/filter/vtable.rs +++ b/vortex-array/src/arrays/filter/vtable.rs @@ -11,6 +11,7 @@ use vortex_error::vortex_panic; use vortex_mask::Mask; use vortex_session::VortexSession; +use crate::AnyCanonical; use crate::ArrayEq; use crate::ArrayHash; use crate::ArrayRef; @@ -34,6 +35,7 @@ use crate::buffer::BufferHandle; use crate::dtype::DType; use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; +use crate::require_child; use crate::scalar::Scalar; use crate::serde::ArrayChildren; use crate::validity::Validity; @@ -142,14 +144,18 @@ impl VTable for Filter { if let Some(canonical) = execute_filter_fast_paths(array.as_view(), ctx)? { return Ok(ExecutionResult::done(canonical)); } + + let array = require_child!(array, array.child(), CHILD_SLOT => AnyCanonical); + let Mask::Values(mask_values) = &array.mask else { unreachable!("`execute_filter_fast_paths` handles AllTrue and AllFalse") }; - // We rely on the optimization pass that runs prior to this execution for filter pushdown, - // so now we can just execute the filter without worrying. + // Child is pre-canonicalized — apply the filter directly. + debug_assert!(array.child().is_canonical()); + let child = array.child().to_canonical()?; Ok(ExecutionResult::done( - execute_filter(array.child().clone().execute(ctx)?, mask_values).into_array(), + execute_filter(child, mask_values).into_array(), )) } diff --git a/vortex-array/src/arrays/list/vtable/mod.rs b/vortex-array/src/arrays/list/vtable/mod.rs index 300ff5c59a2..d5403629765 100644 --- a/vortex-array/src/arrays/list/vtable/mod.rs +++ b/vortex-array/src/arrays/list/vtable/mod.rs @@ -23,9 +23,11 @@ use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayView; use crate::array::VTable; +use crate::arrays::Primitive; use crate::arrays::list::ListArrayExt; use crate::arrays::list::ListData; use crate::arrays::list::array::NUM_SLOTS; +use crate::arrays::list::array::OFFSETS_SLOT; use crate::arrays::list::array::SLOT_NAMES; use crate::arrays::list::compute::PARENT_KERNELS; use crate::arrays::list::compute::rules::PARENT_RULES; @@ -34,6 +36,7 @@ use crate::buffer::BufferHandle; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; +use crate::require_child; use crate::serde::ArrayChildren; use crate::validity::Validity; use crate::vtable; @@ -114,7 +117,7 @@ impl VTable for List { let elements = slots[crate::arrays::list::array::ELEMENTS_SLOT] .as_ref() .vortex_expect("ListArray elements slot"); - let offsets = slots[crate::arrays::list::array::OFFSETS_SLOT] + let offsets = slots[OFFSETS_SLOT] .as_ref() .vortex_expect("ListArray offsets slot"); vortex_ensure!( @@ -180,6 +183,7 @@ impl VTable for List { } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let array = require_child!(array, array.offsets(), OFFSETS_SLOT => Primitive); Ok(ExecutionResult::done( list_view_from_list(array, ctx)?.into_array(), )) diff --git a/vortex-array/src/arrays/masked/vtable/mod.rs b/vortex-array/src/arrays/masked/vtable/mod.rs index 337ce1b8dc3..5ed996c1de9 100644 --- a/vortex-array/src/arrays/masked/vtable/mod.rs +++ b/vortex-array/src/arrays/masked/vtable/mod.rs @@ -13,10 +13,10 @@ use vortex_error::vortex_ensure; use vortex_error::vortex_panic; use vortex_session::VortexSession; +use crate::AnyCanonical; use crate::ArrayEq; use crate::ArrayHash; use crate::ArrayRef; -use crate::Canonical; use crate::IntoArray; use crate::Precision; use crate::array::Array; @@ -35,6 +35,7 @@ use crate::buffer::BufferHandle; use crate::dtype::DType; use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; +use crate::require_child; use crate::scalar::Scalar; use crate::serde::ArrayChildren; use crate::validity::Validity; @@ -166,7 +167,10 @@ impl VTable for Masked { // While we could manually convert the dtype, `mask_validity_canonical` is already O(1) for // `AllTrue` masks (no data copying), so there's no benefit. - let child = array.child().clone().execute::(ctx)?; + let array = require_child!(array, array.child(), CHILD_SLOT => AnyCanonical); + + debug_assert!(array.child().is_canonical()); + let child = array.child().to_canonical()?; Ok(ExecutionResult::done( mask_validity_canonical(child, &validity_mask, ctx)?.into_array(), )) diff --git a/vortex-array/src/arrays/patched/vtable/mod.rs b/vortex-array/src/arrays/patched/vtable/mod.rs index f30094f57fe..64da095f710 100644 --- a/vortex-array/src/arrays/patched/vtable/mod.rs +++ b/vortex-array/src/arrays/patched/vtable/mod.rs @@ -31,10 +31,15 @@ use crate::array::ArrayView; use crate::array::VTable; use crate::array::ValidityChild; use crate::array::ValidityVTableFromChild; +use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::patched::PatchedArrayExt; use crate::arrays::patched::PatchedData; +use crate::arrays::patched::array::INDICES_SLOT; +use crate::arrays::patched::array::INNER_SLOT; +use crate::arrays::patched::array::LANE_OFFSETS_SLOT; use crate::arrays::patched::array::SLOT_NAMES; +use crate::arrays::patched::array::VALUES_SLOT; use crate::arrays::patched::compute::rules::PARENT_RULES; use crate::arrays::patched::vtable::kernels::PARENT_KERNELS; use crate::arrays::primitive::PrimitiveDataParts; @@ -45,6 +50,7 @@ use crate::dtype::DType; use crate::dtype::NativePType; use crate::dtype::PType; use crate::match_each_native_ptype; +use crate::require_child; use crate::serde::ArrayChildren; use crate::vtable; @@ -242,12 +248,18 @@ impl VTable for Patched { SLOT_NAMES[idx].to_string() } - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + let array = require_child!(array, array.base_array(), INNER_SLOT => Primitive); + let array = require_child!(array, array.lane_offsets(), LANE_OFFSETS_SLOT => Primitive); + let array = require_child!(array, array.patch_indices(), INDICES_SLOT => Primitive); + let array = require_child!(array, array.patch_values(), VALUES_SLOT => Primitive); + let inner = array .base_array() .clone() - .execute::(ctx)? - .into_primitive(); + .try_downcast::() + .ok() + .vortex_expect("base_array pre-canonicalized to Primitive"); let PrimitiveDataParts { buffer, @@ -258,17 +270,21 @@ impl VTable for Patched { let lane_offsets = array .lane_offsets() .clone() - .execute::(ctx)?; + .try_downcast::() + .ok() + .vortex_expect("lane_offsets pre-canonicalized to Primitive"); let indices = array .patch_indices() .clone() - .execute::(ctx)?; - - // TODO(aduffy): add support for non-primitive PatchedArray patches application (?) + .try_downcast::() + .ok() + .vortex_expect("patch_indices pre-canonicalized to Primitive"); let values = array .patch_values() .clone() - .execute::(ctx)?; + .try_downcast::() + .ok() + .vortex_expect("patch_values pre-canonicalized to Primitive"); let patched_values = match_each_native_ptype!(values.ptype(), |V| { let offset = array.offset(); diff --git a/vortex-array/src/arrays/slice/vtable.rs b/vortex-array/src/arrays/slice/vtable.rs index b5552cdec3c..1cdb3c49dac 100644 --- a/vortex-array/src/arrays/slice/vtable.rs +++ b/vortex-array/src/arrays/slice/vtable.rs @@ -18,7 +18,6 @@ use crate::AnyCanonical; use crate::ArrayEq; use crate::ArrayHash; use crate::ArrayRef; -use crate::Canonical; use crate::IntoArray; use crate::Precision; use crate::array::Array; @@ -36,6 +35,7 @@ use crate::buffer::BufferHandle; use crate::dtype::DType; use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; +use crate::require_child; use crate::scalar::Scalar; use crate::serde::ArrayChildren; use crate::validity::Validity; @@ -141,20 +141,15 @@ impl VTable for Slice { vortex_bail!("Slice array is not serializable") } - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - // Execute the child to get canonical form, then slice it - let Some(canonical) = array.child().as_opt::() else { - // If the child is not canonical, recurse. - return array - .child() - .clone() - .execute::(ctx)? - .slice(array.slice_range().clone()) - .map(ExecutionResult::done); - }; + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + let array = require_child!(array, array.child(), CHILD_SLOT => AnyCanonical); + // Child is now canonical — slice it. // TODO(ngates): we should inline canonical slice logic here. - Canonical::from(canonical) + debug_assert!(array.child().is_canonical()); + array + .child() + .to_canonical()? .into_array() .slice(array.range.clone()) .map(ExecutionResult::done) diff --git a/vortex-array/src/arrays/varbin/vtable/mod.rs b/vortex-array/src/arrays/varbin/vtable/mod.rs index b1fd1410b68..df7167419d9 100644 --- a/vortex-array/src/arrays/varbin/vtable/mod.rs +++ b/vortex-array/src/arrays/varbin/vtable/mod.rs @@ -18,14 +18,17 @@ use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayView; use crate::array::VTable; +use crate::arrays::Primitive; use crate::arrays::varbin::VarBinArrayExt; use crate::arrays::varbin::VarBinData; use crate::arrays::varbin::array::NUM_SLOTS; +use crate::arrays::varbin::array::OFFSETS_SLOT; use crate::arrays::varbin::array::SLOT_NAMES; use crate::buffer::BufferHandle; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; +use crate::require_child; use crate::serde::ArrayChildren; use crate::validity::Validity; use crate::vtable; @@ -89,7 +92,7 @@ impl VTable for VarBin { "VarBinArray expected {NUM_SLOTS} slots, found {}", slots.len() ); - let offsets = slots[crate::arrays::varbin::array::OFFSETS_SLOT] + let offsets = slots[OFFSETS_SLOT] .as_ref() .vortex_expect("VarBinArray offsets slot"); vortex_ensure!( @@ -187,6 +190,7 @@ impl VTable for VarBin { } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let array = require_child!(array, array.offsets(), OFFSETS_SLOT => Primitive); Ok(ExecutionResult::done( varbin_to_canonical(array.as_view(), ctx)?.into_array(), )) diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 623410938d2..0688fe533f6 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -105,6 +105,9 @@ impl ArrayRef { }); let mut current = self.optimize()?; + // self is now unused + drop(self); + // Stack frames: (parent, slot_idx, done_predicate_for_slot) let mut stack: Vec<(ArrayRef, usize, DonePredicate)> = Vec::new(); @@ -120,7 +123,7 @@ impl ArrayRef { return Ok(current); } Some((parent, slot_idx, _)) => { - current = parent.with_slot(slot_idx, current)?; + current = unsafe { parent.with_slot_unchecked(slot_idx, current) }; current = current.optimize()?; continue; } @@ -136,7 +139,7 @@ impl ArrayRef { return Ok(current); } Some((parent, slot_idx, _)) => { - current = parent.with_slot(slot_idx, current)?; + current = unsafe { parent.with_slot_unchecked(slot_idx, current) }; current = current.optimize()?; continue; } @@ -335,7 +338,7 @@ impl Executable for ArrayRef { // replacing it, and returning the updated array. let child = array.slots()[i].clone().vortex_expect("valid slot index"); let executed_child = child.execute::(ctx)?; - array.with_slot(i, executed_child) + Ok(unsafe { array.with_slot_unchecked(i, executed_child) }) } } }