DataFusion Rewrites Its Plan Trees In Place: Arc::make_mut and the End of the Clone-Collect-Rebuild Cycle

TL;DR: TreeNode::rewrite is a lovely API and a bad fit for an optimizer that runs 26 rules over a plan tree. Every rule that says “no change” still pays for a full clone-and-rebuild of the tree, because the ownership-based API cannot know in advance whether a node will change. DataFusion added a second, in-place traversal built on Arc::make_mut: if nobody else shares a node, it is mutated for free; if they do, it is cloned exactly once. On top of that, the optimizer checks whether the plan contains any subqueries and, if not, uses a much cheaper traversal that does not need to visit expression subtrees at all.

Query optimizers are tree rewriters. You have a plan, you have a rule, the rule looks at a node and maybe replaces it. The naive Rust implementation of “maybe replaces it” is awkward, because the rule needs an owned node to return an owned node, and the tree is a graph of Arcs.

DataFusion’s TreeNode trait solves this with a consuming rewrite:

rust
fn rewrite_with_subqueries(
    &self,
    _f: &mut dyn FnMut(&LogicalPlan) -> Result<Transformed<LogicalPlan>>,
) -> Result<Transformed<LogicalPlan>>;

The rule gets an owned LogicalPlan, returns a Transformed<LogicalPlan>, and the framework rebuilds the tree around it. It is clean, and it is expensive: every node that a rule inspects and declines to change is still cloned, because the rewrite has to produce a new owned tree to reassemble.

With 26 rules and a plan of a few dozen nodes, that is a lot of churn — and the churn is proportional to the size of the plan, not to the amount of actual rewriting.

The in-place traversal

datafusion/optimizer/src/optimizer.rs contains a second traversal written specifically for the optimizer:

rust
/// Rewrites a plan tree in place using `Arc::make_mut` for
/// copy-on-write semantics on `Arc<LogicalPlan>` children.
///
/// This avoids the `Arc::unwrap_or_clone` + `Arc::new` cycle that the
/// ownership-based `TreeNode::rewrite` performs at every child node.
#[cfg_attr(feature = "recursive_protection", recursive::recursive)]
fn rewrite_plan_in_place(
    plan: &mut LogicalPlan,
    apply_order: ApplyOrder,
    rule: &dyn OptimizerRule,
    config: &dyn OptimizerConfig,
) -> Result<bool> {

Arc::make_mut is the whole trick. If the Arc has refcount 1, it hands out a &mut to the value with zero cost. If the refcount is greater than one, it clones the value once and hands out a mutable reference to the clone. Either way the caller gets a mutable slot and never has to write the match Arc::try_unwrap dance by hand.

The children accessor is a big match over the plan enum, and each arm is one line:

rust
fn map_children_mut<F: FnMut(&mut LogicalPlan) -> Result<bool>>(
    plan: &mut LogicalPlan,
    mut f: F,
) -> Result<bool> {
    Ok(match plan {
        LogicalPlan::Projection(Projection { input, .. })
        | LogicalPlan::Filter(Filter { input, .. })
        | LogicalPlan::Repartition(Repartition { input, .. })
        | LogicalPlan::Window(Window { input, .. })
        | LogicalPlan::Aggregate(Aggregate { input, .. })
        | LogicalPlan::Sort(Sort { input, .. })
        | LogicalPlan::Limit(Limit { input, .. })
        | LogicalPlan::SubqueryAlias(SubqueryAlias { input, .. })
        ...
        => f(Arc::make_mut(input))?,
        LogicalPlan::Join(Join { left, right, .. }) => {
            let l = f(Arc::make_mut(left))?;
            let r = f(Arc::make_mut(right))?;
            l || r
        }
        ...
    })
}

The f callback returns bool — “did you change this child?” — and the match combines them with ||. The whole tree walk returns a single bool saying whether anything changed, which is exactly the signal the optimizer needs to know whether to run another pass.

Bridging `&mut` to owned

There is a wrinkle. The OptimizerRule trait is defined in terms of owned plans:

rust
fn rewrite(&self, plan: LogicalPlan, config: &dyn OptimizerConfig)
    -> Result<Transformed<LogicalPlan>>;

The rule consumes the plan and returns a new one. To call it through a &mut LogicalPlan, the in-place traversal has to bridge the two, and the way it does it is a small piece of care:

rust
// f_down phase
let mut changed = false;
if apply_order == ApplyOrder::TopDown {
    // `rule.rewrite()` takes the plan by value, so bridge the `&mut` to an
    // owned value with `std::mem::take`. `LogicalPlan::default()` is a cheap
    // empty placeholder (shared empty schema, no allocation) and is
    // overwritten with the rule's output on the next line.
    let owned = std::mem::take(plan);
    let result = rule.rewrite(owned, config)?;
    *plan = result.data;
    changed |= result.transformed;
    ...
}

std::mem::take(plan) moves the plan out, leaving LogicalPlan::default() behind. default() is documented as cheap — a shared empty schema, no allocation — so the placeholder costs nothing. The ? on the rule’s result is safe because the placeholder is still a valid (if useless) plan, so an early return leaves the tree in a well-defined state rather than with a dangling hole.

The same trick appears in datafusion/common/src/tree_node.rs, at the container level:

rust
fn map_elements<F: FnMut(T) -> Result<Transformed<T>>>(
    mut self,
    f: F,
) -> Result<Transformed<Self>> {
    // Rewrite in place so the existing heap allocation can be reused.
    // `mem::take` hands the inner `C` to `f` while leaving
    // `C::default()` in the slot, so an unwinding drop finds a valid
    // `C` even if `f` panics or the `?` short-circuits.
    let inner = std::mem::take(&mut *self);
    Ok(inner.map_elements(f)?.update_data(|c| {
        *self = c;
        self
    }))
}

The comment is the interesting part: “an unwinding drop finds a valid C even if f panics or the ? short-circuits”. This is panic-safety by construction. If f panics halfway, the Box/Arc slot contains C::default() rather than a partially-initialised value, so the drop glue does the right thing. Rust’s “leave a valid value in place of what you took” pattern is usually discussed as a way to avoid Option; here it is used as a panic guard.

The Arc version adds the CoW step:

rust
fn map_elements<F: FnMut(T) -> Result<Transformed<T>>>(
    mut self,
    f: F,
) -> Result<Transformed<Self>> {
    // Rewrite in place using the same `mem::take` strategy as
    // `Box<C>::map_elements`. `Arc::make_mut` gives us exclusive
    // access (cloning `C` first if we were sharing), after which
    // `get_mut` is infallible.
    let inner = std::mem::take(Arc::make_mut(&mut self));
    Ok(inner.map_elements(f)?.update_data(|c| {
        *Arc::get_mut(&mut self).unwrap() = c;
        self
    }))
}

Note Arc::get_mut(&mut self).unwrap() — infallible because make_mut guaranteed exclusivity.

The subquery fast path

The in-place traversal is not the only optimisation. The optimizer decides which traversal to use based on whether the plan contains subqueries:

rust
let result = match rule.apply_order() {
    Some(apply_order) => {
        if has_subqueries {
            // Plans with subqueries need the full
            // rewrite_with_subqueries traversal to
            // recurse into subquery plans.
            new_plan.rewrite_with_subqueries(...)
        } else {
            // No subqueries: use in-place rewriting
            // with Arc::make_mut for zero-cost CoW on
            // children, avoiding Arc unwrap/rewrap.
            rewrite_plan_in_place(...)
        }
    }
    None => rule.rewrite(new_plan, config),
}

and the predicate is computed by a dedicated traversal:

rust
/// Returns true if the plan contains any subquery expressions
/// (EXISTS, IN subquery, scalar subquery, set comparison).
///
/// Used to determine whether the more expensive `rewrite_with_subqueries`
/// traversal is needed. When the plan has no subqueries, the cheaper
/// `rewrite` traversal is sufficient since all plan nodes are reachable
/// via direct children.
fn plan_has_subqueries(plan: &LogicalPlan) -> bool {

The reason is structural. rewrite_with_subqueries has to walk inside every expression looking for Expr::Exists, Expr::InSubquery, Expr::SetComparison and Expr::ScalarSubquery, each of which carries a nested LogicalPlan. A plan with no subqueries has no such nested plans, so the ordinary child-only traversal reaches everything.

And the check is re-done during the pass, not just once:

rust
if transformed {
    has_subqueries = plan_has_subqueries(&new_plan);
    log_plan(rule.name(), &new_plan);
}

Because a rule such as ScalarSubqueryToJoin can introduce subqueries (well, remove them) or change their shape. The comment in the optimizer explains:

Track subquery presence across the pass. Refresh after changed rules so decorrelation can move later rules onto the in-place path; that path refreshes parent schemas after child schemas change.

Note that the refresh only happens when a rule actually changed something (if transformed). An unchanged plan cannot have changed its subquery content, so the re-check would be pure waste.

The catch

Arc::make_mut is only free if nobody else holds a reference. If a plan node is shared (a CTE used twice, a cached plan, a plan held by an Arc in the session state), the first mutation clones it. That clone is a deep clone of the node’s contents but a shallow clone of its children — which is exactly the right granularity.
The error semantics are different, and the code says so. The in-place traversal documents:
On Err, *plan is left in an unspecified state and must not be used. Note this is different than consuming APIs such as TreeNode::rewrite where the original plan is freed and no longer available on error
The optimizer handles this by cloning the plan before each rule when skip_failed_rules is enabled:
rust
```
let prev_plan = options
    .optimizer
    .skip_failed_rules
    .then(|| new_plan.clone());
```
Note this clone is only paid when the option is on — the default path does not clone at all, because in-place rewriting leaves the plan recoverable in the normal case.
The Arc::make_mut approach does not generalise. The source is candid:
the Arc::make_mut approach does not generalize to the other tree types (Expr children are Boxed; PhysicalExpr/ExecutionPlan children are Arc<dyn _>, which Arc::make_mut cannot handle)
Arc::make_mut requires a concrete T, not a trait object. So the in-place trick works for LogicalPlan and not for the physical plan tree. That is a genuine limit of the technique, and the code documents it rather than pretending otherwise.
The map_children_mut match must be kept in sync with the enum. Adding a new LogicalPlan variant without adding an arm is a compile-time error only because the match is exhaustive — a good property, but it is the only thing keeping this correct.

Why it matters

Optimisation time is not free. A DataFusion deployment serving short interactive queries pays the optimiser cost on every request. The difference between “clone the plan 26 times” and “mutate the plan in place” is the difference between an optimiser that shows up in a latency histogram and one that does not.

And the same reasoning scales up: the physical planner, the rule engine, and the proto serialisation layer all have the same shape of problem, and each has grown its own local solution. The in-place traversal is the generalised version of that idea, applied where it pays the most.

The in-place traversal

Bridging &mut to owned

The subquery fast path

The catch

Why it matters

Bridging `&mut` to owned