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

vortex_array/expr/
session.rs
