refactor: propagate bound expressions through layouts

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

@@ -8,7 +8,10 @@ use vortex_error::VortexResult;
 use vortex_utils::aliases::hash_map::HashMap;
 use vortex_utils::aliases::hash_set::HashSet;
 
+use crate::expr::BoundExpression;
+use crate::expr::ExactBoundExpr;
 use crate::expr::Expression;
+use crate::expr::traversal::Node;
 use crate::expr::traversal::NodeExt;
 use crate::expr::traversal::NodeVisitor;
 use crate::expr::traversal::TraversalOrder;
@@ -17,28 +20,38 @@ pub trait Annotation: Clone + Hash + Eq {}
 
 impl<A> Annotation for A where A: Clone + Hash + Eq {}
 
-pub trait AnnotationFn: Fn(&Expression) -> Vec<Self::Annotation> {
+pub trait AnnotationFn<N = Expression>: Fn(&N) -> Vec<Self::Annotation> {
     type Annotation: Annotation;
 }
 
-impl<A, F> AnnotationFn for F
+impl<N, A, F> AnnotationFn<N> for F
 where
     A: Annotation,
-    F: Fn(&Expression) -> Vec<A>,
+    F: Fn(&N) -> Vec<A>,
 {
     type Annotation = A;
 }
 
-pub type Annotations<'a, A> = HashMap<&'a Expression, HashSet<A>>;
+pub type Annotations<'a, A, N = Expression> = HashMap<&'a N, HashSet<A>>;
+
+/// Annotations keyed by bound-tree identity.
+///
+/// Identity keys avoid structurally hashing every node's dtype. That matters when a bound root
+/// carries a lazy schema whose structural hash would deserialize every field.
+pub type BoundAnnotations<A> = HashMap<ExactBoundExpr, HashSet<A>>;
 
 /// Walk the expression tree and annotate each expression with zero or more annotations.
 ///
 /// Returns a map of each expression to all annotations that any of its descendent (child)
 /// expressions are annotated with.
-pub fn descendent_annotations<A: AnnotationFn>(
-    expr: &Expression,
+pub fn descendent_annotations<'a, N, A>(
+    expr: &'a N,
     annotate: A,
-) -> Annotations<'_, A::Annotation> {
+) -> Annotations<'a, A::Annotation, N>
+where
+    N: Node + Eq + Hash,
+    A: AnnotationFn<N>,
+{
     let mut visitor = AnnotationVisitor {
         annotations: Default::default(),
         annotate,
@@ -53,10 +66,11 @@ pub fn descendent_annotations<A: AnnotationFn>(
 ///
 /// Returns a map of each expression to all annotations. Annotations of
 /// children are not propagated to parents.
-pub fn direct_annotations<A: AnnotationFn>(
-    expr: &Expression,
-    annotate: A,
-) -> Annotations<'_, A::Annotation> {
+pub fn direct_annotations<'a, N, A>(expr: &'a N, annotate: A) -> Annotations<'a, A::Annotation, N>
+where
+    N: Node + Eq + Hash,
+    A: AnnotationFn<N>,
+{
     let mut visitor = AnnotationVisitor {
         annotations: Default::default(),
         annotate,
@@ -66,14 +80,66 @@ pub fn direct_annotations<A: AnnotationFn>(
     visitor.annotations
 }
 
-struct AnnotationVisitor<'a, A: AnnotationFn> {
-    annotations: Annotations<'a, A::Annotation>,
+/// Annotate a bound expression and propagate each annotation to its ancestors.
+///
+/// Unlike [`descendent_annotations`], this uses [`ExactBoundExpr`] keys to preserve the cheap
+/// identity semantics of an already-bound tree.
+pub fn descendent_bound_annotations<A>(
+    expr: &BoundExpression,
+    annotate: A,
+) -> BoundAnnotations<A::Annotation>
+where
+    A: AnnotationFn<BoundExpression>,
+{
+    bound_annotations(expr, annotate, true)
+}
+
+/// Annotate each bound-expression node without propagating annotations to its ancestors.
+///
+/// The returned map uses [`ExactBoundExpr`] keys so lookups do not structurally hash node dtypes.
+pub fn direct_bound_annotations<A>(
+    expr: &BoundExpression,
+    annotate: A,
+) -> BoundAnnotations<A::Annotation>
+where
+    A: AnnotationFn<BoundExpression>,
+{
+    bound_annotations(expr, annotate, false)
+}
+
+fn bound_annotations<A>(
+    expr: &BoundExpression,
+    annotate: A,
+    propagate_up: bool,
+) -> BoundAnnotations<A::Annotation>
+where
+    A: AnnotationFn<BoundExpression>,
+{
+    let mut visitor = BoundAnnotationVisitor {
+        annotations: Default::default(),
+        annotate,
+        propagate_up,
+    };
+    expr.accept(&mut visitor).vortex_expect("Infallible");
+    visitor.annotations
+}
+
+struct AnnotationVisitor<'a, N, A>
+where
+    N: Node + Eq + Hash,
+    A: AnnotationFn<N>,
+{
+    annotations: Annotations<'a, A::Annotation, N>,
     annotate: A,
     propagate_up: bool,
 }
 
-impl<'a, A: AnnotationFn> NodeVisitor<'a> for AnnotationVisitor<'a, A> {
-    type NodeTy = Expression;
+impl<'a, N, A> NodeVisitor<'a> for AnnotationVisitor<'a, N, A>
+where
+    N: Node + Eq + Hash,
+    A: AnnotationFn<N>,
+{
+    type NodeTy = N;
 
     fn visit_down(&mut self, node: &'a Self::NodeTy) -> VortexResult<TraversalOrder> {
         let annotations = (self.annotate)(node);
@@ -89,20 +155,74 @@ impl<'a, A: AnnotationFn> NodeVisitor<'a> for AnnotationVisitor<'a, A> {
         }
     }
 
-    fn visit_up(&mut self, node: &'a Expression) -> VortexResult<TraversalOrder> {
+    fn visit_up(&mut self, node: &'a N) -> VortexResult<TraversalOrder> {
         if !self.propagate_up {
             return Ok(TraversalOrder::Continue);
         }
+        let child_annotations = node.iter_children(|children| {
+            children
+                .filter_map(|child| self.annotations.get(child).cloned())
+                .collect::<Vec<_>>()
+        });
+
+        let annotations = self.annotations.entry(node).or_default();
+        child_annotations
+            .into_iter()
+            .for_each(|ps| annotations.extend(ps.iter().cloned()));
+
+        Ok(TraversalOrder::Continue)
+    }
+}
+
+struct BoundAnnotationVisitor<A>
+where
+    A: AnnotationFn<BoundExpression>,
+{
+    annotations: BoundAnnotations<A::Annotation>,
+    annotate: A,
+    propagate_up: bool,
+}
+
+impl<'a, A> NodeVisitor<'a> for BoundAnnotationVisitor<A>
+where
+    A: AnnotationFn<BoundExpression>,
+{
+    type NodeTy = BoundExpression;
+
+    fn visit_down(&mut self, node: &'a Self::NodeTy) -> VortexResult<TraversalOrder> {
+        let annotations = (self.annotate)(node);
+        if annotations.is_empty() {
+            return Ok(TraversalOrder::Continue);
+        }
+
+        self.annotations
+            .entry(ExactBoundExpr(node.clone()))
+            .or_default()
+            .extend(annotations);
+        Ok(TraversalOrder::Skip)
+    }
+
+    fn visit_up(&mut self, node: &'a Self::NodeTy) -> VortexResult<TraversalOrder> {
+        if !self.propagate_up {
+            return Ok(TraversalOrder::Continue);
+        }
+
         let child_annotations = node
             .children()
             .iter()
-            .filter_map(|c| self.annotations.get(c).cloned())
+            .filter_map(|child| {
+                self.annotations
+                    .get(&ExactBoundExpr(child.clone()))
+                    .cloned()
+            })
             .collect::<Vec<_>>();
-
-        let annotations = self.annotations.entry(node).or_default();
+        let annotations = self
+            .annotations
+            .entry(ExactBoundExpr(node.clone()))
+            .or_default();
         child_annotations
             .into_iter()
-            .for_each(|ps| annotations.extend(ps.iter().cloned()));
+            .for_each(|child| annotations.extend(child));
 
         Ok(TraversalOrder::Continue)
     }

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

@@ -6,6 +6,7 @@ use vortex_utils::aliases::hash_set::HashSet;
 
 use crate::dtype::FieldName;
 use crate::dtype::StructFields;
+use crate::expr::BoundExpression;
 use crate::expr::Expression;
 use crate::expr::analysis::AnnotationFn;
 use crate::expr::analysis::Annotations;
@@ -43,7 +44,7 @@ pub type FieldAccesses<'a> = Annotations<'a, FieldName>;
 /// - The full expression has free fields `{a, d}` (not `b`, only top-level fields are tracked).
 pub fn make_free_field_annotator(
     scope: &StructFields,
-) -> impl AnnotationFn<Annotation = FieldName> {
+) -> impl AnnotationFn<Expression, Annotation = FieldName> {
     move |expr: &Expression| {
         if let Some(selection) = expr.as_opt::<Select>() {
             if expr.child(0).is::<Root>() {
@@ -65,6 +66,33 @@ pub fn make_free_field_annotator(
     }
 }
 
+/// Returns the free top-level fields for bound expression nodes.
+pub fn make_bound_free_field_annotator(
+    scope: &StructFields,
+) -> impl AnnotationFn<BoundExpression, Annotation = FieldName> {
+    move |expr: &BoundExpression| {
+        let Some(scalar_fn) = expr.as_scalar() else {
+            return scope.names().iter().cloned().collect();
+        };
+
+        if let Some(selection) = scalar_fn.as_opt::<Select>() {
+            if expr.children()[0].is_root() {
+                return selection
+                    .normalize_to_included_fields(scope.names())
+                    .vortex_expect("Select fields must be valid for scope")
+                    .into_iter()
+                    .collect();
+            }
+        } else if let Some(field_name) = scalar_fn.as_opt::<GetItem>()
+            && expr.children()[0].is_root()
+        {
+            return vec![field_name.clone()];
+        }
+
+        vec![]
+    }
+}
+
 /// For all subexpressions in an expression, find the fields that are accessed directly from the
 /// scope, but not any fields in those fields
 /// e.g. scope = {a: {b: .., c: ..}, d: ..}, expr = root().a.b + root().d accesses {a,d} (not b).