Introduce OptimizerRule::try_optimize_owned to rewrite in place

Author: alamb

datafusion/optimizer/src/optimizer.rs

@@ -48,7 +48,7 @@ use crate::utils::log_plan;
 use datafusion_common::alias::AliasGenerator;
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::instant::Instant;
-use datafusion_common::{DFSchema, DataFusionError, Result};
+use datafusion_common::{internal_err, DFSchema, DataFusionError, Result};
 use datafusion_expr::logical_plan::LogicalPlan;
 
 use chrono::{DateTime, Utc};
@@ -80,12 +80,28 @@ pub trait OptimizerRule {
     /// A human readable name for this optimizer rule
     fn name(&self) -> &str;
 
-    /// How should the rule be applied by the optimizer? See comments on [`ApplyOrder`] for details.
+    /// How should the rule be applied by the optimizer? See comments on
+    /// [`ApplyOrder`] for details.
     ///
-    /// If a rule use default None, it should traverse recursively plan inside itself
+    /// If returns `None`, the default, the rule must handle recursion itself
     fn apply_order(&self) -> Option<ApplyOrder> {
         None
     }
+
+    /// Does this rule support rewriting owned plans (rather than by reference)?
+    fn supports_owned(&self) -> bool {
+        false
+    }
+
+    /// if supports_owned returns true, the Optimizer calls
+    /// [`Self::try_optimize_owned`] instead of [`Self::try_optimize`]
+    fn try_optimize_owned(
+        &self,
+        _plan: LogicalPlan,
+        _config: &dyn OptimizerConfig,
+    ) -> Result<Transformed<LogicalPlan>, DataFusionError> {
+        internal_err!("try_optimized_owned is not implemented for {}", self.name())
+    }
 }
 
 /// Options to control the DataFusion Optimizer.
@@ -297,7 +313,10 @@ fn optimize_plan_node(
     rule: &dyn OptimizerRule,
     config: &dyn OptimizerConfig,
 ) -> Result<Transformed<LogicalPlan>> {
-    // TODO: introduce a better API to OptimizerRule to allow rewriting by ownership
+    if rule.supports_owned() {
+        return rule.try_optimize_owned(plan, config);
+    }
+
     rule.try_optimize(&plan, config).map(|maybe_plan| {
         match maybe_plan {
             Some(new_plan) => {