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