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, FieldPathSet};
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 PartialEq<str> for Identifier {
31    fn eq(&self, other: &str) -> bool {
32        match self {
33            Identifier::Identity => other.is_empty(),
34            Identifier::Other(str) => str.as_ref() == other,
35        }
36    }
37}
38
39impl From<&str> for Identifier {
40    fn from(value: &str) -> Self {
41        if value.is_empty() {
42            Identifier::Identity
43        } else {
44            Identifier::Other(Arc::from(value))
45        }
46    }
47}
48
49impl From<&Arc<str>> for Identifier {
50    fn from(value: &Arc<str>) -> Self {
51        if value.as_ref() == "" {
52            Identifier::Identity
53        } else {
54            Identifier::Other(value.clone())
55        }
56    }
57}
58
59impl Identifier {
60    pub fn is_identity(&self) -> bool {
61        matches!(self, Self::Identity)
62    }
63}
64
65impl std::fmt::Display for Identifier {
66    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
67        match self {
68            Identifier::Identity => write!(f, ""),
69            Identifier::Other(v) => write!(f, "{}", v),
70        }
71    }
72}
73
74/// Scope define the evaluation context/scope that an expression uses when being evaluated.
75/// There is a special `Identifier` (`Identity`) which is used to bind the initial array being evaluated
76///
77/// Other identifier can be bound with variables either before execution or while executing (see `Let`).
78/// Values can be extracted from the scope using the `Var` expression.
79///
80/// ```code
81/// <let x = lit(1) in var(Identifier::Identity) + var(x), { Identity -> Primitive[1,2,3]> ->
82/// <var(Identifier::Identity) + var(x), { Identity -> Primitive[1,2,3], x -> ConstantArray(1)> ->
83/// <Primitive[1,2,3] + var(x), { Identity -> Primitive[1,2,3], x -> ConstantArray(1)> ->
84/// <Primitive[1,2,3] + ConstantArray(1), { Identity -> Primitive[1,2,3], x -> ConstantArray(1)> ->
85/// <Primitive[2,3,4], { Identity -> Primitive[1,2,3], x -> ConstantArray(1)>
86/// ```
87///
88/// Other values can be bound before execution e.g.
89///  `<var("x") + var("y") + var("z"), x -> ..., y -> ..., z -> ...>`
90#[derive(Clone, Default)]
91pub struct Scope {
92    array_len: usize,
93    root_scope: Option<ArrayRef>,
94    /// A map from identifiers to arrays
95    arrays: ExprScope<ArrayRef>,
96    /// A map identifiers to opaque values used by expressions, but
97    /// cannot affect the result type/shape.
98    vars: ExprScope<Arc<dyn Any + Send + Sync>>,
99}
100
101impl Scope {
102    pub fn new(arr: ArrayRef) -> Self {
103        Self {
104            array_len: arr.len(),
105            root_scope: Some(arr),
106            ..Default::default()
107        }
108    }
109
110    pub fn empty(len: usize) -> Self {
111        Self {
112            array_len: len,
113            ..Default::default()
114        }
115    }
116
117    /// Get a value out of the scope by its [`Identifier`]
118    pub fn array(&self, id: &Identifier) -> Option<&ArrayRef> {
119        if id.is_identity() {
120            return self.root_scope.as_ref();
121        }
122        self.arrays.get(id)
123    }
124
125    pub fn vars(&self, id: Identifier) -> VortexResult<&Arc<dyn Any + Send + Sync>> {
126        self.vars
127            .get(&id)
128            .ok_or_else(|| vortex_err!("cannot find {} in var scope", id))
129    }
130
131    pub fn is_empty(&self) -> bool {
132        self.array_len == 0
133    }
134
135    pub fn len(&self) -> usize {
136        self.array_len
137    }
138
139    pub fn copy_with_array(&self, ident: Identifier, value: ArrayRef) -> Self {
140        self.clone().with_array(ident, value)
141    }
142
143    /// Register an array with an identifier in the scope, overriding any existing value stored in it.
144    pub fn with_array(mut self, ident: Identifier, value: ArrayRef) -> Self {
145        assert_eq!(value.len(), self.len());
146
147        if ident.is_identity() {
148            self.root_scope = Some(value);
149        } else {
150            self.arrays.insert(ident, value);
151        }
152        self
153    }
154
155    pub fn with_var(mut self, ident: Identifier, var: Arc<dyn Any + Send + Sync>) -> Self {
156        self.vars.insert(ident, var);
157        self
158    }
159
160    pub fn iter(&self) -> impl Iterator<Item = (&Identifier, &ArrayRef)> {
161        let values = self.arrays.iter();
162
163        self.root_scope
164            .iter()
165            .map(|s| (&Identifier::Identity, s))
166            .chain(values)
167    }
168}
169
170impl From<ArrayRef> for Scope {
171    fn from(value: ArrayRef) -> Self {
172        Self::new(value)
173    }
174}
175
176#[derive(Clone, Default, Debug)]
177pub struct ScopeDType {
178    root: Option<DType>,
179    types: ExprScope<DType>,
180}
181
182impl From<&Scope> for ScopeDType {
183    fn from(ctx: &Scope) -> Self {
184        Self {
185            root: ctx.root_scope.as_ref().map(|s| s.dtype().clone()),
186            types: HashMap::from_iter(
187                ctx.arrays
188                    .iter()
189                    .map(|(k, v)| (k.clone(), v.dtype().clone())),
190            ),
191        }
192    }
193}
194
195impl ScopeDType {
196    pub fn new(dtype: DType) -> Self {
197        Self {
198            root: Some(dtype),
199            ..Default::default()
200        }
201    }
202
203    pub fn dtype(&self, id: &Identifier) -> Option<&DType> {
204        if id.is_identity() {
205            return self.root.as_ref();
206        }
207        self.types.get(id)
208    }
209
210    pub fn copy_with_dtype(&self, ident: Identifier, dtype: DType) -> Self {
211        self.clone().with_dtype(ident, dtype)
212    }
213
214    pub fn with_dtype(mut self, ident: Identifier, dtype: DType) -> Self {
215        if ident.is_identity() {
216            self.root = Some(dtype);
217        } else {
218            self.types.insert(ident, dtype);
219        }
220        self
221    }
222}
223
224#[derive(Default, Clone, Debug)]
225pub struct ScopeFieldPathSet {
226    root: Option<FieldPathSet>,
227    sets: ExprScope<FieldPathSet>,
228}
229
230impl ScopeFieldPathSet {
231    pub fn new(path_set: FieldPathSet) -> Self {
232        Self {
233            root: Some(path_set),
234            ..Default::default()
235        }
236    }
237
238    pub fn set(&self, id: &Identifier) -> Option<&FieldPathSet> {
239        if id.is_identity() {
240            return self.root.as_ref();
241        }
242        self.sets.get(id)
243    }
244
245    pub fn copy_with_set(&self, ident: Identifier, set: FieldPathSet) -> Self {
246        self.clone().with_set(ident, set)
247    }
248
249    pub fn with_set(mut self, ident: Identifier, set: FieldPathSet) -> Self {
250        if ident.is_identity() {
251            self.root = Some(set);
252        } else {
253            self.sets.insert(ident, set);
254        }
255        self
256    }
257}
vortex_expr/scope.rs

vortex_expr/
scope.rs
