vortex_expr/transform/
immediate_access.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_dtype::{FieldName, StructFields};
5use vortex_error::VortexExpect;
6use vortex_utils::aliases::hash_set::HashSet;
7
8use crate::transform::annotations::{AnnotationFn, Annotations, descendent_annotations};
9use crate::{ExprRef, GetItemVTable, SelectVTable, is_root};
10
11pub type FieldAccesses<'a> = Annotations<'a, FieldName>;
12
13/// An [`AnnotationFn`] for annotating scope accesses.
14pub fn annotate_scope_access(scope: &StructFields) -> impl AnnotationFn<Annotation = FieldName> {
15    move |expr: &ExprRef| {
16        assert!(
17            !expr.is::<SelectVTable>(),
18            "cannot analyse select, simplify the expression"
19        );
20
21        if let Some(get_item) = expr.as_opt::<GetItemVTable>() {
22            if is_root(get_item.child()) {
23                return vec![get_item.field().clone()];
24            }
25        } else if is_root(expr) {
26            return scope.names().iter().cloned().collect();
27        }
28
29        vec![]
30    }
31}
32
33/// For all subexpressions in an expression, find the fields that are accessed directly from the
34/// scope, but not any fields in those fields
35/// e.g. scope = {a: {b: .., c: ..}, d: ..}, expr = root().a.b + root().d accesses {a,d} (not b).
36///
37/// Note: This is a very naive, but simple analysis to find the fields that are accessed directly on an
38/// identity node. This is combined to provide an over-approximation of the fields that are accessed
39/// by an expression.
40pub fn immediate_scope_accesses<'a>(
41    expr: &'a ExprRef,
42    scope: &'a StructFields,
43) -> FieldAccesses<'a> {
44    descendent_annotations(expr, annotate_scope_access(scope))
45}
46
47/// This returns the immediate scope_access (as explained `immediate_scope_accesses`) for `expr`.
48pub fn immediate_scope_access<'a>(
49    expr: &'a ExprRef,
50    scope: &'a StructFields,
51) -> HashSet<FieldName> {
52    immediate_scope_accesses(expr, scope)
53        .get(expr)
54        .vortex_expect("Expression missing from scope accesses, this is a internal bug")
55        .clone()
56}
vortex_expr/transform/immediate_access.rs

vortex_expr/transform/
immediate_access.rs
