diff --git a/vortex-array/public-api.lock b/vortex-array/public-api.lock index 7d8520ce85b..0e221cafa75 100644 --- a/vortex-array/public-api.lock +++ b/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 +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 + pub struct vortex_array::arrays::scalar_fn::ExactScalarFn(_) impl core::default::Default for vortex_array::arrays::scalar_fn::ExactScalarFn @@ -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 +impl vortex_array::matcher::OwnedMatcher for vortex_array::matcher::AnyArray + +pub type vortex_array::matcher::AnyArray::OwnedMatch = alloc::sync::Arc + +pub fn vortex_array::matcher::AnyArray::maybe_match(array: vortex_array::ArrayRef) -> core::option::Option + 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 +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 + +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 + +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 + +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 + +impl vortex_array::matcher::OwnedMatcher for vortex_array::matcher::AnyArray + +pub type vortex_array::matcher::AnyArray::OwnedMatch = alloc::sync::Arc + +pub fn vortex_array::matcher::AnyArray::maybe_match(array: vortex_array::ArrayRef) -> core::option::Option + +impl vortex_array::matcher::OwnedMatcher for V + +pub type V::OwnedMatch = ::Array + +pub fn V::maybe_match(array: vortex_array::ArrayRef) -> core::option::Option + 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 +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 + 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 +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 + #[repr(transparent)] pub struct vortex_array::ArrayAdapter(_) impl vortex_array::ArrayAdapter diff --git a/vortex-array/src/array/mod.rs b/vortex-array/src/array/mod.rs index 95643eff95d..4707d789b7a 100644 --- a/vortex-array/src/array/mod.rs +++ b/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_(self: Arc) -> M::OwnedMatch { + self.try_into_matched::() + .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(self: Arc) -> Option { + M::maybe_match(self) + } + pub fn as_constant(&self) -> Option { self.as_opt::().map(|a| a.scalar().clone()) } @@ -1144,3 +1157,37 @@ impl Matcher for V { .map(|adapter| adapter.as_inner()) } } + +/// Implement an owned matcher for a specific VTable type. +/// +/// During the migration, this tries both `Array` (new path) and `ArrayAdapter` +/// (legacy path). Returns `V::Array`. +impl OwnedMatcher for V { + type OwnedMatch = V::Array; + + fn maybe_match(array: ArrayRef) -> Option { + if !::matches(&*array) { + return None; + } + let any_arc = array.as_any_arc(); + // Try new Array first. + match any_arc.downcast::>() { + 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. + match any_arc.downcast::>() { + 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") + } + } + } + } + } +} diff --git a/vortex-array/src/arrays/dict/vtable/mod.rs b/vortex-array/src/arrays/dict/vtable/mod.rs index d9bfdcefc8c..c84d8aed445 100644 --- a/vortex-array/src/arrays/dict/vtable/mod.rs +++ b/vortex-array/src/arrays/dict/vtable/mod.rs @@ -221,9 +221,9 @@ impl VTable for Dict { .try_into::() .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::() + .vortex_expect("must be canonical"); Ok(ExecutionResult::done( take_canonical(values, &codes, ctx)?.into_array(), diff --git a/vortex-array/src/arrays/scalar_fn/vtable/mod.rs b/vortex-array/src/arrays/scalar_fn/vtable/mod.rs index d6eb44f9e65..c915332f2d6 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/mod.rs +++ b/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 { + array.try_into::().ok() + } +} + /// A matcher that matches a specific scalar function expression. #[derive(Debug, Default)] pub struct ExactScalarFn(PhantomData); diff --git a/vortex-array/src/canonical.rs b/vortex-array/src/canonical.rs index fe0c563ad1e..b8ef348f560 100644 --- a/vortex-array/src/canonical.rs +++ b/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 { + if !::matches(&*array) { + return None; + } + let array = match array.try_into::() { + Ok(a) => return Some(Canonical::Null(a)), + Err(array) => array, + }; + let array = match array.try_into::() { + Ok(a) => return Some(Canonical::Bool(a)), + Err(array) => array, + }; + let array = match array.try_into::() { + Ok(a) => return Some(Canonical::Primitive(a)), + Err(array) => array, + }; + let array = match array.try_into::() { + Ok(a) => return Some(Canonical::Decimal(a)), + Err(array) => array, + }; + let array = match array.try_into::() { + Ok(a) => return Some(Canonical::Struct(a)), + Err(array) => array, + }; + let array = match array.try_into::() { + Ok(a) => return Some(Canonical::List(a)), + Err(array) => array, + }; + let array = match array.try_into::() { + Ok(a) => return Some(Canonical::FixedSizeList(a)), + Err(array) => array, + }; + let array = match array.try_into::() { + Ok(a) => return Some(Canonical::VarBinView(a)), + Err(array) => array, + }; + let array = match array.try_into::() { + Ok(a) => return Some(Canonical::Variant(a)), + Err(array) => array, + }; + match array.try_into::() { + Ok(a) => Some(Canonical::Extension(a)), + Err(_) => None, + } + } +} + #[cfg(test)] mod test { use std::sync::Arc; diff --git a/vortex-array/src/columnar.rs b/vortex-array/src/columnar.rs index 8734a446bd3..37a10776f7e 100644 --- a/vortex-array/src/columnar.rs +++ b/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 { + match array.try_into::() { + Ok(constant) => Some(Columnar::Constant(constant)), + Err(array) => { + ::maybe_match(array).map(Columnar::Canonical) + } + } + } +} diff --git a/vortex-array/src/matcher.rs b/vortex-array/src/matcher.rs index 46c9d79198c..9b4273ba8cb 100644 --- a/vortex-array/src/matcher.rs +++ b/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>; } +/// 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; +} + /// 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 { + Some(array) + } +}