diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs index ef25c559b99..7e220ee9e02 100644 --- a/encodings/fsst/src/array.rs +++ b/encodings/fsst/src/array.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::any::TypeId; use std::fmt::Debug; use std::fmt::Display; use std::fmt::Formatter; @@ -46,7 +47,9 @@ use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityVTable; +use vortex_array::vtable::VarBinExportable; use vortex_array::vtable::child_to_validity; +use vortex_array::vtable::has_capability; use vortex_array::vtable::validity_to_child; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -114,6 +117,8 @@ impl ArrayEq for FSSTData { } } +impl VarBinExportable for FSST {} + impl VTable for FSST { type TypedArrayData = FSSTData; type OperationsVTable = Self; @@ -318,6 +323,10 @@ impl VTable for FSST { canonicalize_fsst(array.as_view(), ctx).map(ExecutionResult::done) } + fn has_capability(&self, capability: TypeId) -> bool { + has_capability::(self, capability) + } + fn append_to_builder( array: ArrayView<'_, Self>, builder: &mut dyn ArrayBuilder, diff --git a/encodings/onpair/src/array.rs b/encodings/onpair/src/array.rs index e57bed91827..0e636138eec 100644 --- a/encodings/onpair/src/array.rs +++ b/encodings/onpair/src/array.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::any::TypeId; use std::fmt::Debug; use std::fmt::Display; use std::fmt::Formatter; @@ -40,7 +41,9 @@ use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityVTable; +use vortex_array::vtable::VarBinExportable; use vortex_array::vtable::child_to_validity; +use vortex_array::vtable::has_capability; use vortex_array::vtable::validity_to_child; use vortex_buffer::Buffer; use vortex_buffer::ByteBuffer; @@ -415,6 +418,8 @@ fn validate_parts( Ok(()) } +impl VarBinExportable for OnPair {} + impl VTable for OnPair { type TypedArrayData = OnPairData; type OperationsVTable = Self; @@ -586,6 +591,10 @@ impl VTable for OnPair { canonicalize_onpair(array.as_view(), ctx).map(ExecutionResult::done) } + fn has_capability(&self, capability: TypeId) -> bool { + has_capability::(self, capability) + } + fn append_to_builder( array: ArrayView<'_, Self>, builder: &mut dyn ArrayBuilder, diff --git a/encodings/sparse/src/lib.rs b/encodings/sparse/src/lib.rs index f3a488ea9ba..b500010c77c 100644 --- a/encodings/sparse/src/lib.rs +++ b/encodings/sparse/src/lib.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::any::TypeId; use std::fmt::Debug; use std::fmt::Display; use std::fmt::Formatter; @@ -47,6 +48,8 @@ use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityVTable; +use vortex_array::vtable::VarBinExportable; +use vortex_array::vtable::has_capability; use vortex_buffer::Buffer; use vortex_buffer::ByteBufferMut; use vortex_error::VortexExpect as _; @@ -187,6 +190,8 @@ impl ArrayEq for SparseData { } } +impl VarBinExportable for Sparse {} + impl VTable for Sparse { type TypedArrayData = SparseData; @@ -359,6 +364,10 @@ impl VTable for Sparse { execute_sparse(parts, ctx).map(ExecutionResult::done) } + fn has_capability(&self, capability: TypeId) -> bool { + has_capability::(self, capability) + } + fn append_to_builder( array: ArrayView<'_, Self>, builder: &mut dyn ArrayBuilder, diff --git a/encodings/zstd/src/array.rs b/encodings/zstd/src/array.rs index 0d831e58775..42fed71d182 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::any::TypeId; use std::fmt::Debug; use std::fmt::Display; use std::fmt::Formatter; @@ -45,7 +46,9 @@ use vortex_array::validity::Validity; use vortex_array::vtable::OperationsVTable; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityVTable; +use vortex_array::vtable::VarBinExportable; use vortex_array::vtable::child_to_validity; +use vortex_array::vtable::has_capability; use vortex_array::vtable::validity_to_child; use vortex_buffer::Alignment; use vortex_buffer::Buffer; @@ -136,6 +139,8 @@ impl ArrayEq for ZstdData { } } +impl VarBinExportable for Zstd {} + impl VTable for Zstd { type TypedArrayData = ZstdData; @@ -279,6 +284,10 @@ impl VTable for Zstd { .map(ExecutionResult::done) } + fn has_capability(&self, capability: TypeId) -> bool { + has_capability::(self, capability) + } + fn append_to_builder( array: ArrayView<'_, Self>, builder: &mut dyn ArrayBuilder, diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index 6674c4d9429..7541b146d5a 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::any::TypeId; use std::any::type_name; use std::fmt::Debug; use std::fmt::Formatter; @@ -385,19 +386,19 @@ impl ArrayRef { /// Does the array match the given matcher. #[inline] - pub fn is(&self) -> bool { + pub fn is(&self) -> bool { M::matches(self) } /// Returns the array downcast by the given matcher. #[inline] - pub fn as_(&self) -> M::Match<'_> { + pub fn as_(&self) -> M::Match<'_> { self.as_opt::().vortex_expect("Failed to downcast") } /// Returns the array downcast by the given matcher. #[inline] - pub fn as_opt(&self) -> Option> { + pub fn as_opt(&self) -> Option> { M::try_match(self) } @@ -443,6 +444,13 @@ impl ArrayRef { self.is::() } + /// Whether this array's encoding implements the capability trait `C`. + /// + /// `false` unless the encoding reports `C` from [`VTable::has_capability`]. + pub fn has_capability(&self) -> bool { + self.0.data.has_capability(TypeId::of::()) + } + /// 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 8b9b2812bfa..0e467a54ca1 100644 --- a/vortex-array/src/array/mod.rs +++ b/vortex-array/src/array/mod.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::any::Any; +use std::any::TypeId; use std::fmt::Debug; use std::fmt::Formatter; use std::hash::Hasher; @@ -135,6 +136,10 @@ pub(crate) trait DynArrayData: 'static + private::Sealed + Send + Sync + Debug { ctx: &mut ExecutionCtx, ) -> VortexResult<()>; + /// Whether this array's encoding implements the given capability trait. See + /// [`VTable::has_capability`]. + fn has_capability(&self, capability: TypeId) -> bool; + // --- Visitor methods (formerly in ArrayVisitor) --- /// Returns the buffers of the array. @@ -307,6 +312,10 @@ impl DynArrayData for ArrayData { Ok(()) } + fn has_capability(&self, capability: TypeId) -> bool { + V::has_capability(&self.vtable, capability) + } + fn buffers(&self, this: &ArrayRef) -> Vec { let view = unsafe { ArrayView::new_unchecked(this, &self.data) }; (0..V::nbuffers(view)) diff --git a/vortex-array/src/array/vtable/capability.rs b/vortex-array/src/array/vtable/capability.rs new file mode 100644 index 00000000000..82f4efbcb42 --- /dev/null +++ b/vortex-array/src/array/vtable/capability.rs @@ -0,0 +1,38 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::any::TypeId; + +/// Whether `capability` names the capability trait `C`. +/// +/// Called once per capability from [`VTable::has_capability`](super::VTable::has_capability). +/// Passing the vtable as `&C` is what proves it implements `C`, so a vtable cannot claim a +/// capability it does not have. +pub fn has_capability(_vtable: &C, capability: TypeId) -> bool { + capability == TypeId::of::() +} + +#[cfg(test)] +mod tests { + use super::*; + + trait Greet {} + + trait Absent {} + + struct Encoding; + + impl Greet for Encoding {} + + #[test] + fn has_capability_matches_only_the_requested_trait() { + assert!(has_capability::( + &Encoding, + TypeId::of::() + )); + assert!(!has_capability::( + &Encoding, + TypeId::of::() + )); + } +} diff --git a/vortex-array/src/array/vtable/mod.rs b/vortex-array/src/array/vtable/mod.rs index 3078a6ae9f5..3e2398f0d54 100644 --- a/vortex-array/src/array/vtable/mod.rs +++ b/vortex-array/src/array/vtable/mod.rs @@ -11,16 +11,21 @@ //! into these traits. Implementations should focus on encoding-specific work and uphold the //! documented postconditions. +mod capability; mod operations; mod validity; +mod varbin_exportable; +use std::any::TypeId; use std::fmt::Debug; use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hasher; +pub use capability::*; pub use operations::*; pub use validity::*; +pub use varbin_exportable::*; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -140,10 +145,28 @@ pub trait VTable: 'static + Clone + Sized + Send + Sync + Debug { session: &VortexSession, ) -> VortexResult>; + /// Whether this encoding implements the capability trait identified by `capability`. + /// + /// Capabilities are optional interfaces, queried through + /// [`ArrayRef::has_capability`](crate::ArrayRef::has_capability). Report each one with + /// [`has_capability`]: + /// + /// ```ignore + /// fn has_capability(&self, capability: TypeId) -> bool { + /// has_capability::(self, capability) + /// || has_capability::(self, capability) + /// } + /// ``` + fn has_capability(&self, capability: TypeId) -> bool { + _ = capability; + false + } + /// Writes the array's logical values into a canonical builder. /// /// The default implementation executes the full array to [`Canonical`] and appends that result. - /// Encodings may override this to avoid materializing an intermediate canonical array. + /// Encodings may override this to avoid materializing an intermediate canonical array. An + /// override targeting a `VarBinBuilder` should also claim [`VarBinExportable`]. fn append_to_builder( array: ArrayView<'_, Self>, builder: &mut dyn ArrayBuilder, diff --git a/vortex-array/src/array/vtable/varbin_exportable.rs b/vortex-array/src/array/vtable/varbin_exportable.rs new file mode 100644 index 00000000000..b1a15d88b1f --- /dev/null +++ b/vortex-array/src/array/vtable/varbin_exportable.rs @@ -0,0 +1,53 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use crate::ArrayRef; +use crate::matcher::Matcher; + +/// Capability for encodings whose [`append_to_builder`](super::VTable::append_to_builder) writes +/// `Utf8`/`Binary` values straight into a [`VarBinBuilder`](crate::builders::VarBinBuilder). +/// +/// Callers filling a `VarBinBuilder` should `execute_until::`: executing past +/// these encodings reaches a canonical `VarBinView` that the builder has to re-lay out. +/// +/// Claim it only where stopping is the better trade. `Dict` does not: it appends canonical values, +/// but executing through it lets kernels such as FSST's `Dict` parent kernel decode through the +/// dictionary instead. +pub trait VarBinExportable: 'static + Send + Sync {} + +/// Matches every array whose encoding offers [`VarBinExportable`]. +impl Matcher for dyn VarBinExportable { + type Match<'a> = &'a ArrayRef; + + #[inline] + fn try_match(array: &ArrayRef) -> Option> { + array + .has_capability::() + .then_some(array) + } +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use super::*; + use crate::IntoArray; + use crate::arrays::DictArray; + use crate::arrays::PrimitiveArray; + use crate::arrays::VarBinViewArray; + + #[test] + fn encodings_report_the_capability_through_their_vtable() -> VortexResult<()> { + let values = VarBinViewArray::from_iter_str(["a", "b"]).into_array(); + assert!(values.has_capability::()); + assert!(values.is::()); + + // Dict gathers canonical values into the builder but deliberately does not claim the + // capability, so execution continues through it. + let codes = PrimitiveArray::from_iter([0u8, 1, 0]).into_array(); + let dict = DictArray::try_new(codes, values)?.into_array(); + assert!(!dict.has_capability::()); + Ok(()) + } +} diff --git a/vortex-array/src/arrays/chunked/vtable/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index 59e89f031e1..6af976fa490 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::any::TypeId; use std::hash::Hasher; use itertools::Itertools; @@ -41,6 +42,8 @@ use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; use crate::serde::ArrayChildren; +use crate::vtable::VarBinExportable; +use crate::vtable::has_capability; mod canonical; mod operations; mod validity; @@ -66,6 +69,8 @@ impl ArrayEq for ChunkedData { } } +impl VarBinExportable for Chunked {} + impl VTable for Chunked { type TypedArrayData = ChunkedData; @@ -234,6 +239,10 @@ impl VTable for Chunked { .with_slots(slots)) } + fn has_capability(&self, capability: TypeId) -> bool { + has_capability::(self, capability) + } + fn append_to_builder( array: ArrayView<'_, Self>, builder: &mut dyn ArrayBuilder, diff --git a/vortex-array/src/arrays/constant/vtable/mod.rs b/vortex-array/src/arrays/constant/vtable/mod.rs index 2e4c982a8ea..64eb6ca5d5a 100644 --- a/vortex-array/src/arrays/constant/vtable/mod.rs +++ b/vortex-array/src/arrays/constant/vtable/mod.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::any::TypeId; use std::fmt::Debug; use std::hash::Hash; use std::hash::Hasher; @@ -45,6 +46,8 @@ use crate::scalar::DecimalValue; use crate::scalar::Scalar; use crate::scalar::ScalarValue; use crate::serde::ArrayChildren; +use crate::vtable::VarBinExportable; +use crate::vtable::has_capability; pub(crate) mod canonical; mod operations; mod validity; @@ -67,6 +70,8 @@ impl ArrayEq for ConstantData { } } +impl VarBinExportable for Constant {} + impl VTable for Constant { type TypedArrayData = ConstantData; @@ -178,6 +183,10 @@ impl VTable for Constant { )?)) } + fn has_capability(&self, capability: TypeId) -> bool { + has_capability::(self, capability) + } + fn append_to_builder( array: ArrayView<'_, Self>, builder: &mut dyn ArrayBuilder, diff --git a/vortex-array/src/arrays/varbin/vtable/mod.rs b/vortex-array/src/arrays/varbin/vtable/mod.rs index 397b0f8084b..dd56e40eaec 100644 --- a/vortex-array/src/arrays/varbin/vtable/mod.rs +++ b/vortex-array/src/arrays/varbin/vtable/mod.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::any::TypeId; use std::hash::Hasher; use prost::Message; @@ -47,6 +48,8 @@ use crate::EqMode; use crate::arrays::varbin::compute::rules::PARENT_RULES; use crate::hash::ArrayEq; use crate::hash::ArrayHash; +use crate::vtable::VarBinExportable; +use crate::vtable::has_capability; /// A [`VarBin`]-encoded Vortex array. pub type VarBinArray = Array; @@ -73,6 +76,8 @@ impl ArrayEq for VarBinData { } } +impl VarBinExportable for VarBin {} + impl VTable for VarBin { type TypedArrayData = VarBinData; @@ -208,6 +213,10 @@ impl VTable for VarBin { PARENT_RULES.evaluate(array, parent, child_idx) } + fn has_capability(&self, capability: TypeId) -> bool { + has_capability::(self, capability) + } + fn append_to_builder( array: ArrayView<'_, Self>, builder: &mut dyn ArrayBuilder, diff --git a/vortex-array/src/arrays/varbinview/vtable/mod.rs b/vortex-array/src/arrays/varbinview/vtable/mod.rs index 33e9ce289a5..37d1847c221 100644 --- a/vortex-array/src/arrays/varbinview/vtable/mod.rs +++ b/vortex-array/src/arrays/varbinview/vtable/mod.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::any::TypeId; use std::hash::Hasher; use std::mem::size_of; use std::sync::Arc; @@ -36,6 +37,8 @@ use crate::hash::ArrayHash; use crate::match_each_varbin_builder; use crate::serde::ArrayChildren; use crate::validity::Validity; +use crate::vtable::VarBinExportable; +use crate::vtable::has_capability; mod kernel; mod operations; mod validity; @@ -70,6 +73,8 @@ impl ArrayEq for VarBinViewData { } } +impl VarBinExportable for VarBinView {} + impl VTable for VarBinView { type TypedArrayData = VarBinViewData; @@ -241,6 +246,10 @@ impl VTable for VarBinView { PARENT_RULES.evaluate(array, parent, child_idx) } + fn has_capability(&self, capability: TypeId) -> bool { + has_capability::(self, capability) + } + fn append_to_builder( array: ArrayView<'_, Self>, builder: &mut dyn ArrayBuilder, diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 4e87ff5e7dd..2d65ecc6efb 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -165,7 +165,10 @@ impl ArrayRef { /// parent rewrite would observe inconsistent state and could discard accumulated builder /// data. #[allow(clippy::cognitive_complexity)] - pub fn execute_until(self, ctx: &mut ExecutionCtx) -> VortexResult { + pub fn execute_until( + self, + ctx: &mut ExecutionCtx, + ) -> VortexResult { let mut current_array = self; let mut current_builder: Option> = None; let mut stack: Vec = Vec::new(); @@ -774,7 +777,7 @@ impl ExecutionResult { /// 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. - pub fn execute_slot(array: impl IntoArray, slot_idx: usize) -> Self { + pub fn execute_slot(array: impl IntoArray, slot_idx: usize) -> Self { let array = array.into_array(); Self { array, diff --git a/vortex-array/src/matcher.rs b/vortex-array/src/matcher.rs index c3f9e542f13..90da0d68c9c 100644 --- a/vortex-array/src/matcher.rs +++ b/vortex-array/src/matcher.rs @@ -4,6 +4,9 @@ use crate::ArrayRef; /// Trait for matching array types. +/// +/// Implemented for encoding vtable types, for logical matchers like `AnyCanonical`, and for +/// capability traits such as `dyn VarBinExportable`. pub trait Matcher { type Match<'a>; diff --git a/vortex-array/src/test_harness/trace/mod.rs b/vortex-array/src/test_harness/trace/mod.rs index 0890f85f067..a34b75ade4f 100644 --- a/vortex-array/src/test_harness/trace/mod.rs +++ b/vortex-array/src/test_harness/trace/mod.rs @@ -561,7 +561,7 @@ fn record_parent_reduce_applied( }); } -pub(crate) fn record_execute_until_start(root: &ArrayRef) { +pub(crate) fn record_execute_until_start(root: &ArrayRef) { record(TraceEvent::ExecuteUntilStart { target: short_type_name::(), root: ArraySummary::new(root), @@ -816,7 +816,7 @@ fn adapter_field<'a>(label: &'a str, field: &str) -> Option<&'a str> { Some(&rest[..end]) } -fn short_type_name() -> String { +fn short_type_name() -> String { std::any::type_name::() .rsplit("::") .next()