fixes

Signed-off-by: Robert Kruszewski <github@robertk.io>

Author: robert3005

Cargo.lock

@@ -89,7 +89,7 @@ version = "0.1.0"
 aho-corasick = "1.1.3"
 anyhow = "1.0.97"
 arbitrary = "1.3.2"
-arc-swap = "1.8"
+arc-swap = "1.9"
 arcref = "0.2.0"
 arrow-arith = "58"
 arrow-array = "58"

Cargo.toml

@@ -88,7 +88,6 @@ mod tests {
     use crate::IntoArray;
     #[expect(deprecated)]
     use crate::ToCanonical as _;
-    use crate::VortexSessionExecute;
     use crate::arrays::Constant;
     use crate::arrays::ConstantArray;
     use crate::arrays::Extension;

vortex-array/src/arrays/extension/compute/rules.rs

@@ -0,0 +1,87 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Session state for pluggable parent-reduce rules.
+//!
+//! [`ArrayKernels`] wraps an [`FnRegistry`] keyed by `(parent_encoding_id, child_encoding_id)`
+//! and is consulted by the optimizer during execution, before the child encoding's static
+//! `PARENT_RULES` are tried. Entries are typed as [`ReduceParentFn`](super::ReduceParentFn).
+//!
+//! The registry is empty by default. Downstream crates register `ReduceParentFn` values to add
+//! new parent-reduce rules or override ones that the child encoding would otherwise run from its
+//! static `PARENT_RULES`.
+
+use std::hash::BuildHasher;
+use std::sync::Arc;
+use std::sync::LazyLock;
+
+use vortex_error::VortexResult;
+use vortex_session::Ref;
+use vortex_session::SessionExt;
+use vortex_session::registry::FnRegistry;
+use vortex_session::registry::Id;
+use vortex_utils::aliases::DefaultHashBuilder;
+
+use crate::ArrayRef;
+
+static FN_HASHER: LazyLock<DefaultHashBuilder> = LazyLock::new(DefaultHashBuilder::default);
+
+/// Pluggable parent-reduce function signature used by [`ArrayKernels`].
+///
+/// A function of this type rewrites the parent array that holds `child` at `child_idx`, given
+/// the child itself and its parent. Returns `Ok(None)` when the function doesn't apply.
+///
+/// Registered under `(parent_encoding_id, child_encoding_id)`; callers downcast the erased
+/// `child`/`parent` to their expected types before applying logic.
+pub type ReduceParentFn =
+    fn(child: &ArrayRef, parent: &ArrayRef, child_idx: usize) -> VortexResult<Option<ArrayRef>>;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+enum FnKind {
+    Reduce,
+    ReduceParent,
+    ExecuteParent,
+    Execute,
+}
+
+/// Session state for pluggable parent-reduce dispatch keyed by `(parent_id, child_id)`.
+#[derive(Debug, Default)]
+pub struct ArrayKernels {
+    registry: FnRegistry,
+}
+
+impl ArrayKernels {
+    /// Create an empty session with no rules registered.
+    pub fn empty() -> Self {
+        Self::default()
+    }
+
+    pub fn register_reduce_parent(&self, outer: Id, child: Id, f: ReduceParentFn) {
+        self.registry
+            .register(self.hash_fn_ids(outer, child, FnKind::ReduceParent), f)
+    }
+
+    pub fn find_reduce_parent(&self, outer: Id, child: Id) -> Option<Arc<ReduceParentFn>> {
+        self.registry
+            .find(self.hash_fn_ids(outer, child, FnKind::ReduceParent))
+    }
+
+    pub fn contains_reduce_parent(&self, outer: Id, child: Id) -> bool {
+        self.registry
+            .contains(self.hash_fn_ids(outer, child, FnKind::ReduceParent))
+    }
+
+    fn hash_fn_ids(&self, outer: Id, child: Id, fn_kind: FnKind) -> u64 {
+        FN_HASHER.hash_one((outer, child, fn_kind))
+    }
+}
+
+/// Extension trait for accessing the optimizer registry from a Vortex session.
+pub trait ArrayKernelsExt: SessionExt {
+    /// Returns the optimizer session variable.
+    fn kernels(&self) -> Ref<'_, ArrayKernels> {
+        self.get::<ArrayKernels>()
+    }
+}
+
+impl<S: SessionExt> ArrayKernelsExt for S {}

vortex-array/src/optimizer/kernels.rs

@@ -14,7 +14,7 @@
 //!   `ArrayBuiltins::cast` and `ArrayRef::slice` that build wrapped expressions and need them
 //!   normalized inline.
 //! * [`ArrayOptimizer::optimize_ctx`] — runs the static rules and additionally consults the
-//!   session's [`OptimizerSession`] registry keyed by `(parent_encoding_id, child_encoding_id)`
+//!   session's [`ArrayKernels`] registry keyed by `(parent_encoding_id, child_encoding_id)`
 //!   before each `reduce_parent` step. The execute loop calls this entry point so plugin-
 //!   registered parent-reduce rules fire during execution.
 
@@ -26,38 +26,29 @@ use vortex_session::SessionExt;
 use vortex_session::VortexSession;
 
 use crate::ArrayRef;
-use crate::optimizer::session::OptimizerSession;
+use crate::optimizer::kernels::ArrayKernels;
+use crate::optimizer::kernels::ReduceParentFn;
 
+pub mod kernels;
 pub mod rules;
-pub mod session;
-
-/// Pluggable parent-reduce function signature used by [`OptimizerSession`].
-///
-/// A function of this type rewrites the parent array that holds `child` at `child_idx`, given
-/// the child itself and its parent. Returns `Ok(None)` when the function doesn't apply.
-///
-/// Registered under `(parent_encoding_id, child_encoding_id)`; callers downcast the erased
-/// `child`/`parent` to their expected types before applying logic.
-pub type ReduceParentFn =
-    fn(child: &ArrayRef, parent: &ArrayRef, child_idx: usize) -> VortexResult<Option<ArrayRef>>;
 
 /// Extension trait for optimizing array trees using reduce/reduce_parent rules.
 pub trait ArrayOptimizer {
     /// Optimize the root array node by running reduce and reduce_parent rules to fixpoint.
     ///
     /// Uses only the child encoding's static `PARENT_RULES`. Use [`Self::optimize_ctx`] from
-    /// inside the execute loop to also consult the session-scoped [`OptimizerSession`] registry.
+    /// inside the execute loop to also consult the session-scoped [`ArrayKernels`] registry.
     fn optimize(&self) -> VortexResult<ArrayRef>;
 
-    /// Like [`Self::optimize`], but additionally consults the [`OptimizerSession`] registered on
+    /// Like [`Self::optimize`], but additionally consults the [`ArrayKernels`] registered on
     /// `session` for each `(parent_encoding_id, child_encoding_id)` pair before the static
-    /// vtable rules. If `session` does not have an [`OptimizerSession`] registered, falls
+    /// vtable rules. If `session` does not have an [`ArrayKernels`] registered, falls
     /// through to the static rules.
     fn optimize_ctx(&self, session: &VortexSession) -> VortexResult<ArrayRef>;
 
     /// Optimize the entire array tree recursively (root and all descendants).
     ///
-    /// Consults the [`OptimizerSession`] registered on `session` for each parent/child pair
+    /// Consults the [`ArrayKernels`] registered on `session` for each parent/child pair
     /// encountered during the recursive walk, so plugin-registered rules apply throughout the
     /// tree. Requires a [`VortexSession`] unconditionally so the registry is always honored
     /// when a recursive optimization is requested.
@@ -80,18 +71,17 @@ impl ArrayOptimizer for ArrayRef {
 
 /// Resolve a pluggable [`ReduceParentFn`] for `(parent, child)` from `session`.
 ///
-/// Returns `None` when no [`OptimizerSession`] is registered, or no function is registered under
+/// Returns `None` when no [`ArrayKernels`] is registered, or no function is registered under
 /// `(parent.encoding_id(), child.encoding_id())`. The returned `Arc` is owned so the caller can
 /// drop the session borrow before invoking it.
 fn plugin_reduce_parent(
     session: &VortexSession,
     parent: &ArrayRef,
     child: &ArrayRef,
 ) -> Option<Arc<ReduceParentFn>> {
-    session.get_opt::<OptimizerSession>().and_then(|s| {
-        s.registry()
-            .find::<ReduceParentFn>(parent.encoding_id(), child.encoding_id())
-    })
+    session
+        .get_opt::<ArrayKernels>()
+        .and_then(|s| s.find_reduce_parent(parent.encoding_id(), child.encoding_id()))
 }
 
 fn try_optimize(

vortex-array/src/optimizer/mod.rs

@@ -20,7 +20,7 @@ all-features = true
 workspace = true
 
 [dependencies]
-arcref = { workspace = true }
+arc-swap = { workspace = true }
 dashmap = { workspace = true }
 lasso = { workspace = true }
 parking_lot = { workspace = true }

vortex-array/src/optimizer/session.rs

@@ -10,16 +10,19 @@ use std::fmt;
 use std::fmt::Debug;
 use std::fmt::Display;
 use std::fmt::Formatter;
+use std::hash::Hash;
 use std::ops::Deref;
 use std::sync::Arc;
 use std::sync::LazyLock;
 use std::sync::OnceLock;
 
+use arc_swap::ArcSwap;
 use lasso::Spur;
 use lasso::ThreadedRodeo;
 use parking_lot::RwLock;
 use vortex_error::VortexExpect;
 use vortex_utils::aliases::dash_map::DashMap;
+use vortex_utils::aliases::hash_map::HashMap;
 
 /// Global string interner for [`Id`] values.
 static INTERNER: LazyLock<ThreadedRodeo> = LazyLock::new(ThreadedRodeo::new);
@@ -299,8 +302,8 @@ impl<T: Clone> Context<T> {
 /// optimizer's parent-reduce registry keys by `(parent_encoding_id, child_encoding_id)` so that
 /// downstream crates can override the rule that would normally run from the child encoding's
 /// static `PARENT_RULES` set.
-#[derive(Clone, Debug, Default)]
-pub struct FnRegistry(Arc<DashMap<(Id, Id), Arc<dyn Any + Send + Sync>>>);
+#[derive(Debug, Default)]
+pub struct FnRegistry(ArcSwap<HashMap<u64, Arc<dyn Any + Send + Sync>>>);
 
 impl FnRegistry {
     /// Create a new, empty registry.
@@ -309,26 +312,35 @@ impl FnRegistry {
     }
 
     /// Register a function under `(outer, inner)`, replacing any existing entry.
-    pub fn register<F: Any + Send + Sync>(&self, outer: Id, inner: Id, f: F) {
-        self.0.insert((outer, inner), Arc::new(f));
+    pub fn register<F: Any + Send + Sync>(&self, id: u64, f: F) {
+        let registry = self.0.load();
+        let mut owned_registry = registry.as_ref().clone();
+        owned_registry.insert(id, Arc::new(f));
+        self.0.store(Arc::new(owned_registry));
     }
 
     /// Look up a function registered under `(outer, inner)`, downcasting to `F`.
     ///
     /// Returns `None` if no function is registered, or if the registered value is not of type `F`.
-    pub fn find<F: Any + Send + Sync>(&self, outer: Id, inner: Id) -> Option<Arc<F>> {
-        let entry = self.0.get(&(outer, inner))?;
-        Arc::clone(entry.value()).downcast::<F>().ok()
+    pub fn find<F: Any + Send + Sync>(&self, id: u64) -> Option<Arc<F>> {
+        let map = self.0.load();
+        let entry = map.get(&id)?;
+        Arc::clone(entry).downcast::<F>().ok()
     }
 
     /// Return `true` if any function is registered under `(outer, inner)`.
-    pub fn contains(&self, outer: Id, inner: Id) -> bool {
-        self.0.contains_key(&(outer, inner))
+    pub fn contains(&self, id: u64) -> bool {
+        let map = self.0.load();
+        map.contains_key(&id)
     }
 }
 
 #[cfg(test)]
 mod fn_registry_tests {
+    use std::hash::BuildHasher;
+
+    use vortex_utils::aliases::DefaultHashBuilder;
+
     use super::FnRegistry;
     use super::Id;
 
@@ -343,12 +355,13 @@ mod fn_registry_tests {
         let registry = FnRegistry::default();
         let outer = Id::new("test.double");
         let inner = Id::new("test.int");
+        let id = DefaultHashBuilder::default().hash_one((outer, inner));
 
-        assert!(!registry.contains(outer, inner));
-        registry.register::<DoubleFn>(outer, inner, double);
+        assert!(!registry.contains(id));
+        registry.register::<DoubleFn>(id, double);
 
-        assert!(registry.contains(outer, inner));
-        let f = registry.find::<DoubleFn>(outer, inner).unwrap();
+        assert!(registry.contains(id));
+        let f = registry.find::<DoubleFn>(id).unwrap();
         assert_eq!(f(21), 42);
     }
 
@@ -357,17 +370,20 @@ mod fn_registry_tests {
         let registry = FnRegistry::default();
         let outer = Id::new("test.double");
         let inner = Id::new("test.int");
-        registry.register::<DoubleFn>(outer, inner, double);
+        let id = DefaultHashBuilder::default().hash_one((outer, inner));
+
+        registry.register::<DoubleFn>(id, double);
 
         type OtherFn = fn(i32) -> i32;
-        assert!(registry.find::<OtherFn>(outer, inner).is_none());
+        assert!(registry.find::<OtherFn>(id).is_none());
     }
 
     #[test]
     fn missing_entry_returns_none() {
         let registry = FnRegistry::default();
         let outer = Id::new("test.missing");
         let inner = Id::new("test.int");
-        assert!(registry.find::<DoubleFn>(outer, inner).is_none());
+        let id = DefaultHashBuilder::default().hash_one((outer, inner));
+        assert!(registry.find::<DoubleFn>(id).is_none());
     }
 }

vortex-session/Cargo.toml

@@ -12,7 +12,7 @@ use vortex_array::dtype::session::DTypeSession;
 // vortex::expr is in the process of having its dependencies inverted, and will eventually be
 // pulled back out into a vortex_expr crate.
 pub use vortex_array::expr;
-use vortex_array::optimizer::session::OptimizerSession;
+use vortex_array::optimizer::kernels::ArrayKernels;
 pub use vortex_array::scalar_fn;
 use vortex_array::scalar_fn::session::ScalarFnSession;
 use vortex_array::session::ArraySession;
@@ -166,7 +166,7 @@ impl VortexSessionDefault for VortexSession {
             .with::<ArraySession>()
             .with::<LayoutSession>()
             .with::<ScalarFnSession>()
-            .with::<OptimizerSession>()
+            .with::<ArrayKernels>()
             .with::<AggregateFnSession>()
             .with::<RuntimeSession>();