feat: make execution transfomation application rule generic

Currently, it depends on the particular execution transformation
that resolves placeholders. However, it would be more convenient
to be able to specify the rule which is required to be applied to
the plan.

Author: askalt

datafusion/core/tests/physical_optimizer/physical_expr_resolver.rs

@@ -29,7 +29,7 @@ use datafusion_physical_expr::{
     expressions::{BinaryExpr, col, lit, placeholder},
 };
 use datafusion_physical_optimizer::{
-    PhysicalOptimizerRule, physical_expr_resolver::PhysicalExprResolver,
+    PhysicalOptimizerRule, exec_transform_apply::PhysicalExprResolver,
 };
 use datafusion_physical_plan::{
     ExecutionPlan, filter::FilterExec, get_plan_string,

…-optimizer/src/physical_expr_resolver.rs …al-optimizer/src/exec_transform_apply.rsdatafusion/physical-optimizer/src/physical_expr_resolver.rs renamed to datafusion/physical-optimizer/src/exec_transform_apply.rs datafusion/physical-optimizer/src/physical_expr_resolver.rs renamed to datafusion/physical-optimizer/src/exec_transform_apply.rs

@@ -15,11 +15,10 @@
 // specific language governing permissions and limitations
 // under the License.
 
-//! [`PhysicalExprResolver`] ensures that the physical plan is prepared for placeholder resolution
-//! by wrapping it in a [`TransformPlanExec`] with a [`ResolvePlaceholdersRule`] if the plan
-//! contains any unresolved placeholders. The actual resolution happens during execution.
+//! [`ExecutionTransformationApplier`] ensures that the required execution transformations
+//! are applied to the physical plan.
 
-use std::sync::Arc;
+use std::{borrow::Cow, sync::Arc};
 
 use datafusion_common::{
     Result,
@@ -28,61 +27,67 @@ use datafusion_common::{
 };
 use datafusion_physical_plan::{
     ExecutionPlan,
-    plan_transformer::{ResolvePlaceholdersRule, TransformPlanExec},
+    plan_transformer::{ExecutionTransformationRule, TransformPlanExec},
 };
 
 use crate::PhysicalOptimizerRule;
 
-/// The phase in which the [`PhysicalExprResolver`] rule is applied.
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub enum PhysicalExprResolverPhase {
+/// The phase in which the [`ExecutionTransformationApplier`] rule is applied.
+#[derive(Debug)]
+pub enum ExecutionTransformationApplierPhase {
     /// Optimization that happens before most other optimizations.
     /// This optimization removes all [`TransformPlanExec`] execution plans from the plan
     /// tree.
     Pre,
     /// Optimization that happens after most other optimizations.
-    /// This optimization checks for the presence of placeholders in the optimized plan, and if
-    /// they are present, wraps the plan in a [`TransformPlanExec`] with a [`ResolvePlaceholdersRule`].
-    Post,
+    /// This optimization checks if `rule` requires to transform the plan and wraps the plan with
+    /// [`TransformPlanExec`] if it so, or adds rule to the existing transformation node.
+    Post {
+        rule: Arc<dyn ExecutionTransformationRule>,
+    },
 }
 
-/// Physical optimizer rule that prepares the plan for placeholder resolution during execution.
+/// Physical optimizer rule that wraps the plan with a certain execution-stage transformation.
 #[derive(Debug)]
-pub struct PhysicalExprResolver {
-    phase: PhysicalExprResolverPhase,
+pub struct ExecutionTransformationApplier {
+    phase: ExecutionTransformationApplierPhase,
+    name: Cow<'static, str>,
 }
 
-impl PhysicalExprResolver {
-    /// Creates a new [`PhysicalExprResolver`] optimizer rule that runs in the pre-optimization
-    /// phase. In this phase, the rule removes any existing [`TransformPlanExec`] from the
-    /// plan tree.
+impl ExecutionTransformationApplier {
+    /// Creates a new [`ExecutionTransformationApplier`] optimizer rule that runs in the
+    /// pre-optimization phase.
     pub fn new() -> Self {
         Self {
-            phase: PhysicalExprResolverPhase::Pre,
+            phase: ExecutionTransformationApplierPhase::Pre,
+            name: Cow::Borrowed("ExecutionTransformationApplier"),
         }
     }
 
-    /// Creates a new [`PhysicalExprResolver`] optimizer rule that runs in the post-optimization
-    /// phase. In this phase, the rule wraps the physical plan in a [`TransformPlanExec`] with a
-    /// [`ResolvePlaceholdersRule`] if the plan contains any unresolved placeholders.
-    pub fn new_post_optimization() -> Self {
+    /// Creates a new [`ExecutionTransformationApplier`] optimizer rule that runs in the
+    /// post-optimization phase.
+    pub fn new_post_optimization(rule: Arc<dyn ExecutionTransformationRule>) -> Self {
+        let name = format!("ExecutionTransformationApplier({})", rule.name());
         Self {
-            phase: PhysicalExprResolverPhase::Post,
+            phase: ExecutionTransformationApplierPhase::Post { rule },
+            name: name.into(),
         }
     }
 }
 
-impl Default for PhysicalExprResolver {
+impl Default for ExecutionTransformationApplier {
     fn default() -> Self {
         Self::new()
     }
 }
 
-impl PhysicalOptimizerRule for PhysicalExprResolver {
+impl PhysicalOptimizerRule for ExecutionTransformationApplier {
     fn name(&self) -> &str {
         match self.phase {
-            PhysicalExprResolverPhase::Pre => "PhysicalExprResolver",
-            PhysicalExprResolverPhase::Post => "PhysicalExprResolver(Post)",
+            ExecutionTransformationApplierPhase::Pre => "PhysicalExprResolver",
+            ExecutionTransformationApplierPhase::Post { .. } => {
+                "PhysicalExprResolver(Post)"
+            }
         }
     }
 
@@ -91,8 +96,8 @@ impl PhysicalOptimizerRule for PhysicalExprResolver {
         plan: Arc<dyn ExecutionPlan>,
         _config: &ConfigOptions,
     ) -> Result<Arc<dyn ExecutionPlan>> {
-        match self.phase {
-            PhysicalExprResolverPhase::Pre => plan
+        match &self.phase {
+            ExecutionTransformationApplierPhase::Pre => plan
                 .transform_up(|plan| {
                     if let Some(plan) = plan.as_any().downcast_ref::<TransformPlanExec>()
                     {
@@ -102,26 +107,22 @@ impl PhysicalOptimizerRule for PhysicalExprResolver {
                     }
                 })
                 .map(|t| t.data),
-            PhysicalExprResolverPhase::Post => {
+            ExecutionTransformationApplierPhase::Post { rule } => {
                 if let Some(transformer) =
                     plan.as_any().downcast_ref::<TransformPlanExec>()
                 {
-                    let resolves_placeholders =
-                        transformer.has_rule::<ResolvePlaceholdersRule>();
-
-                    if resolves_placeholders {
+                    let has_rule = transformer.has_dyn_rule(rule);
+                    if has_rule {
+                        // Rule is already applied.
                         Ok(plan)
                     } else {
                         transformer
-                            .add_rule(Box::new(ResolvePlaceholdersRule::new()))
+                            .add_rule(Arc::clone(rule))
                             .map(|r| Arc::new(r) as Arc<_>)
                     }
                 } else {
-                    let transformer = TransformPlanExec::try_new(
-                        plan,
-                        vec![Box::new(ResolvePlaceholdersRule::new())],
-                    )?;
-
+                    let transformer =
+                        TransformPlanExec::try_new(plan, vec![Arc::clone(rule)])?;
                     if transformer.plans_to_transform() > 0 {
                         Ok(Arc::new(transformer))
                     } else {

datafusion/physical-optimizer/src/lib.rs

@@ -37,7 +37,7 @@ pub mod limit_pushdown_past_window;
 pub mod limited_distinct_aggregation;
 pub mod optimizer;
 pub mod output_requirements;
-pub mod physical_expr_resolver;
+pub mod exec_transform_apply;
 pub mod projection_pushdown;
 pub use datafusion_pruning as pruning;
 pub mod pushdown_sort;

datafusion/physical-optimizer/src/optimizer.rs

@@ -30,7 +30,7 @@ use crate::join_selection::JoinSelection;
 use crate::limit_pushdown::LimitPushdown;
 use crate::limited_distinct_aggregation::LimitedDistinctAggregation;
 use crate::output_requirements::OutputRequirements;
-use crate::physical_expr_resolver::PhysicalExprResolver;
+use crate::exec_transform_apply::ExecutionTransformationApplier;
 use crate::projection_pushdown::ProjectionPushdown;
 use crate::sanity_checker::SanityCheckPlan;
 use crate::topk_aggregation::TopKAggregation;
@@ -41,6 +41,7 @@ use crate::pushdown_sort::PushdownSort;
 use datafusion_common::Result;
 use datafusion_common::config::ConfigOptions;
 use datafusion_physical_plan::ExecutionPlan;
+use datafusion_physical_plan::plan_transformer::ResolvePlaceholdersRule;
 
 /// `PhysicalOptimizerRule` transforms one ['ExecutionPlan'] into another which
 /// computes the same results, but in a potentially more efficient way.
@@ -88,7 +89,7 @@ impl PhysicalOptimizer {
             // this information is not lost across different rules during optimization.
             Arc::new(OutputRequirements::new_add_mode()),
             // This rule removes all existing `TransformPlanExec` nodes from the plan tree.
-            Arc::new(PhysicalExprResolver::new()),
+            Arc::new(ExecutionTransformationApplier::new()),
             Arc::new(AggregateStatistics::new()),
             // Statistics-based join selection will change the Auto mode to a real join implementation,
             // like collect left, or hash join, or future sort merge join, which will influence the
@@ -151,7 +152,9 @@ impl PhysicalOptimizer {
             // This rule prepares the physical plan for placeholder resolution by wrapping it in a
             // `TransformPlanExec` with a `ResolvePlaceholdersRule` if it contains any unresolved
             // placeholders.
-            Arc::new(PhysicalExprResolver::new_post_optimization()),
+            Arc::new(ExecutionTransformationApplier::new_post_optimization(
+                Arc::new(ResolvePlaceholdersRule::new()),
+            )),
             // This FilterPushdown handles dynamic filters that may have references to the source ExecutionPlan.
             // Therefore it should be run at the end of the optimization process since any changes to the plan may break the dynamic filter's references.
             // See `FilterPushdownPhase` for more details.

datafusion/physical-plan/benches/plan_transformer.rs

@@ -46,11 +46,7 @@ impl ExecutionTransformationRule for ResetAllRule {
         self
     }
 
-    fn clone_box(&self) -> Box<dyn ExecutionTransformationRule> {
-        Box::new(self.clone())
-    }
-
-    fn matches(&mut self, _node: &Arc<dyn ExecutionPlan>) -> Result<bool> {
+    fn matches(&self, _node: &Arc<dyn ExecutionPlan>) -> Result<bool> {
         Ok(true)
     }
 
@@ -77,11 +73,7 @@ impl ExecutionTransformationRule for ResetByNameRule {
         self
     }
 
-    fn clone_box(&self) -> Box<dyn ExecutionTransformationRule> {
-        Box::new(self.clone())
-    }
-
-    fn matches(&mut self, node: &Arc<dyn ExecutionPlan>) -> Result<bool> {
+    fn matches(&self, node: &Arc<dyn ExecutionPlan>) -> Result<bool> {
         Ok(node.name() == self.node_name)
     }
 
@@ -166,7 +158,7 @@ fn benchmark_with_transformer_exec(
     group: &mut BenchmarkGroup<'_, WallTime>,
     batch_label: &str,
     plan: &Arc<dyn ExecutionPlan>,
-    rules: &[Box<dyn ExecutionTransformationRule>],
+    rules: &[Arc<dyn ExecutionTransformationRule>],
     batch_size: BatchSize,
 ) {
     let ctx = Arc::new(TaskContext::default());
@@ -182,7 +174,7 @@ fn benchmark_with_transformer_exec(
                 || {
                     (
                         Arc::clone(plan),
-                        rules.iter().map(|r| r.clone_box()).collect(),
+                        rules.to_vec(),
                     )
                 },
                 |(plan, rules)| {
@@ -201,8 +193,7 @@ fn benchmark_with_transformer_exec(
         ),
         |b| {
             let plan = Arc::clone(plan);
-            let rules = rules.iter().map(|r| r.clone_box()).collect();
-            let transformer = Arc::new(TransformPlanExec::try_new(plan, rules).unwrap());
+            let transformer = Arc::new(TransformPlanExec::try_new(plan, rules.to_vec()).unwrap());
 
             b.iter_batched(
                 || Arc::clone(&transformer),
@@ -257,14 +248,14 @@ fn criterion_benchmark(c: &mut Criterion) {
 
         for count in rules_count {
             let reset_all_rules = (0..count)
-                .map(|_| Box::new(ResetAllRule {}) as Box<_>)
+                .map(|_| Arc::new(ResetAllRule {}) as _)
                 .collect::<Vec<_>>();
 
             let reset_one_rules = (0..count)
                 .map(|_| {
-                    Box::new(ResetByNameRule {
+                    Arc::new(ResetByNameRule {
                         node_name: "EmptyExec".to_string(),
-                    }) as Box<_>
+                    }) as _
                 })
                 .collect::<Vec<_>>();