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

vortex_expr/
scope.rs
