vortex_array/expr/analysis/
labeling.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexExpect;
5use vortex_error::VortexResult;
6use vortex_utils::aliases::hash_map::HashMap;
7
8use crate::expr::BoundExpression;
9use crate::expr::ExactBoundExpr;
10use crate::expr::traversal::NodeExt;
11use crate::expr::traversal::NodeVisitor;
12use crate::expr::traversal::TraversalOrder;
13
14/// Labels keyed by bound-tree identity.
15pub type BoundLabels<L> = HashMap<ExactBoundExpr, L>;
16
17/// Label each node in a bound expression using identity-keyed lookups.
18///
19/// This avoids structurally hashing bound dtypes, which may deserialize a lazy schema.
20pub fn label_bound_tree<L: Clone>(
21    expr: &BoundExpression,
22    self_label: impl Fn(&BoundExpression) -> L,
23    mut merge_child: impl FnMut(L, &L) -> L,
24) -> BoundLabels<L> {
25    let mut visitor = BoundLabelingVisitor {
26        labels: Default::default(),
27        self_label,
28        merge_child: &mut merge_child,
29    };
30    expr.accept(&mut visitor)
31        .vortex_expect("BoundLabelingVisitor is infallible");
32    visitor.labels
33}
34
35struct BoundLabelingVisitor<'a, L, F, G>
36where
37    F: Fn(&BoundExpression) -> L,
38    G: FnMut(L, &L) -> L,
39{
40    labels: BoundLabels<L>,
41    self_label: F,
42    merge_child: &'a mut G,
43}
44
45impl<'node, 'visitor, L: Clone, F, G> NodeVisitor<'node> for BoundLabelingVisitor<'visitor, L, F, G>
46where
47    F: Fn(&BoundExpression) -> L,
48    G: FnMut(L, &L) -> L,
49{
50    type NodeTy = BoundExpression;
51
52    fn visit_down(&mut self, _node: &'node Self::NodeTy) -> VortexResult<TraversalOrder> {
53        Ok(TraversalOrder::Continue)
54    }
55
56    fn visit_up(&mut self, node: &'node Self::NodeTy) -> VortexResult<TraversalOrder> {
57        let self_label = (self.self_label)(node);
58        let final_label = node.children().iter().fold(self_label, |acc, child| {
59            let child_label = self
60                .labels
61                .get(&ExactBoundExpr(child.clone()))
62                .vortex_expect("child must have label");
63            (self.merge_child)(acc, child_label)
64        });
65        self.labels
66            .insert(ExactBoundExpr(node.clone()), final_label);
67        Ok(TraversalOrder::Continue)
68    }
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74    use crate::expr::col;
75    use crate::expr::eq;
76    use crate::expr::lit;
77    use crate::expr::test_harness::struct_dtype;
78
79    #[test]
80    fn test_tree_depth() -> VortexResult<()> {
81        // Expression: $.col1 = 5
82        // Tree: eq(get_item(root(), "col1"), lit(5))
83        // Depth: root = 1, get_item = 2, lit = 1, eq = 3
84        let expr = eq(col("col1"), lit(5u16)).bind(&struct_dtype())?;
85        let depths = label_bound_tree(
86            &expr,
87            |_node| 1, // Each node has depth 1 by itself
88            |self_depth, child_depth| self_depth.max(*child_depth + 1),
89        );
90
91        // The root (eq) should have depth 3
92        assert_eq!(depths.get(&ExactBoundExpr(expr)), Some(&3));
93        Ok(())
94    }
95
96    #[test]
97    fn test_node_count() -> VortexResult<()> {
98        // Count total nodes in subtree (including self)
99        // Tree: eq(get_item(root(), "col1"), lit(5))
100        // Nodes: eq, get_item, root, lit = 4
101        let expr = eq(col("col1"), lit(5u16)).bind(&struct_dtype())?;
102        let counts = label_bound_tree(
103            &expr,
104            |_node| 1, // Each node counts as 1
105            |self_count, child_count| self_count + *child_count,
106        );
107
108        // Root should have count of 4 (eq, get_item, root, lit)
109        assert_eq!(counts.get(&ExactBoundExpr(expr)), Some(&4));
110        Ok(())
111    }
112}
vortex_array/expr/analysis/labeling.rs

vortex_array/expr/analysis/
labeling.rs
