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::ArcSwapMap;
8use vortex_session::SessionExt;
9use vortex_session::SessionGuard;
10use vortex_session::SessionVar;
11use vortex_session::registry::Id;
12
13use crate::scalar_fn::ScalarFnPluginRef;
14use crate::scalar_fn::ScalarFnVTable;
15use crate::scalar_fn::fns::between::Between;
16use crate::scalar_fn::fns::binary::Binary;
17use crate::scalar_fn::fns::byte_length::ByteLength;
18use crate::scalar_fn::fns::case_when::CaseWhen;
19use crate::scalar_fn::fns::cast::Cast;
20use crate::scalar_fn::fns::ext_storage::ExtStorage;
21use crate::scalar_fn::fns::fill_null::FillNull;
22use crate::scalar_fn::fns::get_item::GetItem;
23use crate::scalar_fn::fns::is_nan::IsNan;
24use crate::scalar_fn::fns::is_not_null::IsNotNull;
25use crate::scalar_fn::fns::is_null::IsNull;
26use crate::scalar_fn::fns::like::Like;
27use crate::scalar_fn::fns::list_contains::ListContains;
28use crate::scalar_fn::fns::list_contains::PreparedSet;
29use crate::scalar_fn::fns::list_length::ListLength;
30use crate::scalar_fn::fns::list_sum::ListSum;
31use crate::scalar_fn::fns::literal::Literal;
32use crate::scalar_fn::fns::mask::Mask;
33use crate::scalar_fn::fns::merge::Merge;
34use crate::scalar_fn::fns::not::Not;
35use crate::scalar_fn::fns::pack::Pack;
36use crate::scalar_fn::fns::replace_time_zone::ReplaceTimeZone;
37use crate::scalar_fn::fns::select::Select;
38use crate::scalar_fn::fns::stat::StatFn;
39use crate::scalar_fn::fns::variant_get::VariantGet;
40use crate::scalar_fn::fns::zip::Zip;
41
42/// Registry of scalar function vtables.
43pub type ScalarFnRegistry = ArcSwapMap<Id, ScalarFnPluginRef>;
44
45/// Session state for scalar function vtables and rewrite rules.
46#[derive(Clone, Debug)]
47pub struct ScalarFnSession {
48    registry: ScalarFnRegistry,
49}
50
51impl ScalarFnSession {
52    pub fn registry(&self) -> &ScalarFnRegistry {
53        &self.registry
54    }
55
56    /// Register a scalar function vtable in the session, replacing any existing vtable with the same ID.
57    pub fn register<V: ScalarFnVTable>(&self, vtable: V) {
58        self.registry
59            .insert(vtable.id(), Arc::new(vtable) as ScalarFnPluginRef);
60    }
61}
62
63impl Default for ScalarFnSession {
64    fn default() -> Self {
65        let this = Self {
66            registry: ScalarFnRegistry::default(),
67        };
68
69        // Register built-in expressions.
70        this.register(Between);
71        this.register(Binary);
72        this.register(ByteLength);
73        this.register(CaseWhen);
74        this.register(Cast);
75        this.register(ExtStorage);
76        this.register(FillNull);
77        this.register(GetItem);
78        this.register(IsNan);
79        this.register(IsNotNull);
80        this.register(IsNull);
81        this.register(Like);
82        this.register(ListContains);
83        this.register(PreparedSet);
84        this.register(ListLength);
85        this.register(ListSum);
86        this.register(Literal);
87        this.register(Mask);
88        this.register(Merge);
89        this.register(Not);
90        this.register(Pack);
91        this.register(ReplaceTimeZone);
92        this.register(Select);
93        this.register(StatFn);
94        this.register(VariantGet);
95        this.register(Zip);
96
97        this
98    }
99}
100
101impl SessionVar for ScalarFnSession {
102    fn as_any(&self) -> &dyn Any {
103        self
104    }
105
106    fn as_any_mut(&mut self) -> &mut dyn Any {
107        self
108    }
109}
110
111/// Extension trait for accessing scalar function session data.
112pub trait ScalarFnSessionExt: SessionExt {
113    /// Returns the scalar function vtable registry.
114    fn scalar_fns(&self) -> SessionGuard<'_, ScalarFnSession> {
115        self.get::<ScalarFnSession>()
116    }
117}
118impl<S: SessionExt> ScalarFnSessionExt for S {}
vortex_array/scalar_fn/session.rs

vortex_array/scalar_fn/
session.rs
