vortex_array/expr/traversal/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Datafusion inspired tree traversal logic.
5//!
6//! Users should want to implement [`Node`] and potentially [`NodeContainer`].
7
8mod fold;
9mod visitor;
10
11use std::marker::PhantomData;
12use std::sync::Arc;
13
14pub use fold::FoldDownContext;
15pub use fold::FoldUp;
16pub use fold::NodeFolderContext;
17pub use visitor::pre_order_visit_down;
18use vortex_error::VortexResult;
19
20use crate::expr::BoundExpression;
21use crate::expr::Expression;
22
23/// Signal to control a traversal's flow
24#[derive(Debug, Clone, PartialEq, Eq)]
25pub enum TraversalOrder {
26    /// In a top-down traversal, skip visiting the children of the current node.
27    /// In the bottom-up phase of the traversal, skip the next step. Either skipping the children of the node,
28    /// moving to its next sibling, or skipping its parent once the children are traversed.
29    Skip,
30    /// Stop visiting any more nodes in the traversal.
31    Stop,
32    /// Continue with the traversal as expected.
33    Continue,
34}
35
36impl TraversalOrder {
37    /// If directed to, continue to visit nodes by running `f`, which should apply on the node's children.
38    pub fn visit_children<F: FnOnce() -> VortexResult<TraversalOrder>>(
39        self,
40        f: F,
41    ) -> VortexResult<TraversalOrder> {
42        match self {
43            Self::Skip => Ok(TraversalOrder::Continue),
44            Self::Stop => Ok(self),
45            Self::Continue => f(),
46        }
47    }
48
49    /// If directed to, continue to visit nodes by running `f`, which should apply on the node's parent.
50    pub fn visit_parent<F: FnOnce() -> VortexResult<TraversalOrder>>(
51        self,
52        f: F,
53    ) -> VortexResult<TraversalOrder> {
54        match self {
55            Self::Continue => f(),
56            Self::Skip | Self::Stop => Ok(self),
57        }
58    }
59}
60
61#[derive(Debug, Clone)]
62pub struct Transformed<T> {
63    /// Value that was being rewritten.
64    pub value: T,
65    /// Controls the flow of rewriting, see [`TraversalOrder`] for more details.
66    pub order: TraversalOrder,
67    /// Was the value changed during rewriting.
68    pub changed: bool,
69}
70
71impl<T> Transformed<T> {
72    pub fn yes(value: T) -> Self {
73        Self {
74            value,
75            order: TraversalOrder::Continue,
76            changed: true,
77        }
78    }
79
80    pub fn no(value: T) -> Self {
81        Self {
82            value,
83            order: TraversalOrder::Continue,
84            changed: false,
85        }
86    }
87
88    pub fn into_inner(self) -> T {
89        self.value
90    }
91
92    /// Apply a function to `value`, changing it without changing the `changed` field.
93    pub fn map<O, F: FnOnce(T) -> O>(self, f: F) -> Transformed<O> {
94        Transformed {
95            value: f(self.value),
96            order: self.order,
97            changed: self.changed,
98        }
99    }
100}
101
102pub trait NodeVisitor<'a> {
103    type NodeTy: Node;
104
105    fn visit_down(&mut self, node: &'a Self::NodeTy) -> VortexResult<TraversalOrder> {
106        _ = node;
107        Ok(TraversalOrder::Continue)
108    }
109
110    fn visit_up(&mut self, node: &'a Self::NodeTy) -> VortexResult<TraversalOrder> {
111        _ = node;
112        Ok(TraversalOrder::Continue)
113    }
114}
115
116pub trait NodeRewriter: Sized {
117    type NodeTy: Node;
118
119    fn visit_down(&mut self, node: Self::NodeTy) -> VortexResult<Transformed<Self::NodeTy>> {
120        Ok(Transformed::no(node))
121    }
122
123    fn visit_up(&mut self, node: Self::NodeTy) -> VortexResult<Transformed<Self::NodeTy>> {
124        Ok(Transformed::no(node))
125    }
126}
127
128pub trait Node: Sized + Clone {
129    /// Walk the node's children by applying `f` to them.
130    ///
131    /// This is a lower level API that other functions rely on for their implementation.
132    fn apply_children<'a, F: FnMut(&'a Self) -> VortexResult<TraversalOrder>>(
133        &'a self,
134        f: F,
135    ) -> VortexResult<TraversalOrder>;
136
137    /// Rewrite the node's children by applying `f` to them.
138    ///
139    /// This is a lower level API that other functions rely on for their implementation.
140    fn map_children<F: FnMut(Self) -> VortexResult<Transformed<Self>>>(
141        self,
142        f: F,
143    ) -> VortexResult<Transformed<Self>>;
144
145    /// This is a lower level API that other functions rely on for their implementation.
146    fn iter_children<T>(&self, f: impl FnOnce(&mut dyn Iterator<Item = &Self>) -> T) -> T;
147
148    /// This is a lower level API that other functions rely on for their implementation.
149    fn children_count(&self) -> usize;
150}
151
152pub trait NodeExt: Node {
153    /// Walk the tree in pre-order (top-down) way, rewriting it as it goes.
154    fn rewrite<R: NodeRewriter<NodeTy = Self>>(
155        self,
156        rewriter: &mut R,
157    ) -> VortexResult<Transformed<Self>> {
158        let mut transformed = rewriter.visit_down(self)?;
159
160        let transformed = match transformed.order {
161            TraversalOrder::Stop => Ok(transformed),
162            TraversalOrder::Skip => {
163                transformed.order = TraversalOrder::Continue;
164                Ok(transformed)
165            }
166            TraversalOrder::Continue => transformed
167                .value
168                .map_children(|c| c.rewrite(rewriter))
169                .map(|mut t| {
170                    t.changed |= transformed.changed;
171                    t
172                }),
173        }?;
174
175        match transformed.order {
176            TraversalOrder::Stop | TraversalOrder::Skip => Ok(transformed),
177            TraversalOrder::Continue => {
178                let mut up_rewrite = rewriter.visit_up(transformed.value)?;
179                up_rewrite.changed |= transformed.changed;
180                Ok(up_rewrite)
181            }
182        }
183    }
184
185    /// A pre-order (top-down) traversal.
186    fn accept<'a, V: NodeVisitor<'a, NodeTy = Self>>(
187        &'a self,
188        visitor: &mut V,
189    ) -> VortexResult<TraversalOrder> {
190        visitor
191            .visit_down(self)?
192            .visit_children(|| self.apply_children(|c| c.accept(visitor)))?
193            .visit_parent(|| visitor.visit_up(self))
194    }
195
196    /// A pre-order transformation
197    fn transform_down<F: FnMut(Self) -> VortexResult<Transformed<Self>>>(
198        self,
199        f: F,
200    ) -> VortexResult<Transformed<Self>> {
201        let mut rewriter = FnRewriter {
202            f,
203            _data: PhantomData,
204        };
205
206        self.rewrite(&mut rewriter)
207    }
208
209    /// applies the `NodeFolderContext` to the Node tree, with an initial `Context`.
210    fn fold_context<R, F: NodeFolderContext<NodeTy = Self, Result = R>>(
211        self,
212        ctx: &F::Context,
213        folder: &mut F,
214    ) -> VortexResult<FoldUp<R>> {
215        let transformed = folder.visit_down(ctx, &self)?;
216        let ctx = match transformed {
217            FoldDownContext::Continue(ctx) => ctx,
218            FoldDownContext::Skip(r) => return Ok(FoldUp::Continue(r)),
219            FoldDownContext::Stop(r) => return Ok(FoldUp::Stop(r)),
220        };
221
222        let mut children = Vec::with_capacity(self.children_count());
223        let mut stop_result = None;
224        self.iter_children(|children_iter| -> VortexResult<()> {
225            for c in children_iter {
226                let t = c.clone().fold_context(&ctx, folder)?;
227                match t {
228                    FoldUp::Stop(r) => {
229                        stop_result = Some(r);
230                        return Ok(());
231                    }
232                    FoldUp::Continue(r) => {
233                        children.push(r);
234                    }
235                }
236            }
237            Ok(())
238        })?;
239
240        if let Some(result) = stop_result {
241            return Ok(FoldUp::Stop(result));
242        }
243
244        folder.visit_up(self, &ctx, children)
245    }
246}
247
248impl<T: Node> NodeExt for T {}
249
250struct FnRewriter<F, T> {
251    f: F,
252    _data: PhantomData<T>,
253}
254
255impl<F, T> NodeRewriter for FnRewriter<F, T>
256where
257    T: Node,
258    F: FnMut(T) -> VortexResult<Transformed<T>>,
259{
260    type NodeTy = T;
261
262    fn visit_down(&mut self, node: Self::NodeTy) -> VortexResult<Transformed<Self::NodeTy>> {
263        (self.f)(node)
264    }
265}
266
267/// A container holding a [`Node`]'s children, which a function can be applied (or mapped) to.
268///
269/// The trait is also implemented to container types in order to make implementing [`Node::map_children`]
270/// and [`Node::apply_children`] easier.
271pub trait NodeContainer<'a, T: 'a>: Sized {
272    /// Applies `f` to all elements of the container, accepting them by reference
273    fn apply_elements<F: FnMut(&'a T) -> VortexResult<TraversalOrder>>(
274        &'a self,
275        f: F,
276    ) -> VortexResult<TraversalOrder>;
277
278    /// Consumes all the children of the node, replacing them with the result of `f`.
279    fn map_elements<F: FnMut(T) -> VortexResult<Transformed<T>>>(
280        self,
281        f: F,
282    ) -> VortexResult<Transformed<Self>>;
283}
284
285pub trait NodeRefContainer<'a, T: 'a>: Sized {
286    fn apply_ref_elements<F: FnMut(&'a T) -> VortexResult<TraversalOrder>>(
287        &self,
288        f: F,
289    ) -> VortexResult<TraversalOrder>;
290}
291
292impl<'a, T: 'a, C: NodeContainer<'a, T>> NodeRefContainer<'a, T> for &'a [C] {
293    fn apply_ref_elements<F: FnMut(&'a T) -> VortexResult<TraversalOrder>>(
294        &self,
295        mut f: F,
296    ) -> VortexResult<TraversalOrder> {
297        let mut order = TraversalOrder::Continue;
298
299        for c in *self {
300            order = c.apply_elements(&mut f)?;
301            match order {
302                TraversalOrder::Continue | TraversalOrder::Skip => {}
303                TraversalOrder::Stop => return Ok(TraversalOrder::Stop),
304            }
305        }
306
307        Ok(order)
308    }
309}
310
311impl<'a, T: 'a, C: NodeContainer<'a, T>> NodeContainer<'a, T> for Box<C> {
312    fn apply_elements<F: FnMut(&'a T) -> VortexResult<TraversalOrder>>(
313        &'a self,
314        f: F,
315    ) -> VortexResult<TraversalOrder> {
316        self.as_ref().apply_elements(f)
317    }
318
319    fn map_elements<F: FnMut(T) -> VortexResult<Transformed<T>>>(
320        self,
321        f: F,
322    ) -> VortexResult<Transformed<Box<C>>> {
323        Ok((*self).map_elements(f)?.map(Box::new))
324    }
325}
326
327impl<'a, T, C> NodeContainer<'a, T> for Arc<C>
328where
329    T: 'a,
330    C: NodeContainer<'a, T> + Clone,
331{
332    fn apply_elements<F: FnMut(&'a T) -> VortexResult<TraversalOrder>>(
333        &'a self,
334        f: F,
335    ) -> VortexResult<TraversalOrder> {
336        self.as_ref().apply_elements(f)
337    }
338
339    fn map_elements<F: FnMut(T) -> VortexResult<Transformed<T>>>(
340        self,
341        f: F,
342    ) -> VortexResult<Transformed<Arc<C>>> {
343        Ok(Arc::unwrap_or_clone(self).map_elements(f)?.map(Arc::new))
344    }
345}
346
347impl<'a, T: 'a, C: NodeContainer<'a, T>> NodeContainer<'a, T> for [C; 2] {
348    fn apply_elements<F: FnMut(&'a T) -> VortexResult<TraversalOrder>>(
349        &'a self,
350        mut f: F,
351    ) -> VortexResult<TraversalOrder> {
352        let [lhs, rhs] = self;
353        match lhs.apply_elements(&mut f)? {
354            TraversalOrder::Skip | TraversalOrder::Continue => rhs.apply_elements(&mut f),
355            TraversalOrder::Stop => Ok(TraversalOrder::Stop),
356        }
357    }
358
359    fn map_elements<F: FnMut(T) -> VortexResult<Transformed<T>>>(
360        self,
361        mut f: F,
362    ) -> VortexResult<Transformed<[C; 2]>> {
363        let [lhs, rhs] = self;
364        let transformed = lhs.map_elements(&mut f)?;
365        match transformed.order {
366            TraversalOrder::Skip | TraversalOrder::Continue => {
367                let mut t = rhs.map_elements(&mut f)?;
368                t.changed |= transformed.changed;
369                Ok(t.map(|new_lhs| [new_lhs, transformed.value]))
370            }
371            TraversalOrder::Stop => Ok(transformed.map(|new_lhs| [new_lhs, rhs])),
372        }
373    }
374}
375
376impl<'a, T: 'a, C: NodeContainer<'a, T>> NodeContainer<'a, T> for Vec<C> {
377    fn apply_elements<F: FnMut(&'a T) -> VortexResult<TraversalOrder>>(
378        &'a self,
379        mut f: F,
380    ) -> VortexResult<TraversalOrder> {
381        let mut order = TraversalOrder::Continue;
382
383        for c in self {
384            order = c.apply_elements(&mut f)?;
385            match order {
386                TraversalOrder::Continue | TraversalOrder::Skip => {}
387                TraversalOrder::Stop => return Ok(TraversalOrder::Stop),
388            }
389        }
390
391        Ok(order)
392    }
393
394    fn map_elements<F: FnMut(T) -> VortexResult<Transformed<T>>>(
395        self,
396        mut f: F,
397    ) -> VortexResult<Transformed<Self>> {
398        let mut order = TraversalOrder::Continue;
399        let mut changed = false;
400
401        let value = self
402            .into_iter()
403            .map(|c| match order {
404                TraversalOrder::Continue | TraversalOrder::Skip => {
405                    c.map_elements(&mut f).map(|result| {
406                        order = result.order;
407                        changed |= result.changed;
408                        result.value
409                    })
410                }
411                TraversalOrder::Stop => Ok(c),
412            })
413            .collect::<VortexResult<Vec<_>>>()?;
414
415        Ok(Transformed {
416            value,
417            order,
418            changed,
419        })
420    }
421}
422
423impl<'a> NodeContainer<'a, Self> for Expression {
424    fn apply_elements<F: FnMut(&'a Self) -> VortexResult<TraversalOrder>>(
425        &'a self,
426        mut f: F,
427    ) -> VortexResult<TraversalOrder> {
428        f(self)
429    }
430
431    fn map_elements<F: FnMut(Self) -> VortexResult<Transformed<Self>>>(
432        self,
433        mut f: F,
434    ) -> VortexResult<Transformed<Self>> {
435        f(self)
436    }
437}
438
439impl Node for Expression {
440    fn apply_children<'a, F: FnMut(&'a Self) -> VortexResult<TraversalOrder>>(
441        &'a self,
442        mut f: F,
443    ) -> VortexResult<TraversalOrder> {
444        self.children().apply_ref_elements(&mut f)
445    }
446
447    fn map_children<F: FnMut(Self) -> VortexResult<Transformed<Self>>>(
448        self,
449        f: F,
450    ) -> VortexResult<Transformed<Self>> {
451        let transformed = self.children().to_vec().map_elements(f)?;
452
453        if transformed.changed {
454            Ok(Transformed {
455                value: self.with_children(transformed.value)?,
456                order: transformed.order,
457                changed: true,
458            })
459        } else {
460            Ok(Transformed::no(self))
461        }
462    }
463
464    fn iter_children<T>(&self, f: impl FnOnce(&mut dyn Iterator<Item = &Self>) -> T) -> T {
465        f(&mut self.children().iter())
466    }
467
468    fn children_count(&self) -> usize {
469        self.children().len()
470    }
471}
472
473impl Node for BoundExpression {
474    fn apply_children<'a, F: FnMut(&'a Self) -> VortexResult<TraversalOrder>>(
475        &'a self,
476        mut f: F,
477    ) -> VortexResult<TraversalOrder> {
478        let BoundExpression::Scalar { children, .. } = self else {
479            return Ok(TraversalOrder::Continue);
480        };
481
482        for child in children.iter() {
483            match f(child)? {
484                TraversalOrder::Continue | TraversalOrder::Skip => {}
485                TraversalOrder::Stop => return Ok(TraversalOrder::Stop),
486            }
487        }
488
489        Ok(TraversalOrder::Continue)
490    }
491
492    fn map_children<F: FnMut(Self) -> VortexResult<Transformed<Self>>>(
493        self,
494        mut f: F,
495    ) -> VortexResult<Transformed<Self>> {
496        let BoundExpression::Scalar { children, .. } = &self else {
497            return Ok(Transformed::no(self));
498        };
499
500        let mut order = TraversalOrder::Continue;
501        // Stays `None` until a child actually changes. Most nodes of a rewritten tree are
502        // untouched.
503        let mut rewritten: Option<Vec<Self>> = None;
504
505        for (index, child) in children.iter().enumerate() {
506            let value = match order {
507                TraversalOrder::Continue | TraversalOrder::Skip => {
508                    let result = f(child.clone())?;
509                    order = result.order;
510                    if result.changed && rewritten.is_none() {
511                        let mut prefix = Vec::with_capacity(children.len());
512                        prefix.extend_from_slice(&children[..index]);
513                        rewritten = Some(prefix);
514                    }
515                    result.value
516                }
517                TraversalOrder::Stop => child.clone(),
518            };
519
520            if let Some(rewritten) = &mut rewritten {
521                rewritten.push(value);
522            }
523        }
524
525        match rewritten {
526            Some(rewritten) => Ok(Transformed {
527                value: self.with_children(rewritten)?,
528                order,
529                changed: true,
530            }),
531            None => Ok(Transformed::no(self)),
532        }
533    }
534
535    fn iter_children<T>(&self, f: impl FnOnce(&mut dyn Iterator<Item = &Self>) -> T) -> T {
536        match self {
537            BoundExpression::Scalar { children, .. } => f(&mut children.iter()),
538            BoundExpression::Root { .. } => f(&mut std::iter::empty()),
539        }
540    }
541
542    fn children_count(&self) -> usize {
543        match self {
544            BoundExpression::Scalar { children, .. } => children.len(),
545            BoundExpression::Root { .. } => 0,
546        }
547    }
548}
549
550#[cfg(test)]
551mod tests {
552    use vortex_error::VortexResult;
553    use vortex_utils::aliases::hash_set::HashSet;
554
555    use super::NodeExt;
556    use super::NodeRewriter;
557    use super::NodeVisitor;
558    use super::Transformed;
559    use super::TraversalOrder;
560    use super::visitor::pre_order_visit_down;
561    use crate::expr::Expression;
562    use crate::expr::and;
563    use crate::expr::col;
564    use crate::expr::eq;
565    use crate::expr::is_root;
566    use crate::expr::lit;
567    use crate::expr::not_eq;
568    use crate::expr::root;
569    use crate::scalar_fn::fns::binary::Binary;
570    use crate::scalar_fn::fns::get_item::GetItem;
571    use crate::scalar_fn::fns::literal::Literal;
572    use crate::scalar_fn::fns::operators::Operator;
573
574    #[derive(Default)]
575    pub struct ExprLitCollector<'a>(pub Vec<&'a Expression>);
576
577    impl<'a> NodeVisitor<'a> for ExprLitCollector<'a> {
578        type NodeTy = Expression;
579
580        fn visit_down(&mut self, node: &'a Expression) -> VortexResult<TraversalOrder> {
581            if node.is::<Literal>() {
582                self.0.push(node)
583            }
584            Ok(TraversalOrder::Continue)
585        }
586
587        fn visit_up(&mut self, _node: &'a Expression) -> VortexResult<TraversalOrder> {
588            Ok(TraversalOrder::Continue)
589        }
590    }
591
592    fn expr_col_to_lit_transform(
593        node: Expression,
594        idx: &mut i32,
595    ) -> VortexResult<Transformed<Expression>> {
596        if node.is::<GetItem>() {
597            let lit_id = *idx;
598            *idx += 1;
599            Ok(Transformed::yes(lit(lit_id)))
600        } else {
601            Ok(Transformed::no(node))
602        }
603    }
604
605    #[derive(Default)]
606    pub struct SkipDownRewriter;
607
608    impl NodeRewriter for SkipDownRewriter {
609        type NodeTy = Expression;
610
611        fn visit_down(&mut self, node: Self::NodeTy) -> VortexResult<Transformed<Self::NodeTy>> {
612            Ok(Transformed {
613                value: node,
614                order: TraversalOrder::Skip,
615                changed: false,
616            })
617        }
618
619        fn visit_up(&mut self, _node: Self::NodeTy) -> VortexResult<Transformed<Self::NodeTy>> {
620            Ok(Transformed::yes(root()))
621        }
622    }
623
624    #[test]
625    fn expr_deep_visitor_test() {
626        let col1: Expression = col("col1");
627        let lit1 = lit(1);
628        let expr = eq(col1, lit1);
629        let lit2 = lit(2);
630        let expr = and(expr, lit2);
631        let mut printer = ExprLitCollector::default();
632        expr.accept(&mut printer).unwrap();
633        assert_eq!(printer.0.len(), 2);
634    }
635
636    #[test]
637    fn expr_deep_mut_visitor_test() {
638        let col1: Expression = col("col1");
639        let col2: Expression = col("col2");
640        let expr = eq(col1, col2);
641        let lit2 = lit(2);
642        let expr = and(expr, lit2);
643
644        let mut idx = 0_i32;
645        let new = expr
646            .transform_down(|node| expr_col_to_lit_transform(node, &mut idx))
647            .unwrap();
648        assert!(new.changed);
649
650        let expr = new.value;
651
652        let mut printer = ExprLitCollector::default();
653        expr.accept(&mut printer).unwrap();
654        assert_eq!(printer.0.len(), 3);
655    }
656
657    #[test]
658    fn expr_skip_test() {
659        let col1: Expression = col("col1");
660        let col2: Expression = col("col2");
661        let expr1 = eq(col1, col2);
662        let col3: Expression = col("col3");
663        let col4: Expression = col("col4");
664        let expr2 = not_eq(col3, col4);
665        let expr = and(expr1, expr2);
666
667        let mut nodes = Vec::new();
668        pre_order_visit_down(&expr, |node: &Expression| {
669            if node.is::<GetItem>() {
670                nodes.push(node)
671            }
672            if let Some(operator) = node.as_opt::<Binary>()
673                && *operator == Operator::Eq
674            {
675                return Ok(TraversalOrder::Skip);
676            }
677            Ok(TraversalOrder::Continue)
678        })
679        .unwrap();
680
681        let nodes: HashSet<Expression> = HashSet::from_iter(nodes.into_iter().cloned());
682        assert_eq!(nodes, HashSet::from_iter([col("col3"), col("col4")]));
683    }
684
685    #[test]
686    fn expr_stop_test() {
687        let col1: Expression = col("col1");
688        let col2: Expression = col("col2");
689        let expr1 = eq(col1, col2);
690        let col3: Expression = col("col3");
691        let col4: Expression = col("col4");
692        let expr2 = not_eq(col3, col4);
693        let expr = and(expr1, expr2);
694
695        let mut nodes = Vec::new();
696        pre_order_visit_down(&expr, |node: &Expression| {
697            if node.is::<GetItem>() {
698                nodes.push(node)
699            }
700            if let Some(operator) = node.as_opt::<Binary>()
701                && *operator == Operator::Eq
702            {
703                return Ok(TraversalOrder::Stop);
704            }
705            Ok(TraversalOrder::Continue)
706        })
707        .unwrap();
708
709        assert!(nodes.is_empty());
710    }
711
712    #[test]
713    fn expr_skip_down_visit_up() {
714        let col = col("col");
715
716        let mut visitor = SkipDownRewriter;
717        let result = col.rewrite(&mut visitor).unwrap();
718
719        assert!(result.changed);
720        assert!(is_root(&result.value));
721    }
722}
vortex_array/expr/traversal/mod.rs

vortex_array/expr/traversal/
mod.rs
