[EPIC] Stop copying `LogicalPlan` during OptimizerPasses

[alamb]

Is your feature request related to a problem or challenge?

Broken out from #9577 where @mustafasrepo @comphead and @jayzhan211 and I were discussing optimizer performance

TLDR is that the datafusion optimizer is slow. When I did some profiling locally by running the following

cargo bench --bench sql_planner -- physical_plan_tpch_all

My analysis is that almost 40% of the planning time is spent in SimplifyExprs and CommonSubexprEliminate and most of that time is related to copying expressions from what I can tell

While those passes themselves internally make a bunch of clones, which we are improving (e.g. @jayzhan211 on #9628) I think there is a more fundamental structural problem

I think a core challenge is that the OptimizerRule trait pretty much requires copying Exprs on each pass, as it gets a &LogicalPlan input, but produces a LogicalPlan output

    // Required methods
    fn try_optimize(
        &self,
        plan: &LogicalPlan,
        config: &dyn OptimizerConfig
    ) -> Result<Option<LogicalPlan>, DataFusionError>;

This mean any pass that works on Exprs must clone all Exprs (by calling LogicalPlan::expressions()) rewrite them, and then then create a new LogicalPlan with those new Exprs.

Here is that pattern in the expression simplifier:

https://github.com/apache/arrow-datafusion/blob/0eec5f8e1d0f55e48f5cdc628fbb5ddd89b91512/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs#L112-L123

Describe the solution you'd like

Find some way to avoid clone'ing exprs during LogicalPlan rewrite

Update: here are the tasks:

Infrastructure Preparation

• Refactor Optimizer to use owned plans and TreeNode API (10% faster planning) #9948
• Avoid LogicalPlan::clone() in LogicalPlan::map_children when possible #9999

Update OptimizerRules to avoid copying

• SimplifyExpressions: Introduce OptimizerRule::rewrite to rewrite in place, rewrite ExprSimplifier (20% faster planning) #9954
• CommonSubexprEliminate: Stop copying Exprs and LogicalPlans so much during Common Subexpression Elimination #9873
• DecorrelatePredicateSubquery: Stop copying LogicalPlan and Exprs in DecorrelatePredicateSubquery  #10289
• EliminateCrossJoin: Stop copying LogicalPlan and Exprs in EliminateCrossJoin #10287
• EliminateDuplicatedExpr: refactor EliminateDuplicatedExpr optimizer pass to avoid clone #10218
• EliminateFilter: Stop copying LogicalPlan and Exprs in EliminateFilter #10288
• EliminateJoin: Implement rewrite for EliminateOneUnion and EliminateJoin #10184
• EliminateLimit: Stop copying LogicalPlan and Exprs in EliminateLimit #10212
• EliminateNestedUnion: Stop copying LogicalPlan and Exprs in EliminateNestedUnion #10296
• EliminateOneUnion: Implement rewrite for EliminateOneUnion and EliminateJoin #10184
• EliminateOuterJoin: Refactor EliminateOuterJoin to implement OptimizerRule::rewrite() #10081
• ExtractEquijoinPredicate: implement rewrite for ExtractEquijoinPredicate and avoid clone in filter #10165
• FilterNullJoinKeys implement rewrite for FilterNullJoinKeys #10166
• OptimizeProjections: Stop copying LogicalPlan and Exprs in OptimizeProjections #10209
• PropagateEmptyRelation: Stop copying LogicalPlan and Exprs in PropagateEmptyRelation #10290
• PushDownFilter: Stop copying LogicalPlan and Exprs in PushDownFilter #10291
• PushDownLimit: Stop copying LogicalPlan and Exprs in PushDownLimit #10292
• ReplaceDistinctWithAggregate: Stop copying LogicalPlan and Exprs in ReplaceDistinctWithAggregate #10293
• RewriteDisjunctivePredicate: Stop copying LogicalPlan and Exprs in RewriteDisjunctivePredicate #10213
• ScalarSubqueryToJoin: Stop copying LogicalPlan and Exprs in ScalarSubqueryToJoin #10294
• SingleDistinctToGroupBy: Stop copying LogicalPlan and Exprs in SingleDistinctToGroupBy #10295
• UnwrapCastInComparison Refactor UnwrapCastInComparison to implement OptimizerRule::rewrite() #10087 / Refactor UnwrapCastInComparison to remove Expr clones #10115

Update AnalyzerRules to avoid copying

• AnalyzerMisc: Minor: Avoid copying all expressions in Analzyer / check_plan #9974
• InlineTableScan: Avoid copies in InlineTableScan via TreeNode API #10038
• ApplyFunctionRewrites: fix NamedStructField should be rewritten in OperatorToFunction in subquery regression (change ApplyFunctionRewrites to use TreeNode API #10032
• TypeCoercion: Stop copying LogicalPlan and Exprs in TypeCoercion #10210
• TypeCoercion more: Onyl recompute schema in TypeCoercion when necessary #10365
• CountWildcardRule: (needs a little reworking to avoid clones) Avoid copies in CountWildcardRule via TreeNode API #10066

Update Other to avoid copying

• Stop copying plans in LogicalPlan::with_param_values #10016

Describe alternatives you've considered

No response

Additional context

We have talked about various other ways to reduce copying of LogicalPlans as well as its challenges in other tickets:

• Remove Arc<LogicalPlan> from LogicalPlan, stop copying LogicalPlans #4628
• Optimizer is slow: Avoid too many string cloning in the optimizer #5157
• Replace Boxes with Arc in the Expr enum. #9577

[jayzhan211]

I have two idea before deep diving to the code

1. rewrite to have owned LogicalPlan then we can have owned expression easily. <- lots of code change
2. lazy clone on rewrite, the spirit similar to copy on write. The idea from simplify() Add a ScalarUDFImpl::simplfy() API, move SimplifyInfo et al to datafusion_expr #9304 may help here too or we can rely on Transformed<Expr>

enum SimpliedExpr {
  Simplified(Expr)
  Original
}

fn simply_expr(&self, expr: &Expr) -> SimpliedExpr

if expr is Original, we can avoid clone.

[alamb]

I have two idea before deep diving to the code

I think 2 sounds interesting

Another think I was thinking was something like LogicalPlan::rewrite(mut self, rewriter)

I think that along with Arc::try_unwrap could be used to minimize the places where cloning was actually needed

Maybe we can prototype with

impl OptimizerRule {
...

  /// does this rule support rewriting owned plans?
  fn supports_owned(&self) -> bool { return false }

  /// if supports_owned returns true, calls try_optimize_owned
  fn try_optimize_owned(
        &self,
        plan: LogicalPlan,
        config: &dyn OptimizerConfig
    ) -> Result<Transformed<LogicalPlan>, DataFusionError> {
      not_implemented_err!("try_optimized_owned is not implemented for this rule")
    }
...

And then play around with the code that calls try_optimize here https://github.com/apache/arrow-datafusion/blob/f4107d47bb4c0260d301294ddfc7c67d96574636/datafusion/optimizer/src/optimizer.rs#L360-L368 to try and use the try_optimize_owned API (without having to rewrite all the optimizers) for SimplifyExprs

If we could show a significant performance improvement for the sql_planner benchmarks then I think it would be worth spending time reworking the other APIs

[jayzhan211]

@alamb I think being able to skip failed rules is the main challenge to rewriting the plan with ownership

https://github.com/apache/arrow-datafusion/blob/37253e57beb25f0f1a4412b75421a489c2cb3c6a/datafusion/optimizer/src/optimizer.rs#L325

We need to restore the original plan if the rewrite fails at some point, however, the old plan is consumed and we lose ownership. Giving up ownership too early is not a good idea if we need to restore the original plan.

Instead of Transformed, I think we need to preserve the old data and fail state to mimic Err in Result but with old data

#[derive(Debug, PartialEq, Eq)]
pub enum OptimizedState {
    Yes,
    No,
    Fail,
}

#[derive(Debug)]
pub struct Optimized<T, E = DataFusionError> {
    pub optimzied_data: Option<T>,
    // Used to store the original data if optimized successfully
    pub original_data: T,
    pub optimized_state: OptimizedState,
    // Used to store the error if optimized failed, so we can early return but preserve the original data
    pub error: Option<E>,
}

impl<T, E> Optimized<T, E> {
    pub fn yes(optimzied_data: T, original_data: T) -> Self {
        Self {
            optimzied_data: Some(optimzied_data),
            original_data,
            optimized_state: OptimizedState::Yes,
            error: None,
        }
    }

    pub fn no(original_data: T) -> Self {
        Self {
            optimzied_data: None,
            original_data,
            optimized_state: OptimizedState::No,
            error: None,
        }
    }

    pub fn fail(original_data: T, e: E) -> Self {
        Self {
            optimzied_data: None,
            original_data,
            optimized_state: OptimizedState::Fail,
            error: Some(e),
        }
    }
}

Then, for every optimization, we return Optimized<LogicalPlan> or Optimzied<Expr>

[alamb]

@alamb I think being able to skip failed rules is the main challenge to rewriting the plan with ownership

This is an excellent point

Instead of Transformed, I think we need to preserve the old data and fail state to mimic Err in Result but with old data

If we need to preserve the old data, there is no choice but to copy the LogicalPlan on each pass 🤔 (well I guess there could be some way to repair the partially rewritten plan, but that sounds very compliated).

Since skip_failed_rules is false by default:

https://github.com/apache/arrow-datafusion/blob/dcfe70987e98da0410146b5e1292ab20f3f118e0/datafusion/common/src/config.rs#L549

Could we maybe only do the copy when skip_failed_rules is enabled? Something like this:

if skip_failed_rules {
  let original_plan = plan.clone();
  let new_plan = optimizer.try_optimize_owned(plan)
   .ok_or_else(|e| original_plan);
} else {
  optimizer.try_optimize_owned(plan)?
}

[jayzhan211]

I think it is possible to preserve unnecessary clones even if we preserve the old plan, and only clone the data for the new expr and new plan, but I agree it is a little complicated, it needs the enum I show above. We can optimize for the default path first, and others later on.

[alamb]

Let's give it a try!

[jayzhan211]

note: I profile my change and find out the time does not improve at all, I find that the cloned inside expr_list_to_fields may be the core issue.

[jayzhan211]

I analyze the sql_planner again and find that exprlist_to_fields and calc_func_dependencies_for_project are the two most time spending func.

so I think #9595 might be an important step for the optimization here.

https://github.com/apache/arrow-datafusion/blob/ad8d552b9f150c3c066b0764e84f72b667a649ff/datafusion/expr/src/logical_plan/plan.rs#L1804-L1813