diff --git a/vortex-array/benches/kleene_bool.rs b/vortex-array/benches/kleene_bool.rs index babd859cec5..b23a0580a5a 100644 --- a/vortex-array/benches/kleene_bool.rs +++ b/vortex-array/benches/kleene_bool.rs @@ -17,6 +17,7 @@ use vortex_array::arrow::ArrowSession; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; +use vortex_array::optimizer::kernels::KernelSession; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_array::session::ArraySession; @@ -29,6 +30,7 @@ fn main() { static SESSION: LazyLock = LazyLock::new(|| { VortexSession::empty() .with::() + .with::() .with::() }); diff --git a/vortex-array/src/arrays/struct_/compute/cast.rs b/vortex-array/src/arrays/struct_/compute/cast.rs index 4ca1febe851..2cadb386e38 100644 --- a/vortex-array/src/arrays/struct_/compute/cast.rs +++ b/vortex-array/src/arrays/struct_/compute/cast.rs @@ -151,9 +151,9 @@ mod tests { use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::StructFields; - use crate::optimizer::kernels::ArrayKernels; use crate::optimizer::kernels::ArrayKernelsExt; use crate::optimizer::kernels::ExecuteParentFn; + use crate::optimizer::kernels::KernelSession; use crate::scalar::Scalar; use crate::scalar_fn::fns::cast::Cast; use crate::validity::Validity; @@ -214,7 +214,7 @@ mod tests { .try_new_array(source.len(), target.clone(), [source]) .unwrap(); let parent_id = cast.encoding_id(); - let session = VortexSession::empty().with::(); + let session = VortexSession::empty().with_some(KernelSession::empty()); session.kernels().register_execute_parent( parent_id, child_id, diff --git a/vortex-array/src/arrays/struct_/compute/rules.rs b/vortex-array/src/arrays/struct_/compute/rules.rs index 287732beb7d..8443c2dd25e 100644 --- a/vortex-array/src/arrays/struct_/compute/rules.rs +++ b/vortex-array/src/arrays/struct_/compute/rules.rs @@ -147,7 +147,7 @@ mod tests { use crate::dtype::StructFields; use crate::executor::VortexSessionExecute; use crate::optimizer::ArrayOptimizer; - use crate::optimizer::kernels::ArrayKernels; + use crate::optimizer::kernels::KernelSession; use crate::optimizer::kernels::ReduceParentFn; use crate::scalar::Scalar; use crate::scalar_fn::ScalarFnVTable; @@ -254,8 +254,8 @@ mod tests { ); let cast = source.cast(target).unwrap(); - let kernels = ArrayKernels::empty(); - kernels.register_reduce_parent( + let kernels = KernelSession::empty(); + kernels.kernels().register_reduce_parent( Cast.id(), Struct.id(), &[no_struct_cast_plugin as ReduceParentFn], diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 50173f95ff9..515e6f7398e 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -860,8 +860,8 @@ mod tests { use crate::VTable as _; use crate::arrays::Bool; use crate::arrays::Primitive; - use crate::optimizer::kernels::ArrayKernels; use crate::optimizer::kernels::ExecuteParentFn; + use crate::optimizer::kernels::KernelSession; use crate::optimizer::kernels::execute_parent_key; fn noop_execute_parent( @@ -875,7 +875,7 @@ mod tests { #[test] fn execution_ctx_snapshots_execute_parent_kernels_at_creation() { - let session = VortexSession::empty().with::(); + let session = VortexSession::empty().with_some(KernelSession::empty()); let key = execute_parent_key(Bool.id(), Primitive.id()); let before_registration = ExecutionCtx::new(session.clone()); diff --git a/vortex-array/src/lib.rs b/vortex-array/src/lib.rs index 151e166b1ba..16d29c69efe 100644 --- a/vortex-array/src/lib.rs +++ b/vortex-array/src/lib.rs @@ -32,6 +32,7 @@ use crate::arrow::ArrowSession; use crate::dtype::session::DTypeSession; use crate::memory::MemorySession; use crate::optimizer::kernels::ArrayKernelsExt; +use crate::optimizer::kernels::KernelSession; use crate::scalar_fn::session::ScalarFnSession; use crate::session::ArraySession; use crate::stats::session::StatsSession; @@ -88,10 +89,8 @@ pub mod flatbuffers { /// Register vortex-array's built-in session-scoped kernels into the active /// [`ArrayKernels`](crate::optimizer::kernels::ArrayKernels) registry. /// -/// If the session contains a standalone [`ArrayKernels`](crate::optimizer::kernels::ArrayKernels), -/// this registers into that registry. Otherwise, if the session contains an [`ArraySession`], this -/// registers into the [`ArraySession`]-owned registry. Sessions that use [`ArraySession::default`] -/// already receive these built-in kernels. +/// If the session contains a [`KernelSession`], this registers into its registry. Sessions that use +/// [`KernelSession::default`] already receive these built-in kernels. pub fn initialize(session: &VortexSession) { if session.kernels_opt().is_some() { arrays::initialize(session); @@ -108,6 +107,7 @@ pub fn initialize(session: &VortexSession) { pub fn array_session() -> VortexSession { VortexSession::builder() .with::() + .with::() .with::() .with::() .with::() diff --git a/vortex-array/src/optimizer/kernels.rs b/vortex-array/src/optimizer/kernels.rs index 9363488df65..4e98e896b96 100644 --- a/vortex-array/src/optimizer/kernels.rs +++ b/vortex-array/src/optimizer/kernels.rs @@ -18,10 +18,10 @@ //! Because registered functions have different signatures for each kernel kind, the registry //! maintains one storage map per function type rather than a single type-erased map. //! -//! [`ArraySession`] owns vortex-array's built-in kernel registry, -//! so sessions that install the default array encodings get their matching built-in kernels too. -//! Sessions can still install a standalone [`ArrayKernels`] registry when they need a kernel-only -//! setup or an explicit override. +//! [`KernelSession`] is the session variable that owns this registry. Its [`Default`] +//! implementation installs vortex-array's built-in parent-reduce and execute-parent kernels, so a +//! session built with [`KernelSession`] participates in the same optimizations and fused execution +//! as the built-in encodings. use std::any::Any; use std::borrow::Borrow; @@ -34,6 +34,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_panic; use vortex_session::SessionExt; use vortex_session::SessionVar; +use vortex_session::VortexSession; use vortex_session::registry::Id; use vortex_utils::aliases::DefaultHashBuilder; use vortex_utils::aliases::hash_map::HashMap; @@ -48,7 +49,6 @@ use crate::kernel::ExecuteParentKernel; use crate::matcher::Matcher; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::cast::Cast; -use crate::session::ArraySession; /// Shared hasher used to combine `(outer, child)` tuples into registry keys. static FN_HASHER: LazyLock = LazyLock::new(DefaultHashBuilder::default); @@ -291,7 +291,47 @@ pub(crate) fn execute_parent_key(parent: Id, child: Id) -> u64 { hash_fn_id(parent, child) } -impl SessionVar for ArrayKernels { +/// Session-scoped holder for the optimizer kernel registry. +/// +/// `KernelSession` is the session variable that owns an [`ArrayKernels`] registry. Its [`Default`] +/// implementation installs vortex-array's built-in parent-reduce and execute-parent kernels, +/// mirroring how [`ScalarFnSession`](crate::scalar_fn::session::ScalarFnSession) and the other +/// session variables register their built-ins. +#[derive(Clone, Debug)] +pub struct KernelSession { + kernels: ArrayKernels, +} + +impl KernelSession { + /// Create a [`KernelSession`] with an empty kernel registry. + pub fn empty() -> Self { + Self { + kernels: ArrayKernels::empty(), + } + } + + /// Returns the [`ArrayKernels`] registry held by this session. + pub fn kernels(&self) -> &ArrayKernels { + &self.kernels + } +} + +impl Default for KernelSession { + fn default() -> Self { + // `ArrayKernels::default` installs the built-in parent-reduce kernels. The execute-parent + // kernels are registered by the per-encoding `initialize` functions, which operate on a + // session. `KernelSession` clones share their registry storage, so kernels registered into + // the temporary session land in `this.kernels`. + let this = Self { + kernels: ArrayKernels::default(), + }; + let session = VortexSession::empty().with_some(this.clone()); + crate::arrays::initialize(&session); + this + } +} + +impl SessionVar for KernelSession { fn as_any(&self) -> &dyn Any { self } @@ -301,25 +341,57 @@ impl SessionVar for ArrayKernels { } } -/// Extension trait for accessing optimizer kernels from a -/// [`VortexSession`](vortex_session::VortexSession). +/// Extension trait for accessing the optimizer kernel registry from a [`VortexSession`]. pub trait ArrayKernelsExt: SessionExt { - /// Returns the active [`ArrayKernels`] registry if one is available. - /// - /// A standalone [`ArrayKernels`] variable takes precedence and is not merged with an - /// [`ArraySession`]-owned registry. Otherwise, sessions that include [`ArraySession`] use its - /// built-in kernel registry. + /// Returns the active [`ArrayKernels`] registry if the session contains a [`KernelSession`]. fn kernels_opt(&self) -> Option<&ArrayKernels> { - self.get_opt::() - .or_else(|| self.get_opt::().map(ArraySession::kernels)) + self.get_opt::().map(KernelSession::kernels) } /// Returns the active [`ArrayKernels`] registry. + /// + /// Panics if the session does not contain a [`KernelSession`]. fn kernels(&self) -> &ArrayKernels { - self.kernels_opt().unwrap_or_else(|| { - vortex_panic!("Session contains neither ArrayKernels nor ArraySession") - }) + self.kernels_opt() + .unwrap_or_else(|| vortex_panic!("Session does not contain a KernelSession")) } } impl ArrayKernelsExt for S {} + +#[cfg(test)] +mod tests { + use vortex_session::VortexSession; + + use super::ArrayKernelsExt; + use super::KernelSession; + use crate::ArrayVTable; + use crate::arrays::Bool; + use crate::scalar_fn::ScalarFnVTable; + use crate::scalar_fn::fns::binary::Binary; + + #[test] + fn kernel_session_default_registers_builtin_kernels() { + let session = VortexSession::empty().with::(); + + assert!(session.kernels().has_execute_parent(Binary.id(), Bool.id())); + } + + #[test] + fn initialize_registers_builtin_kernels_into_empty_kernel_session() { + let session = VortexSession::empty().with_some(KernelSession::empty()); + + assert!(!session.kernels().has_execute_parent(Binary.id(), Bool.id())); + + crate::initialize(&session); + + assert!(session.kernels().has_execute_parent(Binary.id(), Bool.id())); + } + + #[test] + fn kernels_opt_is_none_without_kernel_session() { + let session = VortexSession::empty(); + + assert!(session.kernels_opt().is_none()); + } +} diff --git a/vortex-array/src/optimizer/mod.rs b/vortex-array/src/optimizer/mod.rs index 4bbec8b7c32..77a337098be 100644 --- a/vortex-array/src/optimizer/mod.rs +++ b/vortex-array/src/optimizer/mod.rs @@ -39,9 +39,9 @@ pub trait ArrayOptimizer { /// [`kernels::ArrayKernels`] registry on `session`, if any. /// /// Session kernels are checked for each `(parent_encoding_id, child_encoding_id)` pair before - /// the child's static `PARENT_RULES`. A standalone [`kernels::ArrayKernels`] registry takes - /// precedence over an [`ArraySession`](crate::session::ArraySession)-owned registry. If - /// `session` contains neither registry, this behaves like [`Self::optimize`]. + /// the child's static `PARENT_RULES`. The registry comes from the [`kernels::KernelSession`] on + /// `session`, if any. If `session` does not contain a [`kernels::KernelSession`], this behaves + /// like [`Self::optimize`]. fn optimize_ctx(&self, session: &VortexSession) -> VortexResult; /// Optimize the entire array tree recursively (root and all descendants). diff --git a/vortex-array/src/session/mod.rs b/vortex-array/src/session/mod.rs index eec38d4c15d..6e143eaa4dd 100644 --- a/vortex-array/src/session/mod.rs +++ b/vortex-array/src/session/mod.rs @@ -8,7 +8,6 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_session::SessionExt; use vortex_session::SessionVar; -use vortex_session::VortexSession; use vortex_session::registry::Registry; use crate::ArrayRef; @@ -30,7 +29,6 @@ use crate::arrays::Struct; use crate::arrays::VarBin; use crate::arrays::VarBinView; use crate::arrays::Variant; -use crate::optimizer::kernels::ArrayKernels; pub type ArrayRegistry = Registry; @@ -38,15 +36,12 @@ pub type ArrayRegistry = Registry; pub struct ArraySession { /// The set of registered array encodings. registry: ArrayRegistry, - /// The set of built-in array kernels for the registered encodings. - kernels: ArrayKernels, } impl ArraySession { pub fn empty() -> ArraySession { Self { registry: ArrayRegistry::default(), - kernels: ArrayKernels::empty(), } } @@ -54,14 +49,6 @@ impl ArraySession { &self.registry } - /// Returns this array session's built-in kernel registry. - /// - /// This registry is used when the surrounding [`VortexSession`] does not contain a standalone - /// [`ArrayKernels`] session variable. A standalone [`ArrayKernels`] shadows this registry. - pub fn kernels(&self) -> &ArrayKernels { - &self.kernels - } - /// Register a new array encoding, replacing any existing encoding with the same ID. pub fn register(&self, plugin: P) { self.registry @@ -73,7 +60,6 @@ impl Default for ArraySession { fn default() -> Self { let this = ArraySession { registry: ArrayRegistry::default(), - kernels: ArrayKernels::default(), }; // Register the canonical encodings. @@ -96,9 +82,6 @@ impl Default for ArraySession { this.register(Masked); this.register(VarBin); - let session = VortexSession::empty().with_some(this.kernels.clone()); - crate::arrays::initialize(&session); - this } } @@ -141,43 +124,20 @@ mod tests { use crate::ArrayVTable; use crate::arrays::Bool; - use crate::optimizer::kernels::ArrayKernels; - use crate::optimizer::kernels::ArrayKernelsExt; - use crate::scalar_fn::ScalarFnVTable; - use crate::scalar_fn::fns::binary::Binary; use crate::session::ArraySession; + use crate::session::ArraySessionExt; #[test] - fn array_session_default_registers_builtin_kernels() { + fn array_session_default_registers_encodings() { let session = VortexSession::empty().with::(); - assert!(session.get_opt::().is_none()); - assert!(session.kernels().has_execute_parent(Binary.id(), Bool.id())); + assert!(session.arrays().registry().find(&Bool.id()).is_some()); } #[test] - fn initialize_registers_builtin_kernels_into_empty_array_session() { + fn empty_array_session_registers_no_encodings() { let session = VortexSession::empty().with_some(ArraySession::empty()); - assert!(!session.kernels().has_execute_parent(Binary.id(), Bool.id())); - - crate::initialize(&session); - - assert!(session.kernels().has_execute_parent(Binary.id(), Bool.id())); - } - - #[test] - fn standalone_array_kernels_shadow_array_session_kernels() { - let session = VortexSession::empty() - .with::() - .with::(); - - assert!( - session - .get::() - .kernels() - .has_execute_parent(Binary.id(), Bool.id()) - ); - assert!(!session.kernels().has_execute_parent(Binary.id(), Bool.id())); + assert!(session.arrays().registry().find(&Bool.id()).is_none()); } } diff --git a/vortex-ipc/src/lib.rs b/vortex-ipc/src/lib.rs index 4e7eaba998c..2a198387b6f 100644 --- a/vortex-ipc/src/lib.rs +++ b/vortex-ipc/src/lib.rs @@ -20,6 +20,7 @@ mod test { use std::sync::LazyLock; use vortex_array::dtype::session::DTypeSession; + use vortex_array::optimizer::kernels::KernelSession; use vortex_array::session::ArraySession; use vortex_session::VortexSession; @@ -27,6 +28,7 @@ mod test { VortexSession::builder() .with::() .with::() + .with::() .build() }); } diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index 8b4afceb1fb..ff6bcd39572 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -14,6 +14,7 @@ use vortex_array::dtype::session::DTypeSession; // pulled back out into a vortex_expr crate. pub use vortex_array::expr; use vortex_array::memory::MemorySession; +use vortex_array::optimizer::kernels::KernelSession; pub use vortex_array::scalar_fn; use vortex_array::scalar_fn::session::ScalarFnSession; use vortex_array::session::ArraySession; @@ -167,6 +168,7 @@ impl VortexSessionDefault for VortexSession { let session = VortexSession::empty() .with::() .with::() + .with::() .with::() .with::() .with::()