Skip to main content

vortex_array/scalar_fn/
session.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_session::Ref;
5use vortex_session::SessionExt;
6use vortex_session::registry::Registry;
7
8use crate::scalar_fn::ScalarFnPlugin;
9use crate::scalar_fn::ScalarFnVTable;
10use crate::scalar_fn::fns::between::Between;
11use crate::scalar_fn::fns::binary::Binary;
12use crate::scalar_fn::fns::cast::Cast;
13use crate::scalar_fn::fns::fill_null::FillNull;
14use crate::scalar_fn::fns::get_item::GetItem;
15use crate::scalar_fn::fns::is_null::IsNull;
16use crate::scalar_fn::fns::like::Like;
17use crate::scalar_fn::fns::list_contains::ListContains;
18use crate::scalar_fn::fns::literal::Literal;
19use crate::scalar_fn::fns::merge::Merge;
20use crate::scalar_fn::fns::not::Not;
21use crate::scalar_fn::fns::pack::Pack;
22use crate::scalar_fn::fns::root::Root;
23use crate::scalar_fn::fns::select::Select;
24
25/// Registry of scalar function vtables.
26pub type ScalarFnRegistry = Registry<ScalarFnPlugin>;
27
28/// Session state for scalar function vtables and rewrite rules.
29#[derive(Debug)]
30pub struct ScalarFnSession {
31    registry: ScalarFnRegistry,
32}
33
34impl ScalarFnSession {
35    pub fn registry(&self) -> &ScalarFnRegistry {
36        &self.registry
37    }
38
39    /// Register a scalar function vtable in the session, replacing any existing vtable with the same ID.
40    pub fn register<V: ScalarFnVTable>(&self, vtable: V) {
41        let plugin = ScalarFnPlugin::new(vtable);
42        self.registry.register(plugin.id(), plugin);
43    }
44}
45
46impl Default for ScalarFnSession {
47    fn default() -> Self {
48        let this = Self {
49            registry: ScalarFnRegistry::default(),
50        };
51
52        // Register built-in expressions.
53        this.register(Between);
54        this.register(Binary);
55        this.register(Cast);
56        this.register(FillNull);
57        this.register(GetItem);
58        this.register(IsNull);
59        this.register(Like);
60        this.register(ListContains);
61        this.register(Literal);
62        this.register(Merge);
63        this.register(Not);
64        this.register(Pack);
65        this.register(Root);
66        this.register(Select);
67
68        this
69    }
70}
71
72/// Extension trait for accessing scalar function session data.
73pub trait ScalarFnSessionExt: SessionExt {
74    /// Returns the scalar function vtable registry.
75    fn scalar_fns(&self) -> Ref<'_, ScalarFnSession> {
76        self.get::<ScalarFnSession>()
77    }
78}
79impl<S: SessionExt> ScalarFnSessionExt for S {}