vortex_array/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::expr::Expression;
9use crate::expr::exprs::get_item::GetItem;
10use crate::expr::exprs::select::Select;
11use crate::expr::traversal::{NodeVisitor, TraversalOrder};
12
13#[derive(Default)]
14pub struct ReferenceCollector {
15    fields: HashSet<FieldName>,
16}
17
18impl ReferenceCollector {
19    pub fn new() -> Self {
20        Self {
21            fields: HashSet::new(),
22        }
23    }
24
25    pub fn with_set(set: HashSet<FieldName>) -> Self {
26        Self { fields: set }
27    }
28
29    pub fn into_fields(self) -> HashSet<FieldName> {
30        self.fields
31    }
32}
33
34impl NodeVisitor<'_> for ReferenceCollector {
35    type NodeTy = Expression;
36
37    fn visit_up(&mut self, node: &Expression) -> VortexResult<TraversalOrder> {
38        if let Some(get_item) = node.as_opt::<GetItem>() {
39            self.fields.insert(get_item.data().clone());
40        }
41        if let Some(sel) = node.as_opt::<Select>() {
42            self.fields.extend(sel.data().field_names().iter().cloned());
43        }
44        Ok(TraversalOrder::Continue)
45    }
46}
vortex_array/expr/traversal/references.rs

vortex_array/expr/traversal/
references.rs
