vortex_expr/
scope.rs

1use std::any::Any;
2use std::str::FromStr;
3use std::sync::Arc;
4
5use vortex_array::{Array, ArrayRef};
6use vortex_dtype::DType;
7use vortex_error::{VortexError, VortexResult, vortex_bail, vortex_err};
8use vortex_utils::aliases::hash_map::HashMap;
9
10type ExprScope<T> = HashMap<Identifier, T>;
11
12#[derive(Clone, Debug, Eq, PartialEq, Hash)]
13pub enum Identifier {
14    Identity,
15    Other(Arc<str>),
16}
17
18impl FromStr for Identifier {
19    type Err = VortexError;
20
21    fn from_str(s: &str) -> Result<Self, Self::Err> {
22        if s.is_empty() {
23            vortex_bail!("Empty strings aren't allowed in identifiers")
24        } else {
25            Ok(Identifier::Other(s.into()))
26        }
27    }
28}
29
30impl Identifier {
31    pub fn is_identity(&self) -> bool {
32        matches!(self, Self::Identity)
33    }
34}
35
36impl std::fmt::Display for Identifier {
37    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
38        match self {
39            Identifier::Identity => write!(f, ""),
40            Identifier::Other(v) => write!(f, "{}", v),
41        }
42    }
43}
44
45/// Scope define the evaluation context/scope that an expression uses when being evaluated.
46/// There is a special `Identifier` (`Identity`) which is used to bind the initial array being evaluated
47///
48/// Other identifier can be bound with variables either before execution or while executing (see `Let`).
49/// Values can be extracted from the scope using the `Var` expression.
50///
51/// ```code
52/// <let x = lit(1) in var(Identifier::Identity) + var(x), { Identity -> Primitive[1,2,3]> ->
53/// <var(Identifier::Identity) + var(x), { Identity -> Primitive[1,2,3], x -> ConstantArray(1)> ->
54/// <Primitive[1,2,3] + var(x), { Identity -> Primitive[1,2,3], x -> ConstantArray(1)> ->
55/// <Primitive[1,2,3] + ConstantArray(1), { Identity -> Primitive[1,2,3], x -> ConstantArray(1)> ->
56/// <Primitive[2,3,4], { Identity -> Primitive[1,2,3], x -> ConstantArray(1)>
57/// ```
58///
59/// Other values can be bound before execution e.g.
60///  `<var("x") + var("y") + var("z"), x -> ..., y -> ..., z -> ...>`
61#[derive(Clone, Default)]
62pub struct Scope {
63    array_len: usize,
64    root_scope: Option<ArrayRef>,
65    /// A map from identifiers to arrays
66    arrays: ExprScope<ArrayRef>,
67    /// A map identifiers to opaque values used by expressions, but
68    /// cannot affect the result type/shape.
69    vars: ExprScope<Arc<dyn Any + Send + Sync>>,
70}
71
72impl Scope {
73    pub fn new(arr: ArrayRef) -> Self {
74        Self {
75            array_len: arr.len(),
76            root_scope: Some(arr),
77            ..Default::default()
78        }
79    }
80
81    pub fn empty(len: usize) -> Self {
82        Self {
83            array_len: len,
84            ..Default::default()
85        }
86    }
87
88    /// Get a value out of the scope by its [`Identifier`]
89    pub fn array(&self, id: &Identifier) -> Option<&ArrayRef> {
90        if id.is_identity() {
91            return self.root_scope.as_ref();
92        }
93        self.arrays.get(id)
94    }
95
96    pub fn vars(&self, id: Identifier) -> VortexResult<&Arc<dyn Any + Send + Sync>> {
97        self.vars
98            .get(&id)
99            .ok_or_else(|| vortex_err!("cannot find {} in var scope", id))
100    }
101
102    pub fn is_empty(&self) -> bool {
103        self.array_len == 0
104    }
105
106    pub fn len(&self) -> usize {
107        self.array_len
108    }
109
110    pub fn copy_with_value(&self, ident: Identifier, value: ArrayRef) -> Self {
111        self.clone().with_array(ident, value)
112    }
113
114    /// Register an array with an identifier in the scope, overriding any existing value stored in it.
115    pub fn with_array(mut self, ident: Identifier, value: ArrayRef) -> Self {
116        assert_eq!(value.len(), self.len());
117
118        if ident.is_identity() {
119            self.root_scope = Some(value);
120        } else {
121            self.arrays.insert(ident, value);
122        }
123        self
124    }
125
126    pub fn with_var(mut self, ident: Identifier, var: Arc<dyn Any + Send + Sync>) -> Self {
127        self.vars.insert(ident, var);
128        self
129    }
130
131    pub fn iter(&self) -> impl Iterator<Item = (&Identifier, &ArrayRef)> {
132        let values = self.arrays.iter();
133
134        self.root_scope
135            .iter()
136            .map(|s| (&Identifier::Identity, s))
137            .chain(values)
138    }
139}
140
141impl From<ArrayRef> for Scope {
142    fn from(value: ArrayRef) -> Self {
143        Self::new(value)
144    }
145}
146
147#[derive(Clone, Default)]
148pub struct ScopeDType {
149    root: Option<DType>,
150    types: ExprScope<DType>,
151}
152
153impl From<&Scope> for ScopeDType {
154    fn from(ctx: &Scope) -> Self {
155        Self {
156            root: ctx.root_scope.as_ref().map(|s| s.dtype().clone()),
157            types: HashMap::from_iter(
158                ctx.arrays
159                    .iter()
160                    .map(|(k, v)| (k.clone(), v.dtype().clone())),
161            ),
162        }
163    }
164}
165
166impl ScopeDType {
167    pub fn new(dtype: DType) -> Self {
168        Self {
169            root: Some(dtype),
170            ..Default::default()
171        }
172    }
173
174    pub fn dtype(&self, id: &Identifier) -> Option<&DType> {
175        if id.is_identity() {
176            return self.root.as_ref();
177        }
178        self.types.get(id)
179    }
180
181    pub fn copy_with_value(&self, ident: Identifier, dtype: DType) -> Self {
182        self.clone().with_value(ident, dtype)
183    }
184
185    pub fn with_value(mut self, ident: Identifier, dtype: DType) -> Self {
186        if ident.is_identity() {
187            self.root = Some(dtype);
188        } else {
189            self.types.insert(ident, dtype);
190        }
191        self
192    }
193}
vortex_expr/scope.rs

vortex_expr/
scope.rs
