vortex_array/scalar_fn/
session.rs

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

vortex_array/scalar_fn/
session.rs
