vortex_array/expr/analysis/
referenced_field_paths.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5use vortex_error::vortex_err;
6
7use crate::dtype::Field;
8use crate::dtype::FieldPath;
9use crate::dtype::FieldPathSet;
10use crate::expr::BoundExpression;
11use crate::expr::traversal::FoldDownContext;
12use crate::expr::traversal::FoldUp;
13use crate::expr::traversal::NodeExt;
14use crate::expr::traversal::NodeFolderContext;
15use crate::scalar_fn::fns::get_item::GetItem;
16use crate::scalar_fn::fns::select::Select;
17
18/// Returns the rooted field paths referenced by an expression.
19///
20/// Iterating the returned set (via [`IntoIterator`]) yields the prefix-minimal covering set: when
21/// one referenced path is a prefix of another, only the prefix is kept. A standalone root
22/// expression is represented by [`FieldPath::root`], which conservatively selects all fields.
23/// Scalar functions other than `GetItem` and `Select` conservatively reference each complete child
24/// output.
25pub fn referenced_field_paths(expr: &BoundExpression) -> VortexResult<FieldPathSet> {
26    let mut collector = ReferencedFieldPaths {
27        field_paths: FieldPathSet::default(),
28    };
29    expr.clone()
30        .fold_context(&vec![FieldPath::root()], &mut collector)?;
31    Ok(collector.field_paths)
32}
33
34/// Threads the set of currently-requested field paths down the expression tree, narrowing it at
35/// each `GetItem`/`Select`, and records the rooted paths reached at each `Root` leaf.
36///
37/// Paths are carried reversed so a `GetItem` can `push` its field instead of prepending it; they
38/// are reversed back to rooted order when recorded at a `Root`, and `Select` reads a path's head
39/// from its last element.
40///
41/// Narrowing is only sound through `GetItem` (a genuine field access) and `Select` (a genuine
42/// column projection). Any other function is opaque—we cannot assume it preserves a field's
43/// provenance—so its children conservatively re-request the whole scope, which is what keeps an
44/// expression like `f($).x` reading every field of `$` rather than just `x`.
45struct ReferencedFieldPaths {
46    field_paths: FieldPathSet,
47}
48
49impl NodeFolderContext for ReferencedFieldPaths {
50    type NodeTy = BoundExpression;
51    type Result = ();
52    type Context = Vec<FieldPath>;
53
54    fn visit_down(
55        &mut self,
56        requested: &Self::Context,
57        node: &BoundExpression,
58    ) -> VortexResult<FoldDownContext<Self::Context, ()>> {
59        if node.is_root() {
60            self.field_paths.extend(
61                requested
62                    .iter()
63                    .map(|path| FieldPath::from_iter(path.parts().iter().rev().cloned())),
64            );
65            return Ok(FoldDownContext::Skip(()));
66        }
67
68        if let Some(field_name) = node
69            .as_scalar()
70            .and_then(|scalar_fn| scalar_fn.as_opt::<GetItem>())
71        {
72            let appended = requested
73                .iter()
74                .map(|path| path.clone().push(Field::Name(field_name.clone())))
75                .collect();
76            return Ok(FoldDownContext::Continue(appended));
77        }
78
79        // Keep requested paths whose head is included, expanding a whole-scope request into one
80        // path per included field.
81        if let Some(selection) = node
82            .as_scalar()
83            .and_then(|scalar_fn| scalar_fn.as_opt::<Select>())
84        {
85            let child_fields = node.children()[0]
86                .dtype()
87                .as_struct_fields_opt()
88                .ok_or_else(|| vortex_err!("Select child is not a struct"))?;
89            let included_fields = selection.normalize_to_included_fields(child_fields.names())?;
90
91            let mut narrowed = Vec::with_capacity(requested.len());
92            for path in requested {
93                if path.is_root() {
94                    narrowed.extend(included_fields.iter().cloned().map(FieldPath::from_name));
95                } else if let Some(Field::Name(field_name)) = path.parts().last()
96                    && included_fields
97                        .iter()
98                        .any(|included| included == field_name)
99                {
100                    narrowed.push(path.clone());
101                }
102            }
103
104            // Nothing is requested below this `Select`, so prune the subtree rather than letting an
105            // opaque child re-request the whole scope.
106            if narrowed.is_empty() {
107                return Ok(FoldDownContext::Skip(()));
108            }
109            return Ok(FoldDownContext::Continue(narrowed));
110        }
111
112        // Any other function conservatively references each child's complete output.
113        Ok(FoldDownContext::Continue(vec![FieldPath::root()]))
114    }
115
116    fn visit_up(
117        &mut self,
118        _node: BoundExpression,
119        _requested: &Self::Context,
120        _children: Vec<()>,
121    ) -> VortexResult<FoldUp<()>> {
122        Ok(FoldUp::Continue(()))
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use vortex_utils::aliases::hash_set::HashSet;
129
130    use super::*;
131    use crate::dtype::DType;
132    use crate::dtype::Nullability::NonNullable;
133    use crate::dtype::PType::I32;
134    use crate::dtype::StructFields;
135    use crate::expr::Expression;
136    use crate::expr::get_item;
137    use crate::expr::pack;
138    use crate::expr::root;
139    use crate::expr::select;
140    use crate::expr::select_exclude;
141
142    fn scope() -> DType {
143        DType::Struct(
144            StructFields::from_iter([(
145                "a",
146                DType::Struct(
147                    StructFields::from_iter([("x", I32), ("y", I32)]),
148                    NonNullable,
149                ),
150            )]),
151            NonNullable,
152        )
153    }
154
155    /// Collects the prefix-minimal field paths referenced by `expr` against [`scope`].
156    fn referenced(expr: &Expression) -> VortexResult<HashSet<FieldPath>> {
157        Ok(referenced_field_paths(&expr.bind(&scope())?)?
158            .into_iter()
159            .collect())
160    }
161
162    #[test]
163    fn nested_select_preserves_field_path() -> VortexResult<()> {
164        let expr = select(["x"], get_item("a", root()));
165
166        assert_eq!(
167            referenced(&expr)?,
168            HashSet::from_iter([FieldPath::from_name("a").push("x")])
169        );
170        Ok(())
171    }
172
173    #[test]
174    fn get_item_after_select_only_references_requested_field() -> VortexResult<()> {
175        let expr = get_item("x", select(["x", "y"], get_item("a", root())));
176
177        assert_eq!(
178            referenced(&expr)?,
179            HashSet::from_iter([FieldPath::from_name("a").push("x")])
180        );
181        Ok(())
182    }
183
184    #[test]
185    fn select_exclude_references_included_fields() -> VortexResult<()> {
186        let expr = select_exclude(["y"], get_item("a", root()));
187
188        assert_eq!(
189            referenced(&expr)?,
190            HashSet::from_iter([FieldPath::from_name("a").push("x")])
191        );
192        Ok(())
193    }
194
195    #[test]
196    fn ancestor_path_subsumes_descendant() -> VortexResult<()> {
197        let expr = pack(
198            [
199                ("a", get_item("a", root())),
200                ("x", get_item("x", get_item("a", root()))),
201            ],
202            NonNullable,
203        );
204
205        assert_eq!(
206            referenced(&expr)?,
207            HashSet::from_iter([FieldPath::from_name("a")])
208        );
209        Ok(())
210    }
211
212    #[test]
213    fn get_item_through_opaque_fn_references_all_fields() -> VortexResult<()> {
214        // `pack` is opaque to the path analysis: a `GetItem` of its output must not be pushed down
215        // as a scope field access, so the wrapped `root()` conservatively references all fields.
216        let expr = get_item("x", pack([("x", root())], NonNullable));
217
218        assert_eq!(referenced(&expr)?, HashSet::from_iter([FieldPath::root()]));
219        Ok(())
220    }
221
222    #[test]
223    fn root_references_all_fields() -> VortexResult<()> {
224        assert_eq!(
225            referenced(&root())?,
226            HashSet::from_iter([FieldPath::root()])
227        );
228        Ok(())
229    }
230
231    #[test]
232    fn invalid_get_item_path_returns_error() {
233        let result = get_item("missing", root())
234            .bind(&scope())
235            .and_then(|expr| referenced_field_paths(&expr));
236        assert!(result.is_err());
237    }
238}
vortex_array/expr/analysis/referenced_field_paths.rs

vortex_array/expr/analysis/
referenced_field_paths.rs
