Introduce a common trait TreeNode for ExecutionPlan, PhysicalExpr, LogicalExpr, LogicalPlan #5630

yahoNanJing · 2023-03-17T11:56:50Z

Which issue does this PR close?

Closes #5609.

Rationale for this change

Currently, there are many duplicated code related to the visitor trait and rewrite trait. And the trait methods are sometimes misused.

What changes are included in this PR?

Combine the visitor and rewrite trait to be one trait TreeNode and provide common behavior for ExecutionPlan, PhysicalExpr, LogicalExpr, LogicalPlan. And this trait contains the following methods:

fn get_children(&self) -> Vec<Self>; for getting its children
fn collect<F>(&self, op: &mut F) -> Result<()> where F: FnMut(&Self) -> Result<Recursion>, for collecting info from the tree node by a closure
fn visit<V: TreeNodeVisitor<N = Self>>(&self, visitor: &mut V) -> Result<Recursion> for collecting info from the tree node by a visitor. This is used when need to collect info from a specific post_visit. If only pre_visit is needed, then the collect method is preferred.
fn transform_down<F>(self, op: &F) -> Result<Self> where F: Fn(Self) -> Result<Option<Self>>, for rewriting the node from top to down
fn transform_up<F>(self, op: &F) -> Result<Self> where F: Fn(Self) -> Result<Option<Self>>, for rewriting the node from bottom to up
fn rewrite<R: TreeNodeRewriter<N = Self>>(self, rewriter: &mut R) -> Result<Self> for rewriting the node from top to down. This is used when need to invoke a specific pre_visit. Otherwise, either the transform_down or transform_up is preferred.
fn map_children<F>(self, transform: F) -> Result<Self> where F: FnMut(Self) -> Result<Self> is for replacing all of its children with the rewritted ones

Are these changes tested?

Are there any user-facing changes?

…table to TreeNode

yahoNanJing · 2023-03-17T11:58:33Z

Due to commit 258af4b, I may need to rebase to latest code and do a few refinements.

mingmwang · 2023-03-17T15:08:43Z

I will take a closer look tomorrow.

alamb · 2023-03-18T11:24:53Z

Combine the visitor and rewrite trait to be one trait TreeNode and provide common behavior for ExecutionPlan, PhysicalExpr, LogicalExpr, LogicalPlan. And this trait contains the following methods:

I believe Expr (logical exprs) are also included

alamb

Thank you for cleaning this up @yahoNanJing -- I really like how the same trait is used for visiting all the various trees that are in DataFusion.

It would be awesome to figure out how to avoid the duplication in the physical exprs and physical plans

This would be a fairly non trivial API change (I think for the better) for all users of DataFusion that have rewriters; This is probably fine but I think it would be good to communicate this more broadly with people who have systems built on DataFusion for them to weigh in as well.

When it is ready for review, we can ping them and do some more communocation.

alamb · 2023-03-18T11:27:59Z

datafusion/core/src/datasource/listing/helpers.rs

+/// was performed
+pub fn expr_applicable_for_cols(col_names: &[String], expr: &Expr) -> bool {
+    let mut is_applicable = true;
+    expr.collect(&mut |expr| {


alamb · 2023-03-18T11:31:31Z

datafusion/common/src/tree_node.rs

+/// [`TreeNodeVisitor`] allows keeping the algorithms
+/// separate from the code to traverse the structure of the `TreeNode`
+/// tree and makes it easier to add new types of tree node and
+/// algorithms by.


Suggested change

/// algorithms by.

/// algorithms.

alamb · 2023-03-18T11:32:05Z

datafusion/common/src/tree_node.rs

This API looks very nice 👌

alamb · 2023-03-18T11:32:30Z

datafusion/core/src/physical_optimizer/dist_enforcement.rs

-impl TreeNodeRewritable for PlanWithKeyRequirements {
+impl TreeNode for PlanWithKeyRequirements {
+    fn get_children(&self) -> Vec<Self> {
+        unimplemented!()


Should this be implemented?

I also have the same question.

alamb · 2023-03-18T11:37:02Z

datafusion/physical-expr/src/tree_node/mod.rs

+//! This module provides common traits for visiting or rewriting tree nodes easily.
+//!
+//! It's a duplication of the one in the crate `datafusion-common`.
+//! In the future, if the Orphan rule is relaxed for Arc<dyn T>, these duplicated codes can be removed.


It would be great if we could figure out how to avoid this duplication and the one in physical-expr?

Maybe we could create a blanket impl in datafusion-common like

impl<T: TreeNode> TreeNode for Arc<T> { ... }

?

Thanks @alamb. However, I'm still not sure how to apply to the trait ExecutionPlan and PhysicalExpr

I played around with this and I came up with a solution -- here is a proposed modification to this PR: yahoNanJing#70

metesynnada · 2023-03-19T12:43:09Z

datafusion/common/src/tree_node.rs

+    where
+        F: Fn(Self) -> Result<Option<Self>>,
+    {
+        let node_cloned = self.clone();


Can we remove cloning?

Thanks for your comments. I'm afraid it's not feasible because of the constraint of F: Fn(Self) -> Result<Option<Self>>,

F will consume the Self.

I think this clone could be quite bad for performance. For Arcd PhysicalExprs and ExecutionPlans, it is very cheap, but for Expr it means any code that uses this will effectively deep copy (recursively) the entire tree at each step (N^2 performance).

I think the root cause is the signature

F: Fn(Self) -> Result<Option<Self>>,

Which I think is being used to signal if the node was changed or not

In order to avoid the clone, I think we need to use something other than Option so the original can be returned when the change is not made.

How about something like

enum Transformed<T> { /// The item was transformed / rewritten somehow Yes(T), /// The item was not transformed No(T) } impl<T> From<Transformed<T>> for <T> { fn from(transformed: Transformed<T>) -> T { match transformed { Yes(t) => t, No(t) => t, } }

And then this function would look something like

fn transform_down<F>(self, op: &F) -> Result<Self> where F: Fn(Self) -> Result<Transformed<Self>>, { // this could be cleaner using the `From` impl let after_op = match op(node_cloned)? { Yes(value) => value, No(value) => value, }; after_op.map_children(|node| node.transform_down(op)) }

I agree with @alamb that this is due to the chosen signature, not an unavoidable clone -- and we should fix it. Same problem in exists in transform_up function too. When no change is made, the cloning is done wastefully.

I am particular concerned that this trait will be applied to Exprs (where as LogicalPlans and PhysicalPlans already have Arc's in them so the cost is not that high)

I would be willing to try and change the existing LogicalPlan optimizer (which I think might be causing this problems)

@jackwener and I were just discussing a similar issue with #5705 (comment)

mingmwang · 2023-03-20T02:43:45Z

datafusion/common/src/tree_node.rs

+    /// Use preorder to iterate the node on the tree so that we can stop fast for some cases.
+    ///
+    /// `op` can be used to collect some info from the tree node.
+    fn collect<F>(&self, op: &mut F) -> Result<()>


collect usually refers to apply some func to each element and the returned result should be a container type.

Thanks for your suggestion. Maybe it's better to change the name to for_each.

May be try_for_each ?

Personally, I prefer to use apply. And later I'll remove the get_children and add apply_children.

mingmwang · 2023-03-20T02:51:25Z

datafusion/common/src/tree_node.rs

+                return Err(DataFusionError::Execution(format!(
+                    "Recursion {r:?} is not supported for collect_using"
+                )))


Please change the error msg:

Recursion {r:?} is not supported for visit()

mingmwang · 2023-03-20T03:00:23Z

datafusion/common/src/tree_node.rs

+            }
+        };
+
+        for child in self.get_children() {


Can you use lambda instead of the for loop ?

mingmwang · 2023-03-20T03:13:37Z

datafusion/common/src/tree_node.rs

+        match op(self)? {
+            Recursion::Continue => {}
+            // If the recursion should stop, do not visit children
+            Recursion::Stop => return Ok(()),
+            r => {
+                return Err(DataFusionError::Execution(format!(
+                    "Recursion {r:?} is not supported for collect"
+                )))
+            }
+        };


Maybe you can consider using the ControlFlow from rust std ?

pub enum ControlFlow<B, C = ()> { /// Move on to the next phase of the operation as normal. #[stable(feature = "control_flow_enum_type", since = "1.55.0")] #[lang = "Continue"] Continue(C), /// Exit the operation without running subsequent phases. #[stable(feature = "control_flow_enum_type", since = "1.55.0")] #[lang = "Break"] Break(B), // Yes, the order of the variants doesn't match the type parameters. // They're in this order so that `ControlFlow<A, B>` <-> `Result<B, A>` // is a no-op conversion in the `Try` implementation. }

Good suggestion. I'll refine for it.

mingmwang · 2023-03-20T03:18:46Z

datafusion/core/src/physical_plan/tree_node/mod.rs

+}
+
+/// Trait for potentially recursively transform an [`TreeNode`] node
+/// tree. When passed to `TreeNode::transform_using`, `TreeNodeRewriter::mutate` is
 /// invoked recursively on all nodes of a tree.


There is no transform_using method any more. Please fix the comment.

mingmwang · 2023-03-20T03:30:28Z

datafusion/core/src/physical_optimizer/pipeline_checker.rs

+    fn get_children(&self) -> Vec<Self> {
+        unimplemented!()
+    }


I think this is some kind of anti pattern.

mingmwang · 2023-03-20T03:44:31Z

datafusion/expr/src/tree_node/expr.rs

+use crate::Expr;
+use datafusion_common::{tree_node::TreeNode, Result};
+
+impl TreeNode for Expr {


There is some kind of duplication logic, I think maybe you can provide a default implementation for map_children in the TreeNode trait.

mingmwang · 2023-03-20T03:49:55Z

datafusion/expr/src/tree_node/expr.rs

+
+/// &mut transform a `Vec` of `Expr`s
+fn transform_vec<F>(v: Vec<Expr>, transform: &mut F) -> Result<Vec<Expr>>
+where


You can use slice or IntoIterator.

yahoNanJing · 2023-03-22T07:08:14Z

Hi @alamb, this PR is ready for review now. Could you help ping related guys for helping reviewing this PR?

alamb

Thank you very much @yahoNanJing for pushing this forward. I really like this PR except for:

The cloning by default in transform_down (I have left some comments on how maybe to fix that)
The replication of the trait -- I have a proposal of how to avoid the duplication in Remove duplicated TreeNode traits in physical-expr and physical plan yahoNanJing/arrow-datafusion#70

alamb · 2023-03-23T13:06:30Z

datafusion/common/src/tree_node.rs

+    where
+        F: Fn(Self) -> Result<Option<Self>>,
+    {
+        let node_cloned = self.clone();


I think this clone could be quite bad for performance. For Arcd PhysicalExprs and ExecutionPlans, it is very cheap, but for Expr it means any code that uses this will effectively deep copy (recursively) the entire tree at each step (N^2 performance).

I think the root cause is the signature

F: Fn(Self) -> Result<Option<Self>>,

Which I think is being used to signal if the node was changed or not

In order to avoid the clone, I think we need to use something other than Option so the original can be returned when the change is not made.

How about something like

enum Transformed<T> { /// The item was transformed / rewritten somehow Yes(T), /// The item was not transformed No(T) } impl<T> From<Transformed<T>> for <T> { fn from(transformed: Transformed<T>) -> T { match transformed { Yes(t) => t, No(t) => t, } }

And then this function would look something like

fn transform_down<F>(self, op: &F) -> Result<Self> where F: Fn(Self) -> Result<Transformed<Self>>, { // this could be cleaner using the `From` impl let after_op = match op(node_cloned)? { Yes(value) => value, No(value) => value, }; after_op.map_children(|node| node.transform_down(op)) }

alamb · 2023-03-23T13:55:31Z

datafusion/physical-expr/src/tree_node/mod.rs

+//! This module provides common traits for visiting or rewriting tree nodes easily.
+//!
+//! It's a duplication of the one in the crate `datafusion-common`.
+//! In the future, if the Orphan rule is relaxed for Arc<dyn T>, these duplicated codes can be removed.


I played around with this and I came up with a solution -- here is a proposed modification to this PR: yahoNanJing#70

Remove duplicated TreeNode traits in physical-expr and physical plan

yahoNanJing · 2023-03-24T12:32:30Z

Thanks @alamb. According to your advice, the clone() and duplicated [TreeNode] definition have been removed. If there's no big issues, I suggest to merge this PR as soon as possible so that others can leverage the common trait as soon as possible. And another benefit would be less effort to deal with the conflicts, since there're so many files involved in this PR.

alamb

I think this looks great @yahoNanJing -- thank you so much.

I recommend we merge this PR as soon as we are good with the 21.0.0 release. There is a minor issue reported on the mailing list here: https://lists.apache.org/thread/fnnmrn83spnb2y2l2vdw2v1hd54pfjl7

alamb · 2023-03-25T12:48:15Z

Unless anyone else has thoughts on this PR I plan to merge it monday

crepererum · 2023-03-27T13:32:41Z

datafusion/common/src/tree_node.rs

+    }
+}
+
+/// Helper trait for implementing [`TreeNode`] that have children stored as Arc's


This is a very clever design. I tried this refactor before and somewhat gave up because of Rust's orphan rule (i.e. you cannot impl Trait for Arc<T> if trait is from a different crate).

mingmwang · 2023-03-28T03:52:37Z

Looks really nice !

mingmwang · 2023-03-28T03:59:42Z

@yahoNanJing
There are some other improvement we can do:

In the map_children() method of LogicalPlan, we can avoid construct the plan again and again if the new children has no changes compared to the old inputs.

    fn map_children<F>(self, transform: F) -> Result<Self>
    where
        F: FnMut(Self) -> Result<Self>,
    {
        let children = self.inputs().into_iter().cloned().collect::<Vec<_>>();
        if !children.is_empty() {
            let new_children: Result<Vec<_>> =
                children.into_iter().map(transform).collect();
            self.with_new_inputs(new_children?.as_slice())
        } else {
            Ok(self)
        }
    }

We have the with_new_children_if_necessary utils for Arc<dyn ExecutionPlan> and Arc<dyn PhysicalExpr>

mingmwang · 2023-03-28T04:01:55Z

I think this method can be removed.

pub fn optimize_children(
    optimizer: &impl PhysicalOptimizerRule,
    plan: Arc<dyn ExecutionPlan>,
    config: &ConfigOptions,
) -> Result<Arc<dyn ExecutionPlan>> {
    let children = plan
        .children()
        .iter()
        .map(|child| optimizer.optimize(Arc::clone(child), config))
        .collect::<Result<Vec<_>>>()?;

    if children.is_empty() {
        Ok(Arc::clone(&plan))
    } else {
        with_new_children_if_necessary(plan, children).map(Transformed::into)
    }
}

alamb · 2023-03-28T11:38:34Z

In the map_children() method of LogicalPlan, we can avoid construct the plan again and again if the new children has no changes compared to the old inputs.

in #5761

alamb · 2023-03-28T12:29:02Z

I think this method can be removed.

Filed #5762 to track

Thanks for the comments @mingmwang

alamb · 2023-03-28T15:16:47Z

FWIW I upgraded IOx (which uses the rewriters heavily) to this rewriter code in https://github.com/influxdata/influxdb_iox/pull/7353 -- it was some non trivial API churn but I think resulting code is much nicer 👌

Thanks again @yahoNanJing

kyotoYaho added 7 commits March 16, 2023 11:59

Repalce the TreeNodeVisitor with Closure and change the TreeNodeRewri…

7252d26

…table to TreeNode

Reuse TreeNode for physical expression

3041506

Implement TreeNode for logical Expr

90e31ee

Implement TreeNode for logical plan

5a44a48

Remove ExprRewriter

dad4390

Rename transform_using to rewrite and collect_using to visit in TreeNode

8111d74

Remove PlanVisitor

33138cb

yahoNanJing marked this pull request as draft March 17, 2023 11:57

github-actions bot added core Core DataFusion crate logical-expr Logical plan and expressions optimizer Optimizer rules physical-expr Physical Expressions labels Mar 17, 2023

yahoNanJing changed the title ~~Issue 5609~~ Introduce a common trait TreeNode for ExecutionPlan, PhysicalExpr, LogicalExpr, LogicalPlan Mar 17, 2023

yahoNanJing requested a review from alamb March 17, 2023 11:59

kyotoYaho added 2 commits March 18, 2023 13:37

Merge branch 'main' into issue-5609

7babedd

Fix merge main branch

7e51a74

alamb reviewed Mar 18, 2023

View reviewed changes

metesynnada reviewed Mar 19, 2023

View reviewed changes

Remove the rewrite.rs introduced by 258af4b

8a21d7d

mingmwang reviewed Mar 20, 2023

View reviewed changes

Merge branch 'main' into issue-5609

15f0a78

kyotoYaho and others added 3 commits March 22, 2023 15:13

Minor fix

7937a9c

Remove duplicated TreeNode definition in physical-expr

992104f

Remove duplication in physical_plan

017fa78

alamb mentioned this pull request Mar 23, 2023

Remove duplicated TreeNode traits in physical-expr and physical plan yahoNanJing/arrow-datafusion#70

Merged

alamb reviewed Mar 23, 2023

View reviewed changes

jackwener self-requested a review March 24, 2023 02:52

yahoNanJing and others added 2 commits March 24, 2023 13:18

Merge pull request #70 from alamb/alamb/reduce_duplication

0f1aa4b

Remove duplicated TreeNode traits in physical-expr and physical plan

Introduce enum Transformed to avoid clone in the TreeNode

92cc2f1

jackwener removed their request for review March 24, 2023 11:31

kyotoYaho added 2 commits March 24, 2023 19:41

Rename the trait ArcWithChildren to DynTreeNode

128d6ed

Merge remote-tracking branch 'origin/main' into issue-5609

5a1f73b

github-actions bot added the sql SQL Planner label Mar 24, 2023

alamb approved these changes Mar 24, 2023

View reviewed changes

alamb added the api change Changes the API exposed to users of the crate label Mar 25, 2023

alamb merged commit 8df18ab into apache:main Mar 27, 2023

crepererum reviewed Mar 27, 2023

View reviewed changes

alamb mentioned this pull request Mar 28, 2023

Minor: Avoid an unecessary contruction in map_children some extra plan construction #5761

Merged

alamb mentioned this pull request Mar 28, 2023

Remove optimize_children and replace with map_children #5762

Closed

peter-toth mentioned this pull request Oct 26, 2023

Refactor TreeNode recursions #7942

Closed

andygrove mentioned this pull request Jan 14, 2024

Remove Transformed enum #8835

Closed

Introduce a common trait TreeNode for ExecutionPlan, PhysicalExpr, LogicalExpr, LogicalPlan #5630

Introduce a common trait TreeNode for ExecutionPlan, PhysicalExpr, LogicalExpr, LogicalPlan #5630

Conversation

yahoNanJing commented Mar 17, 2023 • 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?

yahoNanJing commented Mar 17, 2023

mingmwang commented Mar 17, 2023

alamb commented Mar 18, 2023 • edited Loading

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mingmwang Mar 20, 2023 • edited Loading

Choose a reason for hiding this comment

alamb Mar 23, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yahoNanJing Mar 20, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mingmwang Mar 20, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yahoNanJing commented Mar 22, 2023

alamb left a comment

Choose a reason for hiding this comment

alamb Mar 23, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yahoNanJing commented Mar 24, 2023

alamb left a comment

Choose a reason for hiding this comment

alamb commented Mar 25, 2023

Choose a reason for hiding this comment

mingmwang commented Mar 28, 2023

mingmwang commented Mar 28, 2023

mingmwang commented Mar 28, 2023

alamb commented Mar 28, 2023

alamb commented Mar 28, 2023

alamb commented Mar 28, 2023

yahoNanJing commented Mar 17, 2023 •

edited

Loading

alamb commented Mar 18, 2023 •

edited

Loading

mingmwang Mar 20, 2023 •

edited

Loading

alamb Mar 23, 2023 •

edited

Loading

yahoNanJing Mar 20, 2023 •

edited

Loading

mingmwang Mar 20, 2023 •

edited

Loading

alamb Mar 23, 2023 •

edited

Loading