vortex_array/expr/
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::expr::ExprVTable;
9use crate::expr::exprs::between::Between;
10use crate::expr::exprs::binary::Binary;
11use crate::expr::exprs::cast::Cast;
12use crate::expr::exprs::get_item::GetItem;
13use crate::expr::exprs::is_null::IsNull;
14use crate::expr::exprs::like::Like;
15use crate::expr::exprs::list_contains::ListContains;
16use crate::expr::exprs::literal::Literal;
17use crate::expr::exprs::merge::Merge;
18use crate::expr::exprs::not::Not;
19use crate::expr::exprs::pack::Pack;
20use crate::expr::exprs::root::Root;
21use crate::expr::exprs::select::Select;
22
23/// Registry of expression vtables.
24pub type ExprRegistry = Registry<ExprVTable>;
25
26/// Session state for expression vtables and rewrite rules.
27#[derive(Debug)]
28pub struct ExprSession {
29    registry: ExprRegistry,
30}
31
32impl ExprSession {
33    pub fn registry(&self) -> &ExprRegistry {
34        &self.registry
35    }
36
37    /// Register an expression vtable in the session, replacing any existing vtable with the same ID.
38    pub fn register(&self, expr: ExprVTable) {
39        self.registry.register(expr)
40    }
41
42    /// Register expression vtables in the session, replacing any existing vtables with the same IDs.
43    pub fn register_many(&self, exprs: impl IntoIterator<Item = ExprVTable>) {
44        self.registry.register_many(exprs);
45    }
46}
47
48impl Default for ExprSession {
49    fn default() -> Self {
50        let expressions = ExprRegistry::default();
51
52        // Register built-in expressions here if needed.
53        expressions.register_many([
54            ExprVTable::new_static(&Between),
55            ExprVTable::new_static(&Binary),
56            ExprVTable::new_static(&Cast),
57            ExprVTable::new_static(&GetItem),
58            ExprVTable::new_static(&IsNull),
59            ExprVTable::new_static(&Like),
60            ExprVTable::new_static(&ListContains),
61            ExprVTable::new_static(&Literal),
62            ExprVTable::new_static(&Merge),
63            ExprVTable::new_static(&Not),
64            ExprVTable::new_static(&Pack),
65            ExprVTable::new_static(&Root),
66            ExprVTable::new_static(&Select),
67        ]);
68
69        Self {
70            registry: expressions,
71        }
72    }
73}
74
75/// Extension trait for accessing expression session data.
76pub trait ExprSessionExt: SessionExt {
77    /// Returns the expression vtable registry.
78    fn expressions(&self) -> Ref<'_, ExprSession> {
79        self.get::<ExprSession>()
80    }
81}
82impl<S: SessionExt> ExprSessionExt for S {}