vortex_array/expr/analysis/
immediate_access.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexExpect;
5
6use crate::dtype::FieldName;
7use crate::dtype::StructFields;
8use crate::expr::BoundExpression;
9use crate::expr::Expression;
10use crate::expr::analysis::AnnotationFn;
11use crate::scalar_fn::fns::get_item::GetItem;
12use crate::scalar_fn::fns::root::Root;
13use crate::scalar_fn::fns::select::Select;
14
15/// Returns the "free fields" for this expression node.
16///
17/// A "free field" is a top-level field from the root scope that this expression references—not
18/// nested fields within those top-level fields. For example, `root().a.b` has free field `{a}`,
19/// not `{b}`, because `a` is the top-level field being accessed from root.
20///
21/// The term "free" is borrowed from PL theory's "free variables"—variables that reference an
22/// outer scope rather than being introduced locally.
23///
24/// This is useful for column pruning, where we only need to read the top-level fields that an
25/// expression actually touches.
26///
27/// # Annotation Rules
28///
29/// - **[`Select`]**: Returns the included field names if the child is [`Root`].
30/// - **[`GetItem`] on [`Root`]**: Returns `[field_name]` if the child is [`Root`].
31/// - **[`Root`]**: Returns all field names from `scope` (conservative over-approximation).
32/// - **Everything else**: Returns empty (annotations aggregate from children automatically).
33///
34/// # Example
35///
36/// Given `scope = {a: {b: .., c: ..}, d: ..}` and `expr = root().a.b + root().d`:
37/// - `root().a` has free fields `{a}`.
38/// - `root().d` has free fields `{d}`.
39/// - The full expression has free fields `{a, d}` (not `b`, only top-level fields are tracked).
40pub fn make_free_field_annotator(
41    scope: &StructFields,
42) -> impl AnnotationFn<Expression, Annotation = FieldName> {
43    move |expr: &Expression| {
44        if let Some(selection) = expr.as_opt::<Select>() {
45            if expr.child(0).is::<Root>() {
46                return selection
47                    .normalize_to_included_fields(scope.names())
48                    .vortex_expect("Select fields must be valid for scope")
49                    .into_iter()
50                    .collect();
51            }
52        } else if let Some(field_name) = expr.as_opt::<GetItem>() {
53            if expr.child(0).is::<Root>() {
54                return vec![field_name.clone()];
55            }
56        } else if expr.is::<Root>() {
57            return scope.names().iter().cloned().collect();
58        }
59
60        vec![]
61    }
62}
63
64/// Returns the free top-level fields for bound expression nodes.
65pub fn make_bound_free_field_annotator(
66    scope: &StructFields,
67) -> impl AnnotationFn<BoundExpression, Annotation = FieldName> {
68    move |expr: &BoundExpression| {
69        let Some(scalar_fn) = expr.as_scalar() else {
70            return scope.names().iter().cloned().collect();
71        };
72
73        if let Some(selection) = scalar_fn.as_opt::<Select>() {
74            if expr.children()[0].is_root() {
75                return selection
76                    .normalize_to_included_fields(scope.names())
77                    .vortex_expect("Select fields must be valid for scope")
78                    .into_iter()
79                    .collect();
80            }
81        } else if let Some(field_name) = scalar_fn.as_opt::<GetItem>()
82            && expr.children()[0].is_root()
83        {
84            return vec![field_name.clone()];
85        }
86
87        vec![]
88    }
89}
vortex_array/expr/analysis/immediate_access.rs

vortex_array/expr/analysis/
immediate_access.rs
