better execute path

.github/runs-on.yml

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

vortex-array/benches/chunked_dict_builder.rs

@@ -31,34 +31,19 @@ const BENCH_ARGS: &[(usize, usize, usize)] = &[
     (1000, 1000, 100),
 ];
 
-#[divan::bench(types = [u32, u64, f32, f64], args = BENCH_ARGS)]
-fn chunked_dict_primitive_canonical_into<T: NativePType>(
-    bencher: Bencher,
-    (len, unique_values, chunk_count): (usize, usize, usize),
-) where
-    StandardUniform: Distribution<T>,
-{
-    let chunk = gen_dict_primitive_chunks::<T, u16>(len, unique_values, chunk_count);
-
-    bencher.with_inputs(|| &chunk).bench_refs(|chunk| {
-        let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count);
-        chunk
-            .append_to_builder(builder.as_mut(), &mut SESSION.create_execution_ctx())
-            .vortex_expect("append failed");
-        builder.finish()
-    })
-}
-
 #[divan::bench(types = [u32, u64, f32, f64], args = BENCH_ARGS)]
 fn chunked_dict_primitive_into_canonical<T: NativePType>(
     bencher: Bencher,
     (len, unique_values, chunk_count): (usize, usize, usize),
 ) where
     StandardUniform: Distribution<T>,
 {
-    let chunk = gen_dict_primitive_chunks::<T, u16>(len, unique_values, chunk_count);
-
     bencher
-        .with_inputs(|| chunk.clone())
-        .bench_values(|chunk| chunk.execute::<Canonical>(&mut SESSION.create_execution_ctx()))
+        .with_inputs(|| {
+            (
+                gen_dict_primitive_chunks::<T, u16>(len, unique_values, chunk_count),
+                SESSION.create_execution_ctx(),
+            )
+        })
+        .bench_values(|(chunk, mut ctx)| (chunk.execute::<Canonical>(&mut ctx), ctx))
 }

vortex-array/src/array/erased.rs

@@ -50,6 +50,7 @@ use crate::expr::stats::Precision;
 use crate::expr::stats::Stat;
 use crate::expr::stats::StatsProviderExt;
 use crate::matcher::Matcher;
+use crate::matcher::MatcherHint;
 use crate::optimizer::ArrayOptimizer;
 use crate::scalar::Scalar;
 use crate::stats::StatsSetRef;
@@ -151,6 +152,18 @@ impl ArrayRef {
         self.0.encoding_id()
     }
 
+    /// Returns the interned encoding index of the array.
+    #[inline]
+    pub fn encoding_idx(&self) -> u32 {
+        self.0.encoding_idx()
+    }
+
+    /// Returns the category flags of the array's encoding.
+    #[inline]
+    pub fn encoding_categories(&self) -> u32 {
+        self.0.encoding_categories()
+    }
+
     /// Performs a constant-time slice of the array.
     pub fn slice(&self, range: Range<usize>) -> VortexResult<ArrayRef> {
         let len = self.len();
@@ -395,6 +408,35 @@ impl ArrayRef {
         self.is::<AnyCanonical>()
     }
 
+    /// 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.
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure `slot_idx` names an existing slot and `replacement` is a valid
+    /// physical substitute for that slot, meaning it preserves the slot's logical `DType` and
+    /// `len` as expected by the enclosing array encoding.
+    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
@@ -611,6 +653,18 @@ impl IntoArray for ArrayRef {
 impl<V: VTable> Matcher for V {
     type Match<'a> = ArrayView<'a, V>;
 
+    fn dispatch_hint() -> Option<MatcherHint> {
+        let cats = V::category_flags();
+        if cats != 0 {
+            return Some(MatcherHint::Category(cats));
+        }
+        let s = V::static_id();
+        if s.is_empty() {
+            return None;
+        }
+        Some(MatcherHint::Exact(crate::intern(s)))
+    }
+
     fn matches(array: &ArrayRef) -> bool {
         array.0.as_any().is::<ArrayInner<V>>()
     }

vortex-array/src/array/mod.rs

@@ -70,6 +70,12 @@ pub(crate) trait DynArray: 'static + private::Sealed + Send + Sync + Debug {
     /// Returns the encoding ID of the array.
     fn encoding_id(&self) -> ArrayId;
 
+    /// Returns the interned encoding index for this array.
+    fn encoding_idx(&self) -> u32;
+
+    /// Returns the category flags for this array's encoding.
+    fn encoding_categories(&self) -> u32;
+
     /// Fetch the scalar at the given index.
     ///
     /// This method panics if the index is out of bounds for the array.
@@ -141,6 +147,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<Option<ArrayRef>>) -> VortexResult<ArrayRef>;
 
@@ -214,6 +223,14 @@ impl<V: VTable> DynArray for ArrayInner<V> {
         self.vtable.id()
     }
 
+    fn encoding_idx(&self) -> u32 {
+        self.encoding_idx
+    }
+
+    fn encoding_categories(&self) -> u32 {
+        self.encoding_categories
+    }
+
     fn scalar_at(&self, this: &ArrayRef, index: usize) -> VortexResult<Scalar> {
         let view = unsafe { ArrayView::new_unchecked(this, &self.data) };
         <V::OperationsVTable as OperationsVTable<V>>::scalar_at(
@@ -378,6 +395,10 @@ impl<V: VTable> DynArray for ArrayInner<V> {
             })
     }
 
+    fn replace_slot(&mut self, slot_idx: usize, replacement: ArrayRef) {
+        self.slots[slot_idx] = Some(replacement);
+    }
+
     fn with_slots(&self, this: ArrayRef, slots: Vec<Option<ArrayRef>>) -> VortexResult<ArrayRef> {
         let data = self.data.clone();
         let stats = this.statistics().to_owned();

vortex-array/src/array/typed.rs

@@ -78,6 +78,8 @@ impl<V: VTable> TypedArrayRef<V> for ArrayView<'_, V> {}
 #[doc(hidden)]
 pub(crate) struct ArrayInner<V: VTable> {
     pub(crate) vtable: V,
+    pub(crate) encoding_idx: u32,
+    pub(crate) encoding_categories: u32,
     pub(crate) dtype: DType,
     pub(crate) len: usize,
     pub(crate) data: V::ArrayData,
@@ -91,7 +93,15 @@ impl<V: VTable> ArrayInner<V> {
     pub fn try_new(new: ArrayParts<V>) -> VortexResult<Self> {
         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<V>) -> Self {
+        unsafe {
             Self::from_data_unchecked(
                 new.vtable,
                 new.dtype,
@@ -100,7 +110,7 @@ impl<V: VTable> ArrayInner<V> {
                 new.slots,
                 ArrayStats::default(),
             )
-        })
+        }
     }
 
     /// Create without validation.
@@ -115,8 +125,12 @@ impl<V: VTable> ArrayInner<V> {
         slots: Vec<Option<ArrayRef>>,
         stats: ArrayStats,
     ) -> Self {
+        let encoding_idx = crate::intern(vtable.id().as_ref());
+        let encoding_categories = V::category_flags();
         Self {
             vtable,
+            encoding_idx,
+            encoding_categories,
             dtype,
             len,
             data,
@@ -143,6 +157,8 @@ impl<V: VTable> Clone for ArrayInner<V> {
     fn clone(&self) -> Self {
         Self {
             vtable: self.vtable.clone(),
+            encoding_idx: self.encoding_idx,
+            encoding_categories: self.encoding_categories,
             dtype: self.dtype.clone(),
             len: self.len,
             data: self.data.clone(),
@@ -209,16 +225,7 @@ impl<V: VTable> Array<V> {
     /// Caller must ensure the provided parts are logically consistent.
     #[doc(hidden)]
     pub unsafe fn from_parts_unchecked(new: ArrayParts<V>) -> Self {
-        let inner = ArrayRef::from_inner(Arc::new(unsafe {
-            ArrayInner::<V>::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,

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

@@ -59,6 +59,34 @@ pub trait VTable: 'static + Clone + Sized + Send + Sync + Debug {
     /// Returns the ID of the array.
     fn id(&self) -> ArrayId;
 
+    /// Returns the compile-time string ID for this VTable type.
+    ///
+    /// This is used to pre-compute dispatch hints at rule-registration time without
+    /// needing an instance. VTables with instance-dependent IDs (e.g., `ScalarFnVTable`,
+    /// which stores the scalar function's own ID) should NOT override this; instead they
+    /// should set appropriate `category_flags()`.
+    ///
+    /// The default implementation returns an empty string, which causes the dispatch
+    /// system to fall back to the `category_flags()` path (or no hint if both are zero).
+    fn static_id() -> &'static str
+    where
+        Self: Sized,
+    {
+        ""
+    }
+
+    /// Returns the category membership bitmask for this VTable type.
+    ///
+    /// Use the `CATEGORY_*` constants from [`crate::matcher`]. The default is `0`
+    /// (no special category). Override for canonical arrays, constant arrays, and
+    /// scalar-function arrays.
+    fn category_flags() -> u32
+    where
+        Self: Sized,
+    {
+        0
+    }
+
     /// Validates that externally supplied logical metadata matches the array data.
     fn validate(
         &self,

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

@@ -70,6 +70,14 @@ impl VTable for Bool {
         Self::ID
     }
 
+    fn static_id() -> &'static str {
+        "vortex.bool"
+    }
+
+    fn category_flags() -> u32 {
+        crate::matcher::CATEGORY_CANONICAL
+    }
+
     fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
         1
     }

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<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).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::<u64>::with_capacity(len);
 
     for chunk in chunks {
-        let chunk_array = chunk.clone().execute::<ListViewArray>(ctx)?;
+        let chunk_array = chunk
+            .clone()
+            .try_downcast::<ListView>()
+            .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)?;