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

vortex_array/expr/traversal/
references.rs
