Optimise builder execution loop

vortex-array/src/array/mod.rs

@@ -121,9 +121,6 @@ pub(crate) trait DynArrayData: '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;
-
     /// Returns the [`Validity`] of the array.
     fn validity(&self, this: &ArrayRef) -> VortexResult<Validity>;
 
@@ -271,10 +268,6 @@ impl<V: VTable> DynArrayData for ArrayData<V> {
         self
     }
 
-    fn as_any_mut(&mut self) -> &mut dyn Any {
-        self
-    }
-
     fn validity(&self, this: &ArrayRef) -> VortexResult<Validity> {
         if this.dtype().is_nullable() {
             let view = unsafe { ArrayView::new_unchecked(this, &self.data) };

vortex-array/src/array/typed.rs

@@ -306,8 +306,12 @@ impl<V: VTable> Array<V> {
     /// Returns `None` when this handle is not the unique owner of the backing allocation.
     pub fn data_mut(&mut self) -> Option<&mut V::TypedArrayData> {
         let store = self.inner.inner_mut()?;
-        let array_inner = store.data.as_any_mut().downcast_mut::<ArrayData<V>>();
-        Some(&mut array_inner?.data)
+        // NOTE(ngates): use downcast_mut_unchecked when it becomes stable
+        debug_assert!(store.data.as_any().is::<ArrayData<V>>());
+        // SAFETY: `Array<V>` guarantees the inner is `ArrayData<V>`, as `downcast_inner` relies on.
+        let array_inner =
+            unsafe { &mut *std::ptr::from_mut(&mut store.data).cast::<ArrayData<V>>() };
+        Some(&mut array_inner.data)
     }
 
     /// Returns the full typed array construction parts if this handle owns the allocation.

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

@@ -274,10 +274,13 @@ impl VTable for Chunked {
                         array.with_next_builder_slot(slot_idx + 1),
                         slot_idx,
                     ))
-                } else {
+                } else if slot_idx == ChunkedSlots::CHUNKS_OFFSET {
+                    // No chunks, so nothing was appended and there is no builder to finish.
                     Ok(ExecutionResult::done(
                         Canonical::empty(array.dtype()).into_array(),
                     ))
+                } else {
+                    Ok(ExecutionResult::done_into_builder(array))
                 }
             }
         }

vortex-array/src/canonical.rs

@@ -4,6 +4,7 @@
 //! Encodings that enable zero-copy sharing of data with Arrow.
 
 use std::sync::Arc;
+use std::sync::LazyLock;
 
 use vortex_buffer::BitBuffer;
 use vortex_buffer::Buffer;
@@ -17,8 +18,10 @@ use crate::ArraySlots;
 use crate::Executable;
 use crate::ExecutionCtx;
 use crate::IntoArray;
+use crate::array::ArrayId;
 use crate::array::ArrayView;
 use crate::array::child_to_validity;
+use crate::array::vtable::VTable as _;
 use crate::arrays::Bool;
 use crate::arrays::BoolArray;
 use crate::arrays::Decimal;
@@ -1232,63 +1235,88 @@ impl CanonicalView<'_> {
 
 /// A matcher for any canonical array type.
 pub struct AnyCanonical;
-impl Matcher for AnyCanonical {
-    type Match<'a> = CanonicalView<'a>;
-
-    #[inline]
-    fn matches(array: &ArrayRef) -> bool {
-        array.is::<Null>()
-            || array.is::<Bool>()
-            || array.is::<Primitive>()
-            || array.is::<Decimal>()
-            || array.is::<Struct>()
-            || array.is::<Union>()
-            || array.is::<ListView>()
-            || array.is::<Map>()
-            || array.is::<FixedSizeList>()
-            || array.is::<VarBinView>()
-            || array.is::<Variant>()
-            || array.is::<Extension>()
-    }
 
-    #[inline]
-    fn try_match(array: &ArrayRef) -> Option<Self::Match<'_>> {
-        if let Some(a) = array.as_opt::<Null>() {
-            Some(CanonicalView::Null(a))
-        } else if let Some(a) = array.as_opt::<Bool>() {
-            Some(CanonicalView::Bool(a))
-        } else if let Some(a) = array.as_opt::<Primitive>() {
-            Some(CanonicalView::Primitive(a))
-        } else if let Some(a) = array.as_opt::<Decimal>() {
-            Some(CanonicalView::Decimal(a))
-        } else if let Some(a) = array.as_opt::<Struct>() {
-            Some(CanonicalView::Struct(a))
-        } else if let Some(a) = array.as_opt::<Union>() {
-            Some(CanonicalView::Union(a))
-        } else if let Some(a) = array.as_opt::<ListView>() {
-            Some(CanonicalView::List(a))
-        } else if let Some(a) = array.as_opt::<Map>() {
-            Some(CanonicalView::Map(a))
-        } else if let Some(a) = array.as_opt::<FixedSizeList>() {
-            Some(CanonicalView::FixedSizeList(a))
-        } else if let Some(a) = array.as_opt::<VarBinView>() {
-            Some(CanonicalView::VarBinView(a))
-        } else if let Some(a) = array.as_opt::<Variant>() {
-            Some(CanonicalView::Variant(a))
-        } else {
-            array.as_opt::<Extension>().map(CanonicalView::Extension)
+/// The canonical encodings, as `field => vtable => CanonicalView variant` triples.
+macro_rules! with_canonical_encodings {
+    ($mac:ident) => {
+        $mac! {
+            null => Null => Null,
+            bool_ => Bool => Bool,
+            primitive => Primitive => Primitive,
+            decimal => Decimal => Decimal,
+            struct_ => Struct => Struct,
+            union_ => Union => Union,
+            list => ListView => List,
+            map => Map => Map,
+            fixed_size_list => FixedSizeList => FixedSizeList,
+            varbinview => VarBinView => VarBinView,
+            variant => Variant => Variant,
+            extension => Extension => Extension,
         }
-    }
+    };
+}
+
+/// Expands `with_canonical_encodings` into the id cache and [`AnyCanonical`]'s [`Matcher`] impl.
+macro_rules! canonical_matcher {
+    ($($field:ident => $vtable:ident => $variant:ident,)+) => {
+        /// The encoding ids of the canonical encodings, interned once.
+        struct CanonicalIds {
+            $($field: ArrayId,)+
+        }
+
+        static CANONICAL_IDS: LazyLock<CanonicalIds> = LazyLock::new(|| CanonicalIds {
+            $($field: $vtable.id(),)+
+        });
+
+        impl Matcher for AnyCanonical {
+            type Match<'a> = CanonicalView<'a>;
+
+            #[inline]
+            fn matches(array: &ArrayRef) -> bool {
+                // The id selects, the downcast decides: `ForeignArray`, `ScalarFn` and the Python
+                // vtable each return a per-instance `self.id`, so one could be registered under a
+                // canonical encoding's id. Answering yes where `try_match` answers `None` panics
+                // `Canonical::execute`.
+                let ids = &*CANONICAL_IDS;
+                let id = array.encoding_id();
+
+                // One `||` reduction keeps the common rejection branchless.
+                if !($(id == ids.$field ||)+ false) {
+                    return false;
+                }
+
+                $(if id == ids.$field {
+                    return array.is::<$vtable>();
+                })+
+                false
+            }
+
+            #[inline]
+            fn try_match(array: &ArrayRef) -> Option<Self::Match<'_>> {
+                // Selected by id and confirmed by downcast, as in `matches` above.
+                let ids = &*CANONICAL_IDS;
+                let id = array.encoding_id();
+                $(if id == ids.$field {
+                    return array.as_opt::<$vtable>().map(CanonicalView::$variant);
+                })+
+                None
+            }
+        }
+    };
 }
 
+with_canonical_encodings!(canonical_matcher);
+
 #[cfg(test)]
 mod test {
+    use std::sync::Arc;
     use std::sync::LazyLock;
 
     use vortex_error::VortexResult;
     use vortex_error::vortex_err;
     use vortex_session::VortexSession;
 
+    use crate::AnyCanonical;
     use crate::ArrayRef;
     use crate::Canonical;
     use crate::CanonicalValidity;
@@ -1303,12 +1331,88 @@ mod test {
     use crate::arrays::struct_::StructArrayExt;
     use crate::arrays::variant::VariantArraySlotsExt;
     use crate::canonical::StructArray;
+    use crate::dtype::DType;
+    use crate::dtype::DecimalDType;
+    use crate::dtype::MapDType;
     use crate::dtype::Nullability;
+    use crate::dtype::PType;
+    use crate::dtype::StructFields;
+    use crate::dtype::UnionVariants;
+    use crate::extension::datetime::Time;
+    use crate::extension::datetime::TimeUnit;
+    use crate::matcher::Matcher;
     use crate::scalar::Scalar;
 
     /// A shared session for these canonical tests, used to create execution contexts.
     static SESSION: LazyLock<VortexSession> = LazyLock::new(crate::array_session);
 
+    /// One empty array per canonical encoding, covering every arm the matcher generates.
+    fn one_array_per_canonical_encoding() -> VortexResult<Vec<ArrayRef>> {
+        let i32_dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
+        let dtypes = [
+            DType::Null,
+            DType::Bool(Nullability::NonNullable),
+            i32_dtype.clone(),
+            DType::Decimal(DecimalDType::new(10, 2), Nullability::NonNullable),
+            DType::Utf8(Nullability::NonNullable),
+            DType::List(Arc::new(i32_dtype.clone()), Nullability::NonNullable),
+            DType::Map(
+                MapDType::try_new(i32_dtype.clone(), i32_dtype.clone(), false)?,
+                Nullability::NonNullable,
+            ),
+            DType::FixedSizeList(Arc::new(i32_dtype.clone()), 2, Nullability::NonNullable),
+            DType::Struct(
+                StructFields::new(["a"].into(), vec![i32_dtype.clone()]),
+                Nullability::NonNullable,
+            ),
+            DType::Union(
+                UnionVariants::try_new(["a"].into(), vec![i32_dtype], vec![0])?,
+                Nullability::NonNullable,
+            ),
+            DType::Extension(Time::new(TimeUnit::Seconds, Nullability::NonNullable).erased()),
+        ];
+
+        let mut arrays: Vec<ArrayRef> = dtypes
+            .iter()
+            .map(|dtype| Canonical::empty(dtype).into_array())
+            .collect();
+        // `Canonical::empty` rejects `DType::Variant`, so build that one directly.
+        arrays.push(VariantArray::try_new(variant_core_storage(0), None)?.into_array());
+
+        Ok(arrays)
+    }
+
+    /// Every canonical encoding must reach its own arm of both halves of [`AnyCanonical`].
+    ///
+    /// They are separate expansions, and the executor stops on `matches` while
+    /// `Canonical::execute` unwraps `try_match`, so a disagreement between them is a panic.
+    #[test]
+    fn every_canonical_encoding_matches_any_canonical() -> VortexResult<()> {
+        let arrays = one_array_per_canonical_encoding()?;
+        assert_eq!(
+            arrays.len(),
+            12,
+            "expected one array per canonical encoding"
+        );
+
+        for array in arrays {
+            assert!(
+                AnyCanonical::matches(&array),
+                "{} array of dtype {} did not match AnyCanonical",
+                array.encoding_id(),
+                array.dtype(),
+            );
+            assert!(
+                AnyCanonical::try_match(&array).is_some(),
+                "{} array of dtype {} did not view as AnyCanonical",
+                array.encoding_id(),
+                array.dtype(),
+            );
+        }
+
+        Ok(())
+    }
+
     fn variant_core_storage(len: usize) -> ArrayRef {
         ConstantArray::new(
             Scalar::variant(Scalar::primitive(1i32, Nullability::NonNullable)),

vortex-array/src/executor.rs

@@ -12,6 +12,7 @@
 //! See <https://docs.vortex.dev/developer-guide/internals/execution> for the full execution
 //! narrative, diagrams, and walkthroughs.
 
+use std::any::TypeId;
 use std::env::VarError;
 use std::fmt;
 use std::fmt::Display;
@@ -166,7 +167,13 @@ impl ArrayRef {
     /// parent rewrite would observe inconsistent state and could discard accumulated builder
     /// data.
     #[allow(clippy::cognitive_complexity)]
-    pub fn execute_until<M: Matcher>(self, ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
+    pub fn execute_until<M: Matcher + 'static>(
+        self,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<ArrayRef> {
+        // `execute::<Canonical>` passes `AnyCanonical` as the target, making it the same predicate
+        // as the loop's universal stop condition. Folds to a constant per monomorphization.
+        let target_is_any_canonical = TypeId::of::<M>() == TypeId::of::<AnyCanonical>();
         let mut current_array = self;
         let mut current_builder: Option<Box<dyn ArrayBuilder>> = None;
         let mut stack: Vec<StackFrame> = Vec::new();
@@ -186,12 +193,22 @@ impl ArrayRef {
                 current_builder.is_some(),
             ));
 
-            let is_done = stack
-                .last()
-                .map_or(M::matches as DonePredicate, |frame| frame.done);
-
-            let done_target = is_done(&current_array);
-            let done_canonical = AnyCanonical::matches(&current_array);
+            let (done_target, done_canonical) = match stack.last() {
+                // At the root one scan can answer both, rather than scanning the encoding twice.
+                None => {
+                    let done_target = M::matches(&current_array);
+                    let done_canonical = if target_is_any_canonical {
+                        done_target
+                    } else {
+                        AnyCanonical::matches(&current_array)
+                    };
+                    (done_target, done_canonical)
+                }
+                Some(frame) => (
+                    (frame.done)(&current_array),
+                    AnyCanonical::matches(&current_array),
+                ),
+            };
             trace_op!(record_execute_until_done_check(done_target, done_canonical));
 
             if done_target || done_canonical {
@@ -267,7 +284,8 @@ impl ArrayRef {
             }
 
             let expected_len = current_array.len();
-            let expected_dtype = current_array.dtype().clone();
+            // Only `finalize_done` reads this back, and only under debug assertions.
+            let expected_dtype = cfg!(debug_assertions).then(|| current_array.dtype().clone());
             let stats = current_array.statistics().to_array_stats();
             let encoding_id = current_array.encoding_id();
             trace_op!(record_execute_encoding(&current_array));
@@ -591,7 +609,7 @@ fn finalize_done(
     result: ArrayRef,
     mut builder: Option<Box<dyn ArrayBuilder>>,
     expected_len: usize,
-    expected_dtype: DType,
+    expected_dtype: Option<DType>,
     stats: ArrayStats,
     encoding_id: ArrayId,
 ) -> VortexResult<(ArrayRef, Option<Box<dyn ArrayBuilder>>)> {
@@ -601,7 +619,7 @@ fn finalize_done(
         result
     };
 
-    if cfg!(debug_assertions) {
+    if let Some(expected_dtype) = expected_dtype {
         vortex_ensure!(
             output.len() == expected_len,
             "Result length mismatch for {:?}",
@@ -781,6 +799,21 @@ impl ExecutionResult {
         }
     }
 
+    /// Signal that execution is complete and the result is in the executor's active builder.
+    ///
+    /// Pass the consumed parent array: the executor discards it after finishing the builder,
+    /// avoiding an allocation for a placeholder result.
+    ///
+    /// Only valid once at least one [`ExecutionStep::AppendChild`] has been returned, which
+    /// guarantees that the executor has an active builder. The parent may have empty child slots
+    /// because its children have already been appended; it must not escape the executor.
+    pub fn done_into_builder(array: impl IntoArray) -> Self {
+        Self {
+            array: array.into_array(),
+            step: ExecutionStep::Done,
+        }
+    }
+
     /// Request execution of slot at `slot_idx` until it matches the given [`Matcher`].
     ///
     /// The provided array is the (possibly modified) parent that still needs its slot executed.