Improve `TreeNode` and `LogicalPlan` APIs to accept owned closures, deprecate `transform_down_mut()` and `transform_up_mut()` #10126

peter-toth · 2024-04-18T09:57:38Z

Which issue does this PR close?

Closes #10097, part of: #10121.

Rationale for this change

Currently the only difference between TreeNode::transform_down() and TreeNode::transform_down_mut() is that transform_down() accepts an &Fn, while transform_down_mut() accepts an &mut FnMut closure.
As Fn is subtrait of FnMut, it is possible to use transform_down_mut() instead of transform_down() everywhere. But because transform_down() has a better name, let's fix transform_down() and get rid of transform_down_mut().
The same applies to TreeNode::transform_up() vs. TreeNode::transform_up_mut() as well.
Currently TreeNode::apply(), TreeNode::transform(), TreeNode::transform_down(), TreeNode::transform_up(), TreeNode::transform_down_up() and their LogicalPlan::...with_subqueries() variants use &Fn or &mut FnMut closures, which is a bit cumbersome for the API users. This PR proposes to change these APIs to accept owned FnMut closures.

Please note that this is an API breaking change but the fix is straightforward.

What changes are included in this PR?

This PR:

Modifies TreeNode::apply(), TreeNode::transform(), TreeNode::transform_down(), TreeNode::transform_up(), TreeNode::transform_down_up() and their LogicalPlan::...with_subqueries() variants to use owned FnMut closures.
Changes all transform_down_mut() calls to transform_down() and transform_up_mut() calls to transform_up().
Deprecates transform_down_mut() and transform_up_mut().
Removes LogicalPlan::transform_down_mut_with_subqueries() and LogicalPlan::transform_up_mut_with_subqueries() to keep LogicalPlan::..._with_subqueries() APIs in sync with TreeNode APIs. (Please note that transform_down_mut_with_subqueries() and transform_up_mut_with_subqueries() haven't been released yet so we can freely remove them.)

Are these changes tested?

Yes, with existing UTs.

Are there any user-facing changes?

No.

…p_mut()` methods

peter-toth · 2024-04-18T12:45:04Z

@alamb this PR might conflict with #10035, but I can update mine if yours gets merged first.

alamb

Thank you @peter-toth -- this is looking really neat. It would be really cool if we could figure out how to avoid neeidng &mut in front of all the closures (and that might actually make the change backwards compatible too).

However, I realize this may not be possible / easy to do. If that is the case, I think this approach would be an improvement over the current state on master.

What do you think?

alamb · 2024-04-18T22:03:04Z

datafusion/common/src/tree_node.rs

        self,
-        f: &F,
+        f: &mut F,


I tried a bit locally to change this to be f: F (aka take an "owned" F) -- which would make it easier to call this API also works -- so the callsites don't need to be changed to be &mut all over the place

However, I tried it and I got some strange errors when instantiating macros 🤔

error: reached the recursion limit while instantiating `<expr::Expr as TreeNode>::transform_up::<&mut &mut &mut &mut ...>` --> /Users/andrewlamb/Software/arrow-datafusion/datafusion/common/src/tree_node.rs:184:50 | 184 | handle_transform_recursion_up!(self, |c| c.transform_up(&mut f), f) | ^^^^^^^^^^^^^^^^^^^^^^ | note: `transform_up` defined here --> /Users/andrewlamb/Software/arrow-datafusion/datafusion/common/src/tree_node.rs:180:5 | 180 | / fn transform_up<F: FnMut(Self) -> Result<Transformed<Self>>>( 181 | | self, 182 | | mut f: F, 183 | | ) -> Result<Transformed<Self>> { | |__________________________________^ = note: the full type name has been written to '/Users/andrewlamb/Software/arrow-datafusion/target/debug/deps/datafusion_expr-598aae957f8f2540.long-type.txt'

Yes, I did try this already and run into the same issue...
But now I think it is doable with some smart refactor. Let me update this PR today or tomorrow and then we can decide if this change make sense or not

d781ca7 changes TreeNode::apply(), TreeNode::transform(), TreeNode::transform_down(), TreeNode::transform_up(), TreeNode::transform_down_up() and their LogicalPlan::...with_subqueries() variants to use owned closures.

If we like this change I can update the PR description.

alamb · 2024-04-18T22:08:49Z

datafusion/core/src/physical_optimizer/convert_first_last.rs

@@ -60,7 +60,7 @@ impl PhysicalOptimizerRule for OptimizeAggregateOrder {
        plan: Arc<dyn ExecutionPlan>,
        _config: &ConfigOptions,
    ) -> Result<Arc<dyn ExecutionPlan>> {
-        plan.transform_up(&get_common_requirement_of_aggregate_input)
+        plan.transform_up(&mut get_common_requirement_of_aggregate_input)


This is just so weird -- I think most users would expect that this looks like this (though as I mentioned above, I couldn't make it work in a few minutes)

- plan.transform_up(&mut get_common_requirement_of_aggregate_input) + plan.transform_up(get_common_requirement_of_aggregate_input)

alamb · 2024-04-18T22:09:55Z

datafusion/expr/src/logical_plan/tree_node.rs

    }

    /// Similarly to [`Self::transform_up`], rewrites this node and its inputs using `f`,
    /// including subqueries that may appear in expressions such as `IN (SELECT
    /// ...)`.
-    pub fn transform_up_with_subqueries<F: Fn(Self) -> Result<Transformed<Self>>>(


berkaysynnada · 2024-04-19T07:00:57Z

Thanks, @peter-toth—this really simplifies things. I'm wondering if there's an option to avoid changing the transform's Fn to FnMut (thus not requiring &mut for closures) by using the rewrite API for passes that need mutability inevitably. AFAIK rewrite can handle that for the transform_up and transform_down cases.

… transform_down and transform_down_up APIs to accept owned closures

peter-toth · 2024-04-19T10:12:06Z

Thanks, @peter-toth—this really simplifies things. I'm wondering if there's an option to avoid changing the transform's Fn to FnMut (thus not requiring &mut for closures) by using the rewrite API for passes that need mutability inevitably. AFAIK rewrite can handle that for the transform_up and transform_down cases.

My concern is that it would require to refactor all current transform_down_mut() and transform_up_mut() usecases to use the rewrite() API, which is a bigger change required from the API users than just change to transform_down() (or transform_up()) and & to &mut.

alamb

Thank you @peter-toth -- I think this looks like a major step forward in TreeNode API usability (both due to removing the transform_***_mut variants, as well as accepting closures directly -- very nice 👌 🏆

Thanks, @peter-toth—this really simplifies things. I'm wondering if there's an option to avoid changing the transform's Fn to FnMut (thus not requiring &mut for closures) by using the rewrite API for passes that need mutability inevitably. AFAIK rewrite can handle that for the transform_up and transform_down cases.

@berkaysynnada I think the most recent change to this PR addresses your concern (closures do not need a &mut anymore). Can you take another look?

@alamb this PR might conflict with #10035, but I can update mine if yours gets merged first.

Thanks for the heads up. Since #10126 isn't approved / merged yet I think we can simply merge this one first (and then I can clean up that one with the changes)

cc @liukun4515 @mustafasrepo @ozankabak

alamb · 2024-04-19T12:19:42Z

datafusion-examples/examples/function_factory.rs

@@ -164,7 +164,7 @@ impl ScalarUDFImpl for ScalarFunctionWrapper {
 impl ScalarFunctionWrapper {
    // replaces placeholders such as $1 with actual arguments (args[0]
    fn replacement(expr: &Expr, args: &[Expr]) -> Result<Expr> {
-        let result = expr.clone().transform(&|e| {
+        let result = expr.clone().transform(|e| {


I actually think this is a really nice improvement in usability -- to not have to put a & in front of the closure is 💯

peter-toth · 2024-04-19T12:34:25Z

I've updated the PR title and description to reflect the latest changes.

alamb · 2024-04-21T11:12:43Z

I plan to merge this PR on Monday unless anyone else would like additional time to review or has additional comments

alamb · 2024-04-22T19:19:43Z

Thanks again @peter-toth -- epic work

…eprecate `transform_down_mut()` and `transform_up_mut()` (apache#10126) * Deprecate `TreeNode::transform_down_mut()` and `TreeNode::transform_up_mut()` methods * Refactor `TreeNode` and `LogicalPlan` apply, transform, transform_up, transform_down and transform_down_up APIs to accept owned closures

Deprecate TreeNode::transform_down_mut() and `TreeNode::transform_u…

b2f6309

…p_mut()` methods

github-actions bot added sql SQL Planner logical-expr Logical plan and expressions physical-expr Physical Expressions optimizer Optimizer rules core Core DataFusion crate labels Apr 18, 2024

alamb reviewed Apr 18, 2024

View reviewed changes

Refactor TreeNode and LogicalPlan apply, transform, transform_up,…

d781ca7

… transform_down and transform_down_up APIs to accept owned closures

alamb approved these changes Apr 19, 2024

View reviewed changes

peter-toth changed the title ~~Deprecate TreeNode::transform_down_mut() and TreeNode::transform_up_mut() methods~~ Improve TreeNode and LogicalPlan APIs to accept owned closures Apr 19, 2024

alamb changed the title ~~Improve TreeNode and LogicalPlan APIs to accept owned closures~~ Improve TreeNode and LogicalPlan APIs to accept owned closures, deprecate transform_down_mut() and transform_up_mut() Apr 21, 2024

alamb merged commit 3c3cb87 into apache:main Apr 22, 2024
28 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Improve `TreeNode` and `LogicalPlan` APIs to accept owned closures, deprecate `transform_down_mut()` and `transform_up_mut()` #10126

Improve `TreeNode` and `LogicalPlan` APIs to accept owned closures, deprecate `transform_down_mut()` and `transform_up_mut()` #10126

peter-toth commented Apr 18, 2024 •

edited

Loading

peter-toth commented Apr 18, 2024

alamb left a comment

alamb Apr 18, 2024

peter-toth Apr 19, 2024

peter-toth Apr 19, 2024

peter-toth Apr 19, 2024

alamb Apr 18, 2024

alamb Apr 18, 2024

berkaysynnada commented Apr 19, 2024

peter-toth commented Apr 19, 2024

alamb left a comment

alamb Apr 19, 2024

peter-toth commented Apr 19, 2024

alamb commented Apr 21, 2024

alamb commented Apr 22, 2024

Improve TreeNode and LogicalPlan APIs to accept owned closures, deprecate transform_down_mut() and transform_up_mut() #10126

Improve TreeNode and LogicalPlan APIs to accept owned closures, deprecate transform_down_mut() and transform_up_mut() #10126

Conversation

peter-toth commented Apr 18, 2024 • edited Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

peter-toth commented Apr 18, 2024

alamb left a comment

Choose a reason for hiding this comment

alamb Apr 18, 2024

Choose a reason for hiding this comment

peter-toth Apr 19, 2024

Choose a reason for hiding this comment

peter-toth Apr 19, 2024

Choose a reason for hiding this comment

peter-toth Apr 19, 2024

Choose a reason for hiding this comment

alamb Apr 18, 2024

Choose a reason for hiding this comment

alamb Apr 18, 2024

Choose a reason for hiding this comment

berkaysynnada commented Apr 19, 2024

peter-toth commented Apr 19, 2024

alamb left a comment

Choose a reason for hiding this comment

alamb Apr 19, 2024

Choose a reason for hiding this comment

peter-toth commented Apr 19, 2024

alamb commented Apr 21, 2024

alamb commented Apr 22, 2024

Improve `TreeNode` and `LogicalPlan` APIs to accept owned closures, deprecate `transform_down_mut()` and `transform_up_mut()` #10126

Improve `TreeNode` and `LogicalPlan` APIs to accept owned closures, deprecate `transform_down_mut()` and `transform_up_mut()` #10126

peter-toth commented Apr 18, 2024 •

edited

Loading