TreeNode: The Little Traversal Framework That Holds DataFusion's Optimizer Together

TL;DR: A query optimizer is a pile of tree rewrites. If every rewrite has to re-implement “visit my children, bottom-up, and let me change things”, you get a thousand lines of near-duplicate code and a thousand chances to get the traversal order subtly wrong. DataFusion factored that out into one trait, TreeNode, with a three-state recursion control (Continue / Jump / Stop), a Transformed<T> result that says whether anything actually changed, and — the clever bit — a rewrite implementation that does not clone subtrees it did not change.

DataFusion has four tree-shaped data structures: LogicalPlan, ExecutionPlan, Expr and PhysicalExpr. Every optimizer rule, every physical planner method, every expression simplifier needs to walk them, rewrite them, and know whether it changed anything.

If each of those wrote its own recursion, the codebase would be a minefield of subtly different traversals. Instead, they all implement one trait: datafusion/common/src/tree_node.rs, 2,495 lines.

The three categories

The trait’s documentation organises the API into three families:

1. “Inspecting” APIs to traverse a tree of &TreeNodes: apply, visit, exists.

2. “Transforming” APIs that traverse and consume a tree of TreeNodes producing possibly changed TreeNodes: transform, transform_up, transform_down, transform_down_up, and rewrite.

3. Internal APIs used to implement the TreeNode API: apply_children, and map_children.

and then a table of which traversal order pairs with which API:

Traversal Order	Inspecting	Transforming
top-down	apply, exists	transform_down
bottom-up		transform, transform_up
combined with separate f_down and f_up closures		transform_down_up
combined with f_down() and f_up() in an object	visit	rewrite

That table is the whole API design in six lines. A rule that wants to rewrite bottom-up calls transform_up; a rule that needs to see the children before it decides anything implements TreeNodeRewriter with both f_down and f_up.

TreeNodeRecursion: three states, not two

Most tree-walk APIs return a bool: “keep going” or “stop”. DataFusion has three states:

pub enum TreeNodeRecursion {
    /// Continue recursion with the next node.
    Continue,
    /// In top-down traversals, skip recursing into children but continue with
    /// the next node, which actually means pruning of the subtree.
    ///
    /// In bottom-up traversals, bypass calling bottom-up closures till the next
    /// leaf node.
    ///
    /// In combined traversals, if it is the `f_down` (pre-order) phase, execution
    /// "jumps" to the next `f_up` (post-order) phase by shortcutting its children.
    /// If it is the `f_up` (post-order) phase, execution "jumps" to the next `f_down"
    /// (pre-order) phase by shortcutting its parent nodes until the first parent node
    /// having unvisited children path.
    Jump,
    /// Stop recursion.
    Stop,
}

Stop is the obvious one. Jump is the interesting one, and its behaviour is phase-dependent: in a top-down traversal it prunes the subtree; in a bottom-up traversal it skips the up-calls until the next leaf; in a combined traversal it short-circuits to the other phase.

The mechanism is two small combinators:

impl TreeNodeRecursion {
    /// Continues visiting nodes with `f` depending on the current `TreeNodeRecursion`
    /// value and the fact that `f` is visiting the current node's children.
    pub fn visit_children<F: FnOnce() -> Result<TreeNodeRecursion>>(
        self,
        f: F,
    ) -> Result<TreeNodeRecursion> {
        match self {
            TreeNodeRecursion::Continue => f(),
            TreeNodeRecursion::Jump => Ok(TreeNodeRecursion::Continue),
            TreeNodeRecursion::Stop => Ok(self),
        }
    }

    /// Continues visiting nodes with `f` depending on the current `TreeNodeRecursion`
    /// value and the fact that `f` is visiting the current node's sibling.
    pub fn visit_sibling<F: FnOnce() -> Result<TreeNodeRecursion>>(
        self,
        f: F,
    ) -> Result<TreeNodeRecursion> {
        match self {
            TreeNodeRecursion::Continue | TreeNodeRecursion::Jump => f(),
            ...

visit_children and visit_sibling are the two places where “the same recursion value” means two different things, and splitting them into two methods makes the phase-dependence explicit instead of implicit.

Jump is genuinely needed. Consider the decorrelator from the previous article: when it reaches a Subquery node, the subquery is a scope boundary — its outer references belong to a different decorrelation level. The decorrelator wants to skip the subtree’s children but keep walking:

// Subquery nodes are scope boundaries for correlation. A nested
// Subquery's outer references belong to a different decorrelation
// level and must not be pulled up into the current scope.
LogicalPlan::Subquery(_) => {
    Ok(Transformed::new(plan, false, TreeNodeRecursion::Jump))
}

With a two-state API, the only options would be “prune the whole subtree” (losing the siblings) or “keep descending” (wrong semantics). Jump gives the third option.

Transformed<T>: did anything change?

Every rewrite returns a Transformed<T>:

pub struct Transformed<T> {
    pub data: T,
    pub transformed: bool,
    pub tnr: TreeNodeRecursion,
}

The transformed flag is the load-bearing field. It answers “did I actually change anything?”, and the framework uses it to decide whether to keep the new subtree or throw it away and use the old one. This is the mechanism behind the “avoid cloning” claim.

The convenience constructors are worth noting because they encode the idioms:

/// Wrapper for transformed data with `Continue` statement.
pub fn yes(data: T) -> Self {
    Self::new(data, true, TreeNodeRecursion::Continue)
}

/// Wrapper for unchanged data with `Continue` statement.
pub fn no(data: T) -> Self {
    Self::new(data, false, TreeNodeRecursion::Continue)
}

Transformed::no(x) is what a rule returns when it looked at a node and decided to leave it alone. The framework then keeps the original node rather than the one the rule produced. For a rule that inspects a 50-node plan and changes one node, that means 49 subtrees are never cloned.

The transform_children machinery is where that decision is made:

f(node)?
    .transform_children(|n| n.map_children(|c| transform_down_impl(c, f)))?
    .transform_parent(f_up)

transform_children only replaces a child if the child’s own Transformed::transformed is true; otherwise it puts the original child back. Combined with map_children’s laziness — the closure is only called for children that exist — the effect is that an unchanged subtree costs a traversal and nothing else.

The rewrite API: one object, two phases

For rules that need both directions, TreeNodeRewriter is the interface:

fn rewrite<R: TreeNodeRewriter<Node = Self>>(
    self,
    rewriter: &mut R,
) -> Result<Transformed<Self>> {
    handle_transform_recursion!(rewriter.f_down(self), |c| c.rewrite(rewriter), |n| {
        rewriter.f_up(n)
    })
}

with the macro at the top of the file:

macro_rules! handle_transform_recursion {
    ($F_DOWN:expr, $F_CHILD:expr, $F_UP:expr) => {{
        $F_DOWN?
            .transform_children(|n| n.map_children($F_CHILD))?
            .transform_parent($F_UP)
    }};
}

That is the entire traversal engine in five lines, and the macro exists so that the Transformed-aware plumbing is written once.

The f_down / f_up split is the classic “bottom-up rewriting” pattern in a form that does not need Rc<RefCell<>> or interior mutability. The rewriter object holds the state (the join filters it has collected, the columns it has seen) across the whole traversal, and Rust’s borrow checker is happy because the state lives in &mut R, not in the tree.

The PullUpCorrelatedExpr from the previous article is exactly this shape: a TreeNodeRewriter with f_down and f_up and about a dozen state fields.

Plans that carry payloads

One more wrinkle: sometimes you want to walk a plan and some associated data — a DFSchema, an EquivalenceProperties, a PlanContext. The TreeNode framework handles this with a generic container:

This trait is implemented for plans (ExecutionPlan, LogicalPlan) and expression trees (PhysicalExpr, Expr) as well as Plan+Payload combinations PlanContext and ExprContext.

So PlanContext is itself a TreeNode, and the same rewrite machinery works on it. That means a rule written against LogicalPlan can be applied to a PlanContext unchanged — no separate implementation, no subtle divergence in behaviour between the two.

The catch

• The recursion is #[cfg_attr(feature = "recursive_protection", recursive::recursive)]. A deeply nested plan (a chain of a thousand unions, or a pathological expression) can blow the stack. The recursive_protection feature exists to swap in the recursive crate’s trampolined version. It is opt-in because it costs something — but the fact that the same code can be either recursive or trampolined is a nice property of expressing the traversal as a trait method.
• The API surface is large. apply, exists, visit, transform, transform_up, transform_down, transform_down_up, rewrite, apply_children, map_children — ten methods, with subtle differences in traversal order and in whether the closure gets an owned or a borrowed node. Choosing the wrong one is easy.
• Transformed is a three-field struct that callers can misuse. Constructing Transformed::new(data, false, TreeNodeRecursion::Stop) — claiming “not transformed” while also stopping the walk — is contradictory, and the type does not prevent it.
• The “avoid cloning” property is a promise about the implementation, not the signature. transform consumes self and returns Transformed<Self>; there is no way to see from the signature that an unchanged subtree is not cloned. It is an invariant of the implementation, maintained by transform_children and map_children.

Why it matters

The TreeNode trait is not a utility — it is the reason DataFusion’s optimizer can have 26 rules, a physical planner with dozens of node constructors, and an expression simplifier, without the codebase exploding. Every one of them is written as “here is what I do to a node”, and the framework handles the traversal, the change detection, and the copy avoidance.

And the design has an interesting property: the same three-state recursion control and the same Transformed protocol work for inspecting trees (apply), rewriting trees (transform_up), and rewriting trees with rules that carry state (rewrite). One idea, three shapes, one implementation.