vortex_expr/transform/
immediate_access.rs

1use vortex_dtype::{FieldName, StructFields};
2use vortex_error::{VortexResult, vortex_err};
3use vortex_utils::aliases::hash_map::HashMap;
4use vortex_utils::aliases::hash_set::HashSet;
5
6use crate::transform::access_analysis::AccessesAnalysis;
7use crate::traversal::TraversalOrder;
8use crate::{ExprRef, GetItem, Select, is_root};
9
10pub type FieldAccesses<'a> = HashMap<&'a ExprRef, HashSet<FieldName>>;
11
12/// For all subexpressions in an expression, find the fields that are accessed directly from the
13/// scope, but not any fields in those fields
14/// e.g. scope = {a: {b: .., c: ..}, d: ..}, expr = ident().a.b + ident().d accesses {a,d} (not b).
15///
16/// Note: This is a very naive, but simple analysis to find the fields that are accessed directly on an
17/// identity node. This is combined to provide an over-approximation of the fields that are accessed
18/// by an expression.
19pub fn immediate_scope_accesses<'a>(
20    expr: &'a ExprRef,
21    scope_dtype: &'a StructFields,
22) -> VortexResult<FieldAccesses<'a>> {
23    AccessesAnalysis::analyze(expr, move |node| {
24        assert!(
25            !node.as_any().is::<Select>(),
26            "cannot analyse select, simplify the expression"
27        );
28        if let Some(get_item) = node.as_any().downcast_ref::<GetItem>() {
29            if is_root(get_item.child()) {
30                return (TraversalOrder::Skip, vec![get_item.field().clone()]);
31            }
32        } else if is_root(node) {
33            let st_dtype = &scope_dtype;
34            return (
35                TraversalOrder::Skip,
36                st_dtype.names().iter().cloned().collect(),
37            );
38        }
39
40        (TraversalOrder::Continue, vec![])
41    })
42}
43
44/// This returns the immediate scope_access (as explained `immediate_scope_accesses`) for `expr`.
45pub fn immediate_scope_access<'a>(
46    expr: &'a ExprRef,
47    scope_dtype: &'a StructFields,
48) -> VortexResult<HashSet<FieldName>> {
49    immediate_scope_accesses(expr, scope_dtype)?
50        .get(expr)
51        .ok_or_else(|| {
52            vortex_err!("Expression missing from scope accesses, this is a internal bug")
53        })
54        .cloned()
55}
vortex_expr/transform/immediate_access.rs

vortex_expr/transform/
immediate_access.rs
