alamb commented on code in PR #22298:
URL: https://github.com/apache/datafusion/pull/22298#discussion_r3276226221
##########
datafusion/expr/src/logical_plan/tree_node.rs:
##########
@@ -402,6 +402,14 @@ impl LogicalPlan {
/// When the refcount is >1, it clones the inner value first.
///
/// Returns `Ok(true)` if any child was modified by `f`.
+ ///
+ /// # Error semantics
+ ///
+ /// If `f` returns `Err` for a child, this method returns that error
+ /// immediately. Children visited earlier in the same call keep whatever
+ /// modifications `f` already applied to them — they are **not** rolled
Review Comment:
👍
##########
datafusion/optimizer/src/optimizer.rs:
##########
@@ -360,6 +366,121 @@ impl TreeNodeRewriter for Rewriter<'_> {
}
}
+/// Applies `f` to each child (input) of `plan` in place, using
+/// [`Arc::make_mut`] for copy-on-write semantics on `Arc<LogicalPlan>`
+/// children. When the `Arc` refcount is 1 (the common case here)
+/// `Arc::make_mut` hands out a `&mut` without cloning; when it is >1 the
+/// inner value is cloned first.
+///
+/// Returns `Ok(true)` if any child was modified by `f`.
+///
+/// This is deliberately private to the optimizer rather than a method on
+/// [`LogicalPlan`]: it is an implementation detail of in-place rewriting, and
+/// the `Arc::make_mut` approach does not generalize to the other tree types
+/// (`Expr` children are `Box`ed; `PhysicalExpr`/`ExecutionPlan` children are
+/// `Arc<dyn _>`, which `Arc::make_mut` cannot handle). If `TreeNode` ever
+/// grows an in-place traversal this logic can move there.
+///
+/// # Error semantics
+///
+/// If `f` returns `Err` for a child, that error is returned immediately;
+/// children visited earlier keep whatever modifications `f` already applied
+/// to them — they are **not** rolled back.
+fn map_children_mut<F: FnMut(&mut LogicalPlan) -> Result<bool>>(
Review Comment:
I thought this was fine as a new public API but it is also ok here too
I wonder if you think we should make an issue to track improving / adding
other APIs in this area
##########
datafusion/optimizer/src/optimizer.rs:
##########
@@ -357,6 +360,95 @@ impl TreeNodeRewriter for Rewriter<'_> {
}
}
+/// 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.
+/// When the `Arc` refcount is 1 (always true in the optimizer),
+/// `Arc::make_mut` is essentially free.
+///
+/// The `rule.rewrite()` API takes ownership, so we 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 very next line.
+#[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> {
+ // f_down phase
+ let mut changed = false;
+ if apply_order == ApplyOrder::TopDown {
+ let owned = std::mem::take(plan);
+ let result = rule.rewrite(owned, config)?;
+ *plan = result.data;
+ changed |= result.transformed;
+ // Respect TreeNodeRecursion::Stop/Jump from the rule
+ if result.tnr == TreeNodeRecursion::Stop {
+ return Ok(changed);
+ }
+ }
+
+ // Recurse into children using Arc::make_mut (zero-cost when refcount == 1)
+ changed |= plan.map_children_mut(|child| {
+ rewrite_plan_in_place(child, apply_order, rule, config)
+ })?;
+
+ // f_up phase
+ if apply_order == ApplyOrder::BottomUp {
+ let owned = std::mem::take(plan);
+ let result = rule.rewrite(owned, config)?;
+ *plan = result.data;
+ changed |= result.transformed;
+ }
+
+ Ok(changed)
+}
+
+/// 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 {
Review Comment:
@adriangb do you think it is worth filing a ticket to explore this
potential option?
##########
datafusion/optimizer/src/optimizer.rs:
##########
@@ -365,13 +365,22 @@ impl TreeNodeRewriter for Rewriter<'_> {
///
/// This avoids the `Arc::unwrap_or_clone` + `Arc::new` cycle that the
/// ownership-based `TreeNode::rewrite` performs at every child node.
-/// When the `Arc` refcount is 1 (always true in the optimizer),
-/// `Arc::make_mut` is essentially free.
///
-/// The `rule.rewrite()` API takes ownership, so we 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 very next line.
+/// # Error semantics
+///
+/// On `Err`, `*plan` is left in an **unspecified** state and must not be used.
+/// Because `rule.rewrite()` consumes the plan by value, a failing rule drops
+/// the node it was handed and the [`std::mem::take`] placeholder
+/// (`LogicalPlan::default()`, an `EmptyRelation`) is left in its place — at
the
+/// root for a top-level failure, or somewhere in a subtree for a failure
deeper
+/// in the recursion. The pre-rule plan is **not** recoverable here: restoring
it
+/// would require cloning it before every rule invocation, which is exactly the
+/// allocation this function exists to avoid.
+///
+/// Callers must therefore discard `*plan` on `Err`, or restore it from a copy
+/// saved beforehand. [`Optimizer::optimize`] does this: with
`skip_failed_rules`
+/// it restores `new_plan` from the `prev_plan` clone it already holds, and
+/// without it the error aborts the pass and the plan is dropped unobserved.
Review Comment:
I think this is now overly wordy and hard to understand. The key point is
the first sentence and maybe one more sentence to add some additional context.
The rest I found to be hard to read
```suggestion
/// 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
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
