wip

Signed-off-by: Joe Isaacs <[email protected]>

Author: joseph-isaacs

vortex-array/src/expr/analysis/annotations.rs

@@ -6,7 +6,9 @@ use vortex_error::VortexExpect;
 use vortex_utils::aliases::hash_set::HashSet;
 
 use crate::expr::Expression;
-use crate::expr::analysis::annotations::{AnnotationFn, Annotations, descendent_annotations};
+use crate::expr::analysis::annotation_union_set::{
+    AnnotationFn, Annotations, descendent_annotation_union_set,
+};
 use crate::expr::exprs::get_item::GetItem;
 use crate::expr::exprs::root::Root;
 use crate::expr::exprs::select::Select;
@@ -44,7 +46,7 @@ pub fn immediate_scope_accesses<'a>(
     expr: &'a Expression,
     scope: &'a StructFields,
 ) -> FieldAccesses<'a> {
-    descendent_annotations(expr, annotate_scope_access(scope))
+    descendent_annotation_union_set(expr, annotate_scope_access(scope))
 }
 
 /// This returns the immediate scope_access (as explained `immediate_scope_accesses`) for `expr`.

vortex-array/src/expr/analysis/immediate_access.rs

@@ -0,0 +1,145 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::{VortexExpect, VortexResult};
+use vortex_utils::aliases::hash_map::HashMap;
+
+use crate::expr::Expression;
+use crate::expr::traversal::{NodeExt, NodeVisitor, TraversalOrder};
+
+/// Label each node in an expression tree using a bottom-up traversal.
+///
+/// This function separates tree labeling into two distinct steps:
+/// 1. **Label Edge**: Compute a label for each node based only on the node itself
+/// 2. **Merge Child**: Fold/accumulate labels from children into the node's self-label
+///
+/// The labeling process:
+/// - First, `label_edge` is called on the node to produce its self-label
+/// - Then, for each child, `merge_child` is called with `(accumulator, child_label)`
+///   to fold the child label into the accumulator, starting with the self-label
+/// - This produces the final label for the node
+///
+/// This approach avoids allocating a Vec for child labels by processing them one at a time.
+///
+/// # Parameters
+///
+/// - `expr`: The root expression to label
+/// - `label_edge`: Function that computes a label for a single node
+/// - `merge_child`: Mutable function that folds child labels into an accumulator.
+///   Takes `(accumulator, child_label)` and returns the updated accumulator.
+///   Called once per child, with the initial accumulator being the node's self-label.
+///
+/// # Examples
+///
+/// ```ignore
+/// // Count depth of each subtree
+/// let depths = label_tree(
+///     expr,
+///     |_node| 1,  // Each node has depth 1 by itself
+///     |self_depth, child_depth| self_depth.max(*child_depth + 1)
+/// );
+/// ```
+///
+/// ```ignore
+/// // Check if any node in subtree is null-sensitive
+/// let sensitive = label_tree(
+///     expr,
+///     |node| node.is_null_sensitive(),
+///     |acc, child| acc || *child  // OR all children with self
+/// );
+/// ```
+pub fn label_tree<L: Clone>(
+    expr: &Expression,
+    label_edge: impl Fn(&Expression) -> L,
+    mut merge_child: impl FnMut(L, &L) -> L,
+) -> HashMap<&Expression, L> {
+    let mut visitor = LabelingVisitor {
+        labels: Default::default(),
+        label_edge,
+        merge_child: &mut merge_child,
+    };
+    expr.accept(&mut visitor)
+        .vortex_expect("LabelingVisitor is infallible");
+    visitor.labels
+}
+
+struct LabelingVisitor<'a, 'b, L, F, G>
+where
+    F: Fn(&Expression) -> L,
+    G: FnMut(L, &L) -> L,
+{
+    labels: HashMap<&'a Expression, L>,
+    label_edge: F,
+    merge_child: &'b mut G,
+}
+
+impl<'a, 'b, L: Clone, F, G> NodeVisitor<'a> for LabelingVisitor<'a, 'b, L, F, G>
+where
+    F: Fn(&Expression) -> L,
+    G: FnMut(L, &L) -> L,
+{
+    type NodeTy = Expression;
+
+    fn visit_down(&mut self, _node: &'a Self::NodeTy) -> VortexResult<TraversalOrder> {
+        // Continue traversing down
+        Ok(TraversalOrder::Continue)
+    }
+
+    fn visit_up(&mut self, node: &'a Expression) -> VortexResult<TraversalOrder> {
+        // Compute self-label for this node
+        let self_label = (self.label_edge)(node);
+
+        // Fold all child labels into the self label
+        let final_label = node.children().iter().fold(self_label, |acc, child| {
+            let child_label = self
+                .labels
+                .get(child)
+                .vortex_expect("child must have label");
+            (self.merge_child)(acc, child_label)
+        });
+
+        self.labels.insert(node, final_label);
+
+        Ok(TraversalOrder::Continue)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::expr::exprs::binary::eq;
+    use crate::expr::exprs::get_item::col;
+    use crate::expr::exprs::literal::lit;
+
+    #[test]
+    fn test_tree_depth() {
+        // Expression: $.col1 = 5
+        // Tree: eq(get_item(root(), "col1"), lit(5))
+        // Depth: root = 1, get_item = 2, lit = 1, eq = 3
+        let expr = eq(col("col1"), lit(5));
+        let depths = label_tree(
+            &expr,
+            |_node| 1, // Each node has depth 1 by itself
+            |self_depth, child_depth| self_depth.max(*child_depth + 1),
+        );
+
+        // The root (eq) should have depth 3
+        assert_eq!(depths.get(&expr), Some(&3));
+    }
+
+    #[test]
+    fn test_node_count() {
+        // Count total nodes in subtree (including self)
+        // Tree: eq(get_item(root(), "col1"), lit(5))
+        // Nodes: eq, get_item, root, lit = 4
+        let expr = eq(col("col1"), lit(5));
+        let counts = label_tree(
+            &expr,
+            |_node| 1, // Each node counts as 1
+            |self_count, child_count| self_count + *child_count,
+        );
+
+        // Root should have count of 4 (eq, get_item, root, lit)
+        assert_eq!(counts.get(&expr), Some(&4));
+    }
+}

vortex-array/src/expr/analysis/labeling.rs

@@ -1,27 +1,12 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-pub mod annotations;
+pub mod annotation_union_set;
 pub mod immediate_access;
+mod labeling;
 mod null_sensitive;
 
-pub use annotations::*;
+pub use annotation_union_set::*;
 pub use immediate_access::*;
+pub use labeling::*;
 pub use null_sensitive::*;
-use vortex_dtype::FieldPath;
-
-use crate::expr::Expression;
-use crate::stats::Stat;
-
-/// A catalog of available stats that are associated with field paths.
-pub trait StatsCatalog {
-    /// Given a field path and statistic, return an expression that when evaluated over the catalog
-    /// will return that stat for the referenced field.
-    ///
-    /// This is likely to be a column expression, or a literal.
-    ///
-    /// Returns `None` if the stat is not available for the field path.
-    fn stats_ref(&self, _field_path: &FieldPath, _stat: Stat) -> Option<Expression> {
-        None
-    }
-}