vortex_expr/traversal/
references.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_dtype::FieldName;
5use vortex_error::VortexResult;
6use vortex_utils::aliases::hash_set::HashSet;
7
8use crate::traversal::{NodeVisitor, TraversalOrder};
9use crate::{ExprRef, GetItemVTable, SelectVTable};
10
11#[derive(Default)]
12pub struct ReferenceCollector {
13    fields: HashSet<FieldName>,
14}
15
16impl ReferenceCollector {
17    pub fn new() -> Self {
18        Self {
19            fields: HashSet::new(),
20        }
21    }
22
23    pub fn with_set(set: HashSet<FieldName>) -> Self {
24        Self { fields: set }
25    }
26
27    pub fn into_fields(self) -> HashSet<FieldName> {
28        self.fields
29    }
30}
31
32impl NodeVisitor<'_> for ReferenceCollector {
33    type NodeTy = ExprRef;
34
35    fn visit_up(&mut self, node: &ExprRef) -> VortexResult<TraversalOrder> {
36        if let Some(get_item) = node.as_opt::<GetItemVTable>() {
37            self.fields.insert(get_item.field().clone());
38        }
39        if let Some(sel) = node.as_opt::<SelectVTable>() {
40            self.fields
41                .extend(sel.selection().field_names().iter().cloned());
42        }
43        Ok(TraversalOrder::Continue)
44    }
45}
vortex_expr/traversal/references.rs

vortex_expr/traversal/
references.rs
