Add ability to override function behaviour via registry in VortexSession

robert3005 · robert3005 · commit 78a34a012065 · 2026-04-21T22:13:38.000-04:00
Signed-off-by: Robert Kruszewski &lt;github@robertk.io&gt;
diff --git a/vortex-array/public-api.lock b/vortex-array/public-api.lock
@@ -13388,18 +13388,50 @@ pub fn vortex_array::optimizer::rules::ParentReduceRuleAdapter<V, K>::matches(&s
 
 pub fn vortex_array::optimizer::rules::ParentReduceRuleAdapter<V, K>::reduce_parent(&self, child: vortex_array::ArrayView<'_, V>, parent: &vortex_array::ArrayRef, child_idx: usize) -> vortex_error::VortexResult<core::option::Option<vortex_array::ArrayRef>>
 
+pub mod vortex_array::optimizer::session
+
+pub struct vortex_array::optimizer::session::OptimizerSession
+
+impl vortex_array::optimizer::session::OptimizerSession
+
+pub fn vortex_array::optimizer::session::OptimizerSession::empty() -> Self
+
+pub fn vortex_array::optimizer::session::OptimizerSession::registry(&self) -> &vortex_session::registry::FnRegistry
+
+impl core::default::Default for vortex_array::optimizer::session::OptimizerSession
+
+pub fn vortex_array::optimizer::session::OptimizerSession::default() -> vortex_array::optimizer::session::OptimizerSession
+
+impl core::fmt::Debug for vortex_array::optimizer::session::OptimizerSession
+
+pub fn vortex_array::optimizer::session::OptimizerSession::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+pub trait vortex_array::optimizer::session::OptimizerSessionExt: vortex_session::SessionExt
+
+pub fn vortex_array::optimizer::session::OptimizerSessionExt::optimizer(&self) -> vortex_session::Ref<'_, vortex_array::optimizer::session::OptimizerSession>
+
+impl<S: vortex_session::SessionExt> vortex_array::optimizer::session::OptimizerSessionExt for S
+
+pub fn S::optimizer(&self) -> vortex_session::Ref<'_, vortex_array::optimizer::session::OptimizerSession>
+
 pub trait vortex_array::optimizer::ArrayOptimizer
 
 pub fn vortex_array::optimizer::ArrayOptimizer::optimize(&self) -> vortex_error::VortexResult<vortex_array::ArrayRef>
 
+pub fn vortex_array::optimizer::ArrayOptimizer::optimize_ctx(&self, ctx: &vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::ArrayRef>
+
 pub fn vortex_array::optimizer::ArrayOptimizer::optimize_recursive(&self) -> vortex_error::VortexResult<vortex_array::ArrayRef>
 
 impl vortex_array::optimizer::ArrayOptimizer for vortex_array::ArrayRef
 
 pub fn vortex_array::ArrayRef::optimize(&self) -> vortex_error::VortexResult<vortex_array::ArrayRef>
 
+pub fn vortex_array::ArrayRef::optimize_ctx(&self, ctx: &vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::ArrayRef>
+
 pub fn vortex_array::ArrayRef::optimize_recursive(&self) -> vortex_error::VortexResult<vortex_array::ArrayRef>
 
+pub type vortex_array::optimizer::ReduceParentFn = fn(child: &vortex_array::ArrayRef, parent: &vortex_array::ArrayRef, child_idx: usize) -> vortex_error::VortexResult<core::option::Option<vortex_array::ArrayRef>>
+
 pub mod vortex_array::patches
 
 pub struct vortex_array::patches::Patches
@@ -22356,6 +22388,8 @@ impl vortex_array::optimizer::ArrayOptimizer for vortex_array::ArrayRef
 
 pub fn vortex_array::ArrayRef::optimize(&self) -> vortex_error::VortexResult<vortex_array::ArrayRef>
 
+pub fn vortex_array::ArrayRef::optimize_ctx(&self, ctx: &vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::ArrayRef>
+
 pub fn vortex_array::ArrayRef::optimize_recursive(&self) -> vortex_error::VortexResult<vortex_array::ArrayRef>
 
 impl vortex_array::scalar_fn::ReduceNode for vortex_array::ArrayRef
diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs
@@ -107,7 +107,7 @@ impl ArrayRef {
     /// For safety, we will error when the number of execution iterations reaches a configurable
     /// maximum (default 128, override with `VORTEX_MAX_ITERATIONS`).
     pub fn execute_until<M: Matcher>(self, ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
-        let mut current = self.optimize()?;
+        let mut current = self.optimize_ctx(ctx)?;
         let mut stack: Vec<StackFrame> = Vec::new();
 
         for _ in 0..max_iterations() {
@@ -122,7 +122,7 @@ impl ArrayRef {
                         return Ok(current);
                     }
                     Some(frame) => {
-                        current = frame.put_back(current)?.optimize()?;
+                        current = frame.put_back(current)?.optimize_ctx(ctx)?;
                         continue;
                     }
                 }
@@ -138,9 +138,9 @@ impl ArrayRef {
                     "execute_parent rewrote {} -> {}",
                     current, rewritten
                 ));
-                current = rewritten.optimize()?;
+                current = rewritten.optimize_ctx(ctx)?;
                 if let Some(frame) = stack.pop() {
-                    current = frame.put_back(current)?.optimize()?;
+                    current = frame.put_back(current)?.optimize_ctx(ctx)?;
                 }
                 continue;
             }
@@ -159,7 +159,7 @@ impl ArrayRef {
                     ));
                     let frame = StackFrame::new(parent, i, done, &child);
                     stack.push(frame);
-                    current = child.optimize()?;
+                    current = child.optimize_ctx(ctx)?;
                 }
                 ExecutionStep::Done => {
                     ctx.log(format_args!("Done: {}", array));
diff --git a/vortex-array/src/optimizer/mod.rs b/vortex-array/src/optimizer/mod.rs
@@ -6,34 +6,89 @@
 //!
 //! Optimization runs between execution steps, which is what enables cross-step optimizations:
 //! after a child is decoded, new `reduce_parent` rules may match that were previously blocked.
+//!
+//! There are two entry points:
+//!
+//! * [`ArrayOptimizer::optimize`] — runs the static rules only (the child encoding's
+//!   `PARENT_RULES`). It does not require an execution context and is used by helpers like
+//!   `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)`
+//!   before each `reduce_parent` step. The execute loop calls this entry point so plugin-
+//!   registered parent-reduce rules fire during execution.
+
+use std::sync::Arc;
 
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
+use vortex_session::SessionExt;
 
 use crate::ArrayRef;
+use crate::ExecutionCtx;
+use crate::optimizer::session::OptimizerSession;
 
 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 only by running reduce and reduce_parent rules to fixpoint.
+    /// 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.
     fn optimize(&self) -> VortexResult<ArrayRef>;
 
-    /// Optimize the entire array tree recursively (root and all descendants).
+    /// Like [`Self::optimize`], but additionally consults the session's [`OptimizerSession`]
+    /// registry for each `(parent_encoding_id, child_encoding_id)` pair before the static
+    /// vtable rules.
+    fn optimize_ctx(&self, ctx: &ExecutionCtx) -> VortexResult<ArrayRef>;
+
+    /// Optimize the entire array tree recursively (root and all descendants), static rules only.
     fn optimize_recursive(&self) -> VortexResult<ArrayRef>;
 }
 
 impl ArrayOptimizer for ArrayRef {
     fn optimize(&self) -> VortexResult<ArrayRef> {
-        Ok(try_optimize(self)?.unwrap_or_else(|| self.clone()))
+        Ok(try_optimize(self, None)?.unwrap_or_else(|| self.clone()))
+    }
+
+    fn optimize_ctx(&self, ctx: &ExecutionCtx) -> VortexResult<ArrayRef> {
+        Ok(try_optimize(self, Some(ctx))?.unwrap_or_else(|| self.clone()))
     }
 
     fn optimize_recursive(&self) -> VortexResult<ArrayRef> {
         Ok(try_optimize_recursive(self)?.unwrap_or_else(|| self.clone()))
     }
 }
 
-fn try_optimize(array: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
+/// Resolve a pluggable [`ReduceParentFn`] for `(parent, child)` from the session registry.
+///
+/// Returns `None` when no [`OptimizerSession`] is registered, or no function is registered under
+/// `(parent.encoding_id(), child.encoding_id())`. The returned `Arc` is owned so the caller is
+/// free to drop the session borrow before invoking it.
+fn plugin_reduce_parent(
+    ctx: &ExecutionCtx,
+    parent: &ArrayRef,
+    child: &ArrayRef,
+) -> Option<Arc<ReduceParentFn>> {
+    ctx.session().get_opt::<OptimizerSession>().and_then(|s| {
+        s.registry()
+            .find::<ReduceParentFn>(parent.encoding_id(), child.encoding_id())
+    })
+}
+
+fn try_optimize(array: &ArrayRef, ctx: Option<&ExecutionCtx>) -> VortexResult<Option<ArrayRef>> {
     let mut current_array = array.clone();
     let mut any_optimizations = false;
 
@@ -55,6 +110,17 @@ fn try_optimize(array: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
         // Its important to take all slots here, as `current_array` can change inside the loop.
         for (slot_idx, slot) in current_array.slots().iter().enumerate() {
             let Some(child) = slot else { continue };
+
+            // Registry-based override: tried before the child encoding's static PARENT_RULES.
+            if let Some(ctx) = ctx
+                && let Some(plugin) = plugin_reduce_parent(ctx, &current_array, child)
+                && let Some(new_array) = plugin(child, &current_array, slot_idx)?
+            {
+                current_array = new_array;
+                any_optimizations = true;
+                continue 'outer;
+            }
+
             if let Some(new_array) = child.reduce_parent(&current_array, slot_idx)? {
                 // If the parent was replaced, then we attempt to reduce it again.
                 current_array = new_array;
@@ -80,7 +146,7 @@ fn try_optimize_recursive(array: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
     let mut current_array = array.clone();
     let mut any_optimizations = false;
 
-    if let Some(new_array) = try_optimize(&current_array)? {
+    if let Some(new_array) = try_optimize(&current_array, None)? {
         current_array = new_array;
         any_optimizations = true;
     }
diff --git a/vortex-session/public-api.lock b/vortex-session/public-api.lock
@@ -40,6 +40,30 @@ impl<T> core::default::Default for vortex_session::registry::Context<T>
 
 pub fn vortex_session::registry::Context<T>::default() -> Self
 
+pub struct vortex_session::registry::FnRegistry(_)
+
+impl vortex_session::registry::FnRegistry
+
+pub fn vortex_session::registry::FnRegistry::contains(&self, outer: vortex_session::registry::Id, inner: vortex_session::registry::Id) -> bool
+
+pub fn vortex_session::registry::FnRegistry::empty() -> Self
+
+pub fn vortex_session::registry::FnRegistry::find<F: core::any::Any + core::marker::Send + core::marker::Sync>(&self, outer: vortex_session::registry::Id, inner: vortex_session::registry::Id) -> core::option::Option<alloc::sync::Arc<F>>
+
+pub fn vortex_session::registry::FnRegistry::register<F: core::any::Any + core::marker::Send + core::marker::Sync>(&self, outer: vortex_session::registry::Id, inner: vortex_session::registry::Id, f: F)
+
+impl core::clone::Clone for vortex_session::registry::FnRegistry
+
+pub fn vortex_session::registry::FnRegistry::clone(&self) -> vortex_session::registry::FnRegistry
+
+impl core::default::Default for vortex_session::registry::FnRegistry
+
+pub fn vortex_session::registry::FnRegistry::default() -> vortex_session::registry::FnRegistry
+
+impl core::fmt::Debug for vortex_session::registry::FnRegistry
+
+pub fn vortex_session::registry::FnRegistry::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
 pub struct vortex_session::registry::Id(_)
 
 impl vortex_session::registry::Id
diff --git a/vortex-session/src/registry.rs b/vortex-session/src/registry.rs
@@ -4,6 +4,7 @@
 //! Many session types use a registry of objects that can be looked up by name to construct
 //! contexts. This module provides a generic registry type for that purpose.
 
+use std::any::Any;
 use std::cmp::Ordering;
 use std::fmt;
 use std::fmt::Debug;
@@ -287,3 +288,86 @@ impl<T: Clone> Context<T> {
         self.ids.read().clone()
     }
 }
+
+/// A registry of type-erased function values keyed by a pair of [`Id`] values.
+///
+/// Each entry stores an `Arc<dyn Any + Send + Sync>` wrapping a caller-supplied concrete type `F`
+/// (typically a function pointer). Callers recover the original type by passing the same `F` to
+/// [`FnRegistry::find`]; `find` returns `Some(Arc<F>)` on a type match and `None` otherwise.
+///
+/// Used for pluggable dispatch keyed by an `(outer, inner)` identifier pair — for example the
+/// 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>>>);
+
+impl FnRegistry {
+    /// Create a new, empty registry.
+    pub fn empty() -> Self {
+        Self::default()
+    }
+
+    /// 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));
+    }
+
+    /// 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()
+    }
+
+    /// 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))
+    }
+}
+
+#[cfg(test)]
+mod fn_registry_tests {
+    use super::FnRegistry;
+    use super::Id;
+
+    type DoubleFn = fn(i64) -> i64;
+
+    fn double(x: i64) -> i64 {
+        x * 2
+    }
+
+    #[test]
+    fn register_and_find() {
+        let registry = FnRegistry::default();
+        let outer = Id::new("test.double");
+        let inner = Id::new("test.int");
+
+        assert!(!registry.contains(outer, inner));
+        registry.register::<DoubleFn>(outer, inner, double);
+
+        assert!(registry.contains(outer, inner));
+        let f = registry.find::<DoubleFn>(outer, inner).unwrap();
+        assert_eq!(f(21), 42);
+    }
+
+    #[test]
+    fn find_with_wrong_type_returns_none() {
+        let registry = FnRegistry::default();
+        let outer = Id::new("test.double");
+        let inner = Id::new("test.int");
+        registry.register::<DoubleFn>(outer, inner, double);
+
+        type OtherFn = fn(i32) -> i32;
+        assert!(registry.find::<OtherFn>(outer, inner).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());
+    }
+}
diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs
@@ -12,6 +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;
 pub use vortex_array::scalar_fn;
 use vortex_array::scalar_fn::session::ScalarFnSession;
 use vortex_array::session::ArraySession;
@@ -165,6 +166,7 @@ impl VortexSessionDefault for VortexSession {
             .with::<ArraySession>()
             .with::<LayoutSession>()
             .with::<ScalarFnSession>()
+            .with::<OptimizerSession>()
             .with::<AggregateFnSession>()
             .with::<RuntimeSession>();
 
