Move optimizer kernels into a dedicated KernelSession

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<VortexSession> = LazyLock::new(|| {
     VortexSession::empty()
         .with::<ArraySession>()
+        .with::<KernelSession>()
         .with::<ArrowSession>()
 });
 

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::<ArrayKernels>();
+        let session = VortexSession::empty().with_some(KernelSession::empty());
         session.kernels().register_execute_parent(
             parent_id,
             child_id,

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],

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::<ArrayKernels>();
+        let session = VortexSession::empty().with_some(KernelSession::empty());
         let key = execute_parent_key(Bool.id(), Primitive.id());
 
         let before_registration = ExecutionCtx::new(session.clone());

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::<ArraySession>()
+        .with::<KernelSession>()
         .with::<DTypeSession>()
         .with::<ScalarFnSession>()
         .with::<StatsSession>()

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<DefaultHashBuilder> = 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::<ArrayKernels>()
-            .or_else(|| self.get_opt::<ArraySession>().map(ArraySession::kernels))
+        self.get_opt::<KernelSession>().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<S: SessionExt> 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::<KernelSession>();
+
+        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());
+    }
+}

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<ArrayRef>;
 
     /// Optimize the entire array tree recursively (root and all descendants).

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,38 +29,26 @@ 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<ArrayPluginRef>;
 
 #[derive(Clone, Debug)]
 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(),
         }
     }
 
     pub fn registry(&self) -> &ArrayRegistry {
         &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<P: ArrayPlugin>(&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::<ArraySession>();
 
-        assert!(session.get_opt::<ArrayKernels>().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::<ArraySession>()
-            .with::<ArrayKernels>();
-
-        assert!(
-            session
-                .get::<ArraySession>()
-                .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());
     }
 }

vortex-ipc/src/lib.rs

@@ -20,13 +20,15 @@ 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;
 
     pub(crate) static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
         VortexSession::builder()
             .with::<DTypeSession>()
             .with::<ArraySession>()
+            .with::<KernelSession>()
             .build()
     });
 }