vortex_expr/traversal/
references.rs

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

vortex_expr/traversal/
references.rs
