vortex_array/expr/analysis/
annotation.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::hash::Hash;
5
6use vortex_error::VortexExpect;
7use vortex_error::VortexResult;
8use vortex_utils::aliases::hash_map::HashMap;
9use vortex_utils::aliases::hash_set::HashSet;
10
11use crate::expr::BoundExpression;
12use crate::expr::ExactBoundExpr;
13use crate::expr::traversal::NodeExt;
14use crate::expr::traversal::NodeVisitor;
15use crate::expr::traversal::TraversalOrder;
16
17/// A value that can annotate an expression node.
18pub trait Annotation: Clone + Hash + Eq {}
19
20impl<A> Annotation for A where A: Clone + Hash + Eq {}
21
22/// Computes annotations for an expression node, defaulting to bound expressions.
23pub trait AnnotationFn<N = BoundExpression>: Fn(&N) -> Vec<Self::Annotation> {
24    /// The annotation attached to a node.
25    type Annotation: Annotation;
26}
27
28impl<N, A, F> AnnotationFn<N> for F
29where
30    A: Annotation,
31    F: Fn(&N) -> Vec<A>,
32{
33    type Annotation = A;
34}
35
36/// Annotations keyed by bound-tree identity.
37///
38/// Identity keys avoid structurally hashing every node's dtype. That matters when a bound root
39/// carries a lazy schema whose structural hash would deserialize every field.
40pub type BoundAnnotations<A> = HashMap<ExactBoundExpr, HashSet<A>>;
41
42/// Annotate a bound expression and propagate each annotation to its ancestors.
43///
44/// This uses [`ExactBoundExpr`] keys to preserve the cheap identity semantics of a bound tree.
45pub fn descendent_bound_annotations<A>(
46    expr: &BoundExpression,
47    annotate: A,
48) -> BoundAnnotations<A::Annotation>
49where
50    A: AnnotationFn<BoundExpression>,
51{
52    bound_annotations(expr, annotate, true)
53}
54
55/// Annotate each bound-expression node without propagating annotations to its ancestors.
56///
57/// The returned map uses [`ExactBoundExpr`] keys so lookups do not structurally hash node dtypes.
58pub fn direct_bound_annotations<A>(
59    expr: &BoundExpression,
60    annotate: A,
61) -> BoundAnnotations<A::Annotation>
62where
63    A: AnnotationFn<BoundExpression>,
64{
65    bound_annotations(expr, annotate, false)
66}
67
68fn bound_annotations<A>(
69    expr: &BoundExpression,
70    annotate: A,
71    propagate_up: bool,
72) -> BoundAnnotations<A::Annotation>
73where
74    A: AnnotationFn<BoundExpression>,
75{
76    let mut visitor = BoundAnnotationVisitor {
77        annotations: Default::default(),
78        annotate,
79        propagate_up,
80    };
81    expr.accept(&mut visitor).vortex_expect("Infallible");
82    visitor.annotations
83}
84
85struct BoundAnnotationVisitor<A>
86where
87    A: AnnotationFn<BoundExpression>,
88{
89    annotations: BoundAnnotations<A::Annotation>,
90    annotate: A,
91    propagate_up: bool,
92}
93
94impl<'a, A> NodeVisitor<'a> for BoundAnnotationVisitor<A>
95where
96    A: AnnotationFn<BoundExpression>,
97{
98    type NodeTy = BoundExpression;
99
100    fn visit_down(&mut self, node: &'a Self::NodeTy) -> VortexResult<TraversalOrder> {
101        let annotations = (self.annotate)(node);
102        if annotations.is_empty() {
103            return Ok(TraversalOrder::Continue);
104        }
105
106        self.annotations
107            .entry(ExactBoundExpr(node.clone()))
108            .or_default()
109            .extend(annotations);
110        Ok(TraversalOrder::Skip)
111    }
112
113    fn visit_up(&mut self, node: &'a Self::NodeTy) -> VortexResult<TraversalOrder> {
114        if !self.propagate_up {
115            return Ok(TraversalOrder::Continue);
116        }
117
118        let child_annotations = node
119            .children()
120            .iter()
121            .filter_map(|child| {
122                self.annotations
123                    .get(&ExactBoundExpr(child.clone()))
124                    .cloned()
125            })
126            .collect::<Vec<_>>();
127        let annotations = self
128            .annotations
129            .entry(ExactBoundExpr(node.clone()))
130            .or_default();
131        child_annotations
132            .into_iter()
133            .for_each(|child| annotations.extend(child));
134
135        Ok(TraversalOrder::Continue)
136    }
137}
vortex_array/expr/analysis/annotation.rs

vortex_array/expr/analysis/
annotation.rs
