vortex_expr/
scope.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::any::TypeId;
5use std::ops::Deref;
6
7use vortex_array::arrays::StructArray;
8use vortex_array::{ArrayRef, IntoArray};
9
10use crate::scope_vars::{ScopeVar, ScopeVars};
11
12/// Scope define the evaluation context/scope that an expression uses when being evaluated.
13///
14/// Other identifier can be bound with variables either before execution or while executing.
15/// Values can be extracted from the scope by type, see [`ScopeVar`](crate::ScopeVar) for more details.
16///
17/// ```code
18/// <let x = lit(1) in var(Identifier::Identity) + var(x), { Identity -> Primitive[1,2,3]> ->
19/// <var(Identifier::Identity) + var(x), { Identity -> Primitive[1,2,3], x -> ConstantArray(1)> ->
20/// <Primitive[1,2,3] + var(x), { Identity -> Primitive[1,2,3], x -> ConstantArray(1)> ->
21/// <Primitive[1,2,3] + ConstantArray(1), { Identity -> Primitive[1,2,3], x -> ConstantArray(1)> ->
22/// <Primitive[2,3,4], { Identity -> Primitive[1,2,3], x -> ConstantArray(1)>
23/// ```
24///
25/// Other values can be bound before execution e.g.
26///  `<var("x") + var("y") + var("z"), x -> ..., y -> ..., z -> ...>`
27#[derive(Clone)]
28pub struct Scope {
29    root: ArrayRef,
30    /// Variables that can be set on the scope during expression evaluation.
31    scope_vars: ScopeVars,
32}
33
34impl Scope {
35    /// Create a new scope with the given root array.
36    pub fn new(root: ArrayRef) -> Self {
37        Self {
38            root,
39            scope_vars: Default::default(),
40        }
41    }
42
43    /// Create a new scope with the root array set an empty struct.
44    pub fn empty(len: usize) -> Self {
45        Self::new(StructArray::new_with_len(len).into_array())
46    }
47
48    /// Return the root array of the scope.
49    pub fn root(&self) -> &ArrayRef {
50        &self.root
51    }
52
53    /// Returns a new evaluation scope with the given variable applied.
54    pub fn with_scope_var<V: ScopeVar>(mut self, var: V) -> Self {
55        self.scope_vars.insert(TypeId::of::<V>(), Box::new(var));
56        self
57    }
58
59    /// Returns the scope variable of type `V` if it exists.
60    pub fn scope_var<V: ScopeVar>(&self) -> Option<&V> {
61        self.scope_vars
62            .get(&TypeId::of::<V>())
63            .and_then(|boxed| (**boxed).as_any().downcast_ref::<V>())
64    }
65
66    /// Returns the mutable scope variable of type `V` if it exists.
67    pub fn scope_var_mut<V: ScopeVar>(&mut self) -> Option<&mut V> {
68        self.scope_vars
69            .get_mut(&TypeId::of::<V>())
70            .and_then(|boxed| (**boxed).as_any_mut().downcast_mut::<V>())
71    }
72}
73
74impl Deref for Scope {
75    type Target = ArrayRef;
76
77    fn deref(&self) -> &Self::Target {
78        &self.root
79    }
80}
81
82impl From<ArrayRef> for Scope {
83    fn from(value: ArrayRef) -> Self {
84        Self::new(value)
85    }
86}
87
88#[cfg(test)]
89mod test {
90    #[test]
91    fn test_scope_var() {
92        use super::*;
93
94        #[derive(Clone, PartialEq, Eq, Debug)]
95        struct TestVar {
96            value: i32,
97        }
98
99        let scope = Scope::empty(100);
100        assert!(scope.scope_var::<TestVar>().is_none());
101
102        let var = TestVar { value: 42 };
103        let mut scope = scope.with_scope_var(var.clone());
104        assert_eq!(scope.scope_var::<TestVar>(), Some(&var));
105
106        scope.scope_var_mut::<TestVar>().unwrap().value = 43;
107        assert_eq!(scope.scope_var::<TestVar>(), Some(&TestVar { value: 43 }));
108    }
109}
vortex_expr/scope.rs

vortex_expr/
scope.rs
