Add HigherOrderFn infrastructure

Signed-off-by: Matt Katz <mhkatz97@gmail.com>

Author: mhk197

vortex-array/src/higher_order_fn/erased.rs

@@ -0,0 +1,148 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::any::type_name;
+use std::fmt::Debug;
+use std::fmt::Display;
+use std::fmt::Formatter;
+use std::hash::Hash;
+use std::hash::Hasher;
+use std::sync::Arc;
+
+use vortex_error::VortexExpect;
+use vortex_error::VortexResult;
+use vortex_error::vortex_err;
+use vortex_utils::debug_with::DebugWith;
+
+use crate::higher_order_fn::HigherOrderFnId;
+use crate::higher_order_fn::HigherOrderFnOptions;
+use crate::higher_order_fn::HigherOrderFnVTable;
+use crate::higher_order_fn::TypedHigherOrderFnInstance;
+use crate::higher_order_fn::typed::DynHigherOrderFn;
+use crate::scalar_fn::Arity;
+use crate::scalar_fn::ChildName;
+
+/// A type-erased higher-order function, pairing a vtable with per-call options.
+#[derive(Clone)]
+pub struct HigherOrderFnRef(pub(super) Arc<dyn DynHigherOrderFn>);
+
+impl HigherOrderFnRef {
+    /// Bind `vtable` with its per-call `options`.
+    pub fn new<V: HigherOrderFnVTable>(vtable: V, options: V::Options) -> Self {
+        TypedHigherOrderFnInstance::new(vtable, options).erased()
+    }
+
+    /// The function's global identifier.
+    pub fn id(&self) -> HigherOrderFnId {
+        self.0.id()
+    }
+
+    /// Whether this function uses vtable `V`.
+    pub fn is<V: HigherOrderFnVTable>(&self) -> bool {
+        self.0.as_any().is::<TypedHigherOrderFnInstance<V>>()
+    }
+
+    /// Return typed options when this function uses vtable `V`.
+    pub fn as_opt<V: HigherOrderFnVTable>(&self) -> Option<&V::Options> {
+        self.0
+            .as_any()
+            .downcast_ref::<TypedHigherOrderFnInstance<V>>()
+            .map(TypedHigherOrderFnInstance::options)
+    }
+
+    /// Return typed options for vtable `V`.
+    ///
+    /// # Panics
+    ///
+    /// Panics if this function does not use vtable `V`.
+    pub fn as_<V: HigherOrderFnVTable>(&self) -> &V::Options {
+        self.as_opt::<V>()
+            .vortex_expect("higher-order function options type mismatch")
+    }
+
+    /// Return these options behind an opaque type-erased handle.
+    pub fn options(&self) -> HigherOrderFnOptions<'_> {
+        HigherOrderFnOptions { inner: &*self.0 }
+    }
+
+    /// Return the arity of the ordinary arguments.
+    pub fn arity(&self) -> Arity {
+        self.0.arity()
+    }
+
+    /// Return the number of lambda arguments.
+    pub fn lambda_arity(&self) -> usize {
+        self.0.lambda_arity()
+    }
+
+    /// Return the name of an ordinary argument.
+    pub fn child_name(&self, child_idx: usize) -> ChildName {
+        self.0.child_name(child_idx)
+    }
+
+    /// Serialize the per-call options.
+    pub fn serialize(&self) -> VortexResult<Option<Vec<u8>>> {
+        self.0.options_serialize()
+    }
+
+    /// Downcast this function to its typed representation.
+    pub fn try_downcast<V: HigherOrderFnVTable>(
+        self,
+    ) -> Result<Arc<TypedHigherOrderFnInstance<V>>, Self> {
+        if self.is::<V>() {
+            let ptr = Arc::into_raw(self.0) as *const TypedHigherOrderFnInstance<V>;
+            Ok(unsafe { Arc::from_raw(ptr) })
+        } else {
+            Err(self)
+        }
+    }
+
+    /// Downcast this function to its typed representation.
+    ///
+    /// # Panics
+    ///
+    /// Panics if this function does not use vtable `V`.
+    pub fn downcast<V: HigherOrderFnVTable>(self) -> Arc<TypedHigherOrderFnInstance<V>> {
+        self.try_downcast::<V>()
+            .map_err(|function| {
+                vortex_err!(
+                    "failed to downcast higher-order function {} to {}",
+                    function.id(),
+                    type_name::<V>(),
+                )
+            })
+            .vortex_expect("failed to downcast higher-order function")
+    }
+}
+
+impl Debug for HigherOrderFnRef {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("HigherOrderFnRef")
+            .field("vtable", &self.id())
+            .field("options", &DebugWith(|fmt| self.0.options_debug(fmt)))
+            .finish()
+    }
+}
+
+impl Display for HigherOrderFnRef {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}(", self.id())?;
+        self.0.options_display(f)?;
+        write!(f, ")")
+    }
+}
+
+impl PartialEq for HigherOrderFnRef {
+    fn eq(&self, other: &Self) -> bool {
+        self.id() == other.id() && self.0.options_eq(other.0.options_any())
+    }
+}
+
+impl Eq for HigherOrderFnRef {}
+
+impl Hash for HigherOrderFnRef {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.id().hash(state);
+        self.0.options_hash(state);
+    }
+}

vortex-array/src/higher_order_fn/mod.rs

@@ -0,0 +1,29 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Higher-order function vtable machinery.
+//!
+//! Higher-order functions accept ordinary arguments and lambda arguments. This module provides
+//! their identity, type erasure, per-call options, and session registry.
+
+use vortex_session::registry::Id;
+
+mod erased;
+pub use erased::HigherOrderFnRef;
+
+mod options;
+pub use options::HigherOrderFnOptions;
+
+mod typed;
+pub use typed::TypedHigherOrderFnInstance;
+
+mod plugin;
+pub use plugin::*;
+
+pub mod session;
+
+mod vtable;
+pub use vtable::*;
+
+/// A globally unique identifier for a higher-order function.
+pub type HigherOrderFnId = Id;

vortex-array/src/higher_order_fn/options.rs

@@ -0,0 +1,56 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::any::Any;
+use std::fmt::Debug;
+use std::fmt::Display;
+use std::hash::Hash;
+use std::hash::Hasher;
+
+use vortex_error::VortexResult;
+
+use crate::higher_order_fn::typed::DynHigherOrderFn;
+
+/// An opaque handle to the options of a higher-order function.
+pub struct HigherOrderFnOptions<'a> {
+    pub(super) inner: &'a dyn DynHigherOrderFn,
+}
+
+impl Display for HigherOrderFnOptions<'_> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.inner.options_display(f)
+    }
+}
+
+impl Debug for HigherOrderFnOptions<'_> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.inner.options_debug(f)
+    }
+}
+
+impl PartialEq for HigherOrderFnOptions<'_> {
+    fn eq(&self, other: &Self) -> bool {
+        self.inner.id() == other.inner.id() && self.inner.options_eq(other.inner.options_any())
+    }
+}
+
+impl Eq for HigherOrderFnOptions<'_> {}
+
+impl Hash for HigherOrderFnOptions<'_> {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.inner.id().hash(state);
+        self.inner.options_hash(state);
+    }
+}
+
+impl HigherOrderFnOptions<'_> {
+    /// Serialize these options to a byte vector.
+    pub fn serialize(&self) -> VortexResult<Option<Vec<u8>>> {
+        self.inner.options_serialize()
+    }
+
+    /// Return the underlying typed options as [`Any`].
+    pub fn as_any(&self) -> &dyn Any {
+        self.inner.options_any()
+    }
+}

vortex-array/src/higher_order_fn/plugin.rs

@@ -0,0 +1,51 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::sync::Arc;
+
+use vortex_error::VortexResult;
+use vortex_session::VortexSession;
+
+use crate::higher_order_fn::HigherOrderFnId;
+use crate::higher_order_fn::HigherOrderFnRef;
+use crate::higher_order_fn::HigherOrderFnVTable;
+use crate::higher_order_fn::TypedHigherOrderFnInstance;
+
+/// Reference-counted pointer to a higher-order function plugin.
+pub type HigherOrderFnPluginRef = Arc<dyn HigherOrderFnPlugin>;
+
+/// Registry trait for ID-based deserialization of higher-order functions.
+pub trait HigherOrderFnPlugin: 'static + Send + Sync {
+    /// Return the ID for this higher-order function.
+    fn id(&self) -> HigherOrderFnId;
+
+    /// Deserialize a higher-order function from serialized metadata.
+    fn deserialize(
+        &self,
+        metadata: &[u8],
+        session: &VortexSession,
+    ) -> VortexResult<HigherOrderFnRef>;
+}
+
+impl std::fmt::Debug for dyn HigherOrderFnPlugin {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_tuple("HigherOrderFnPlugin")
+            .field(&self.id())
+            .finish()
+    }
+}
+
+impl<V: HigherOrderFnVTable> HigherOrderFnPlugin for V {
+    fn id(&self) -> HigherOrderFnId {
+        V::id(self)
+    }
+
+    fn deserialize(
+        &self,
+        metadata: &[u8],
+        session: &VortexSession,
+    ) -> VortexResult<HigherOrderFnRef> {
+        let options = HigherOrderFnVTable::deserialize(self, metadata, session)?;
+        Ok(TypedHigherOrderFnInstance::new(self.clone(), options).erased())
+    }
+}

vortex-array/src/higher_order_fn/session.rs

@@ -0,0 +1,55 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::any::Any;
+use std::sync::Arc;
+
+use vortex_session::ArcSwapMap;
+use vortex_session::SessionExt;
+use vortex_session::SessionGuard;
+use vortex_session::SessionVar;
+
+use crate::higher_order_fn::HigherOrderFnId;
+use crate::higher_order_fn::HigherOrderFnPluginRef;
+use crate::higher_order_fn::HigherOrderFnVTable;
+
+/// Registry of higher-order function vtables.
+pub type HigherOrderFnRegistry = ArcSwapMap<HigherOrderFnId, HigherOrderFnPluginRef>;
+
+/// Session state for higher-order function vtables.
+#[derive(Clone, Debug, Default)]
+pub struct HigherOrderFnSession {
+    registry: HigherOrderFnRegistry,
+}
+
+impl HigherOrderFnSession {
+    pub fn registry(&self) -> &HigherOrderFnRegistry {
+        &self.registry
+    }
+
+    /// Register a vtable, replacing any existing vtable with the same ID.
+    pub fn register<V: HigherOrderFnVTable>(&self, vtable: V) {
+        self.registry
+            .insert(vtable.id(), Arc::new(vtable) as HigherOrderFnPluginRef);
+    }
+}
+
+impl SessionVar for HigherOrderFnSession {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn as_any_mut(&mut self) -> &mut dyn Any {
+        self
+    }
+}
+
+/// Extension trait for accessing higher-order-function session state.
+pub trait HigherOrderFnSessionExt: SessionExt {
+    /// Return the higher-order function vtable registry.
+    fn higher_order_fns(&self) -> SessionGuard<'_, HigherOrderFnSession> {
+        self.get::<HigherOrderFnSession>()
+    }
+}
+
+impl<S: SessionExt> HigherOrderFnSessionExt for S {}