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