feat[array]: owned matchers

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

Author: joseph-isaacs

vortex-array/public-api.lock

@@ -3618,6 +3618,12 @@ pub fn vortex_array::arrays::scalar_fn::AnyScalarFn::matches(array: &dyn vortex_
 
 pub fn vortex_array::arrays::scalar_fn::AnyScalarFn::try_match(array: &dyn vortex_array::DynArray) -> core::option::Option<Self::Match>
 
+impl vortex_array::matcher::OwnedMatcher for vortex_array::arrays::scalar_fn::AnyScalarFn
+
+pub type vortex_array::arrays::scalar_fn::AnyScalarFn::OwnedMatch = vortex_array::arrays::scalar_fn::ScalarFnArray
+
+pub fn vortex_array::arrays::scalar_fn::AnyScalarFn::maybe_match(array: vortex_array::ArrayRef) -> core::option::Option<Self::OwnedMatch>
+
 pub struct vortex_array::arrays::scalar_fn::ExactScalarFn<F: vortex_array::scalar_fn::ScalarFnVTable>(_)
 
 impl<F: core::default::Default + vortex_array::scalar_fn::ScalarFnVTable> core::default::Default for vortex_array::arrays::scalar_fn::ExactScalarFn<F>
@@ -15068,6 +15074,12 @@ pub fn vortex_array::matcher::AnyArray::matches(_array: &dyn vortex_array::DynAr
 
 pub fn vortex_array::matcher::AnyArray::try_match(array: &dyn vortex_array::DynArray) -> core::option::Option<Self::Match>
 
+impl vortex_array::matcher::OwnedMatcher for vortex_array::matcher::AnyArray
+
+pub type vortex_array::matcher::AnyArray::OwnedMatch = alloc::sync::Arc<dyn vortex_array::DynArray>
+
+pub fn vortex_array::matcher::AnyArray::maybe_match(array: vortex_array::ArrayRef) -> core::option::Option<Self::OwnedMatch>
+
 pub trait vortex_array::matcher::Matcher
 
 pub type vortex_array::matcher::Matcher::Match<'a>
@@ -15124,6 +15136,42 @@ pub fn V::matches(array: &dyn vortex_array::DynArray) -> bool
 
 pub fn V::try_match<'a>(array: &'a dyn vortex_array::DynArray) -> core::option::Option<Self::Match>
 
+pub trait vortex_array::matcher::OwnedMatcher: vortex_array::matcher::Matcher
+
+pub type vortex_array::matcher::OwnedMatcher::OwnedMatch
+
+pub fn vortex_array::matcher::OwnedMatcher::maybe_match(array: vortex_array::ArrayRef) -> core::option::Option<Self::OwnedMatch>
+
+impl vortex_array::matcher::OwnedMatcher for vortex_array::AnyCanonical
+
+pub type vortex_array::AnyCanonical::OwnedMatch = vortex_array::Canonical
+
+pub fn vortex_array::AnyCanonical::maybe_match(array: vortex_array::ArrayRef) -> core::option::Option<Self::OwnedMatch>
+
+impl vortex_array::matcher::OwnedMatcher for vortex_array::AnyColumnar
+
+pub type vortex_array::AnyColumnar::OwnedMatch = vortex_array::Columnar
+
+pub fn vortex_array::AnyColumnar::maybe_match(array: vortex_array::ArrayRef) -> core::option::Option<Self::OwnedMatch>
+
+impl vortex_array::matcher::OwnedMatcher for vortex_array::arrays::scalar_fn::AnyScalarFn
+
+pub type vortex_array::arrays::scalar_fn::AnyScalarFn::OwnedMatch = vortex_array::arrays::scalar_fn::ScalarFnArray
+
+pub fn vortex_array::arrays::scalar_fn::AnyScalarFn::maybe_match(array: vortex_array::ArrayRef) -> core::option::Option<Self::OwnedMatch>
+
+impl vortex_array::matcher::OwnedMatcher for vortex_array::matcher::AnyArray
+
+pub type vortex_array::matcher::AnyArray::OwnedMatch = alloc::sync::Arc<dyn vortex_array::DynArray>
+
+pub fn vortex_array::matcher::AnyArray::maybe_match(array: vortex_array::ArrayRef) -> core::option::Option<Self::OwnedMatch>
+
+impl<V: vortex_array::vtable::VTable> vortex_array::matcher::OwnedMatcher for V
+
+pub type V::OwnedMatch = <V as vortex_array::vtable::VTable>::Array
+
+pub fn V::maybe_match(array: vortex_array::ArrayRef) -> core::option::Option<Self::OwnedMatch>
+
 pub mod vortex_array::normalize
 
 pub enum vortex_array::normalize::Operation
@@ -24098,6 +24146,12 @@ pub fn vortex_array::AnyCanonical::matches(array: &dyn vortex_array::DynArray) -
 
 pub fn vortex_array::AnyCanonical::try_match<'a>(array: &'a dyn vortex_array::DynArray) -> core::option::Option<Self::Match>
 
+impl vortex_array::matcher::OwnedMatcher for vortex_array::AnyCanonical
+
+pub type vortex_array::AnyCanonical::OwnedMatch = vortex_array::Canonical
+
+pub fn vortex_array::AnyCanonical::maybe_match(array: vortex_array::ArrayRef) -> core::option::Option<Self::OwnedMatch>
+
 pub struct vortex_array::AnyColumnar
 
 impl vortex_array::matcher::Matcher for vortex_array::AnyColumnar
@@ -24108,6 +24162,12 @@ pub fn vortex_array::AnyColumnar::matches(array: &dyn vortex_array::DynArray) ->
 
 pub fn vortex_array::AnyColumnar::try_match<'a>(array: &'a dyn vortex_array::DynArray) -> core::option::Option<Self::Match>
 
+impl vortex_array::matcher::OwnedMatcher for vortex_array::AnyColumnar
+
+pub type vortex_array::AnyColumnar::OwnedMatch = vortex_array::Columnar
+
+pub fn vortex_array::AnyColumnar::maybe_match(array: vortex_array::ArrayRef) -> core::option::Option<Self::OwnedMatch>
+
 #[repr(transparent)] pub struct vortex_array::ArrayAdapter<V: vortex_array::vtable::VTable>(_)
 
 impl<V: vortex_array::vtable::VTable> vortex_array::ArrayAdapter<V>

vortex-array/src/array/mod.rs

@@ -51,6 +51,7 @@ use crate::expr::stats::Stat;
 use crate::expr::stats::StatsProviderExt;
 use crate::hash;
 use crate::matcher::Matcher;
+use crate::matcher::OwnedMatcher;
 use crate::optimizer::ArrayOptimizer;
 use crate::scalar::Scalar;
 use crate::scalar_fn::ReduceNode;
@@ -316,6 +317,18 @@ impl dyn DynArray + '_ {
         })
     }
 
+    /// Returns the array matched by the given owned matcher, consuming the `ArrayRef`.
+    pub fn into_<M: OwnedMatcher>(self: Arc<Self>) -> M::OwnedMatch {
+        self.try_into_matched::<M>()
+            .vortex_expect("Failed to downcast")
+    }
+
+    /// Try to match the given array using an owned matcher, returning the owned matched type
+    /// if successful.
+    pub fn try_into_matched<M: OwnedMatcher>(self: Arc<Self>) -> Option<M::OwnedMatch> {
+        M::maybe_match(self)
+    }
+
     pub fn as_constant(&self) -> Option<Scalar> {
         self.as_opt::<Constant>().map(|a| a.scalar().clone())
     }
@@ -1144,3 +1157,37 @@ impl<V: VTable> Matcher for V {
             .map(|adapter| adapter.as_inner())
     }
 }
+
+/// Implement an owned matcher for a specific VTable type.
+///
+/// During the migration, this tries both `Array<V>` (new path) and `ArrayAdapter<V>`
+/// (legacy path). Returns `V::Array`.
+impl<V: VTable> OwnedMatcher for V {
+    type OwnedMatch = V::Array;
+
+    fn maybe_match(array: ArrayRef) -> Option<Self::OwnedMatch> {
+        if !<V as Matcher>::matches(&*array) {
+            return None;
+        }
+        let any_arc = array.as_any_arc();
+        // Try new Array<V> first.
+        match any_arc.downcast::<Array<V>>() {
+            Ok(typed) => Some(match Arc::try_unwrap(typed) {
+                Ok(array) => array.into_inner(),
+                Err(arc) => arc.deref().inner().clone(),
+            }),
+            Err(any_arc) => {
+                // Fall back to legacy ArrayAdapter<V>.
+                match any_arc.downcast::<ArrayAdapter<V>>() {
+                    Ok(adapter) => Some(match Arc::try_unwrap(adapter) {
+                        Ok(adapter) => adapter.into_inner(),
+                        Err(arc) => arc.as_inner().clone(),
+                    }),
+                    Err(_) => {
+                        vortex_panic!("matches returned true but downcast failed")
+                    }
+                }
+            }
+        }
+    }
+}

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

@@ -221,9 +221,9 @@ impl VTable for Dict {
             .try_into::<Primitive>()
             .ok()
             .vortex_expect("must be primitive");
-        debug_assert!(values.is_canonical());
-        // TODO: add canonical owned cast.
-        let values = values.to_canonical()?;
+        let values = values
+            .try_into_matched::<AnyCanonical>()
+            .vortex_expect("must be canonical");
 
         Ok(ExecutionResult::done(
             take_canonical(values, &codes, ctx)?.into_array(),

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

@@ -33,6 +33,7 @@ use crate::executor::ExecutionCtx;
 use crate::executor::ExecutionResult;
 use crate::expr::Expression;
 use crate::matcher::Matcher;
+use crate::matcher::OwnedMatcher;
 use crate::scalar_fn;
 use crate::scalar_fn::Arity;
 use crate::scalar_fn::ChildName;
@@ -259,6 +260,14 @@ impl Matcher for AnyScalarFn {
     }
 }
 
+impl OwnedMatcher for AnyScalarFn {
+    type OwnedMatch = ScalarFnArray;
+
+    fn maybe_match(array: ArrayRef) -> Option<Self::OwnedMatch> {
+        array.try_into::<ScalarFnVTable>().ok()
+    }
+}
+
 /// A matcher that matches a specific scalar function expression.
 #[derive(Debug, Default)]
 pub struct ExactScalarFn<F: scalar_fn::ScalarFnVTable>(PhantomData<F>);

vortex-array/src/canonical.rs

@@ -51,6 +51,7 @@ use crate::dtype::PType;
 use crate::match_each_decimal_value_type;
 use crate::match_each_native_ptype;
 use crate::matcher::Matcher;
+use crate::matcher::OwnedMatcher;
 use crate::validity::Validity;
 
 /// An enum capturing the default uncompressed encodings for each [Vortex type](DType).
@@ -1030,6 +1031,56 @@ impl Matcher for AnyCanonical {
     }
 }
 
+impl OwnedMatcher for AnyCanonical {
+    type OwnedMatch = Canonical;
+
+    fn maybe_match(array: ArrayRef) -> Option<Self::OwnedMatch> {
+        if !<AnyCanonical as Matcher>::matches(&*array) {
+            return None;
+        }
+        let array = match array.try_into::<Null>() {
+            Ok(a) => return Some(Canonical::Null(a)),
+            Err(array) => array,
+        };
+        let array = match array.try_into::<Bool>() {
+            Ok(a) => return Some(Canonical::Bool(a)),
+            Err(array) => array,
+        };
+        let array = match array.try_into::<Primitive>() {
+            Ok(a) => return Some(Canonical::Primitive(a)),
+            Err(array) => array,
+        };
+        let array = match array.try_into::<Decimal>() {
+            Ok(a) => return Some(Canonical::Decimal(a)),
+            Err(array) => array,
+        };
+        let array = match array.try_into::<Struct>() {
+            Ok(a) => return Some(Canonical::Struct(a)),
+            Err(array) => array,
+        };
+        let array = match array.try_into::<ListView>() {
+            Ok(a) => return Some(Canonical::List(a)),
+            Err(array) => array,
+        };
+        let array = match array.try_into::<FixedSizeList>() {
+            Ok(a) => return Some(Canonical::FixedSizeList(a)),
+            Err(array) => array,
+        };
+        let array = match array.try_into::<VarBinView>() {
+            Ok(a) => return Some(Canonical::VarBinView(a)),
+            Err(array) => array,
+        };
+        let array = match array.try_into::<Variant>() {
+            Ok(a) => return Some(Canonical::Variant(a)),
+            Err(array) => array,
+        };
+        match array.try_into::<Extension>() {
+            Ok(a) => Some(Canonical::Extension(a)),
+            Err(_) => None,
+        }
+    }
+}
+
 #[cfg(test)]
 mod test {
     use std::sync::Arc;

vortex-array/src/columnar.rs

@@ -16,6 +16,7 @@ use crate::arrays::Constant;
 use crate::arrays::ConstantArray;
 use crate::dtype::DType;
 use crate::matcher::Matcher;
+use crate::matcher::OwnedMatcher;
 use crate::scalar::Scalar;
 
 /// Represents a columnnar array of data, either in canonical form or as a constant array.
@@ -110,3 +111,16 @@ impl Matcher for AnyColumnar {
         }
     }
 }
+
+impl OwnedMatcher for AnyColumnar {
+    type OwnedMatch = Columnar;
+
+    fn maybe_match(array: ArrayRef) -> Option<Self::OwnedMatch> {
+        match array.try_into::<Constant>() {
+            Ok(constant) => Some(Columnar::Constant(constant)),
+            Err(array) => {
+                <AnyCanonical as OwnedMatcher>::maybe_match(array).map(Columnar::Canonical)
+            }
+        }
+    }
+}

vortex-array/src/matcher.rs

@@ -1,13 +1,14 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use crate::ArrayRef;
 use crate::DynArray;
 
-/// Trait for matching array types.
+/// Trait for matching array types by reference.
 pub trait Matcher {
     type Match<'a>;
 
-    /// Check if the given array matches this matcher type
+    /// Check if the given array matches this matcher type.
     fn matches(array: &dyn DynArray) -> bool {
         Self::try_match(array).is_some()
     }
@@ -16,6 +17,14 @@ pub trait Matcher {
     fn try_match<'a>(array: &'a dyn DynArray) -> Option<Self::Match<'a>>;
 }
 
+/// Trait for matching array types by owned value.
+pub trait OwnedMatcher: Matcher {
+    type OwnedMatch;
+
+    /// Try to match the given array, returning the owned matched type if successful.
+    fn maybe_match(array: ArrayRef) -> Option<Self::OwnedMatch>;
+}
+
 /// Matches any array type (wildcard matcher)
 #[derive(Debug)]
 pub struct AnyArray;
@@ -33,3 +42,12 @@ impl Matcher for AnyArray {
         Some(array)
     }
 }
+
+impl OwnedMatcher for AnyArray {
+    type OwnedMatch = ArrayRef;
+
+    #[inline(always)]
+    fn maybe_match(array: ArrayRef) -> Option<Self::OwnedMatch> {
+        Some(array)
+    }
+}