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

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-//! Integration tests for [`PhysicalExprResolver`] optimizer rule.
+//! Integration tests for [`ExecutionTransformationApplier`] optimizer rule.
 
 use std::{collections::HashMap, sync::Arc};
 
@@ -29,11 +29,14 @@ use datafusion_physical_expr::{
     expressions::{BinaryExpr, col, lit, placeholder},
 };
 use datafusion_physical_optimizer::{
-    PhysicalOptimizerRule, physical_expr_resolver::PhysicalExprResolver,
+    PhysicalOptimizerRule, exec_transform_apply::ExecutionTransformationApplier,
 };
 use datafusion_physical_plan::{
-    ExecutionPlan, filter::FilterExec, get_plan_string,
-    plan_transformer::TransformPlanExec, repartition::RepartitionExec,
+    ExecutionPlan,
+    filter::FilterExec,
+    get_plan_string,
+    plan_transformer::{ResolvePlaceholdersRule, TransformPlanExec},
+    repartition::RepartitionExec,
 };
 
 use crate::physical_optimizer::test_utils::{
@@ -85,6 +88,12 @@ fn repartition_exec(
     )?))
 }
 
+fn placeholder_resolver() -> ExecutionTransformationApplier {
+    ExecutionTransformationApplier::new_post_optimization(Arc::new(
+        ResolvePlaceholdersRule::new(),
+    ))
+}
+
 #[test]
 fn test_noop_if_no_placeholders_found() -> Result<()> {
     let schema = create_schema();
@@ -105,8 +114,7 @@ fn test_noop_if_no_placeholders_found() -> Result<()> {
 
     assert_eq!(initial, expected_initial);
 
-    let after_optimize = PhysicalExprResolver::new_post_optimization()
-        .optimize(plan, &ConfigOptions::new())?;
+    let after_optimize = placeholder_resolver().optimize(plan, &ConfigOptions::new())?;
 
     let optimized_plan_string = get_plan_string(&after_optimize);
     assert_eq!(initial, optimized_plan_string);
@@ -134,8 +142,7 @@ fn test_wrap_plan_with_transformer() -> Result<()> {
 
     assert_eq!(initial, expected_initial);
 
-    let after_optimize = PhysicalExprResolver::new_post_optimization()
-        .optimize(plan, &ConfigOptions::new())?;
+    let after_optimize = placeholder_resolver().optimize(plan, &ConfigOptions::new())?;
 
     let expected_optimized = [
         "TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1]",
@@ -201,7 +208,7 @@ fn test_remove_useless_transformers() -> Result<()> {
     assert_eq!(initial, expected_initial);
 
     let after_optimize =
-        PhysicalExprResolver::new().optimize(plan, &ConfigOptions::new())?;
+        ExecutionTransformationApplier::new().optimize(plan, &ConfigOptions::new())?;
 
     let expected_optimized = [
         "GlobalLimitExec: skip=0, fetch=5",
@@ -244,7 +251,7 @@ fn test_combine_transformers() -> Result<()> {
     assert_eq!(initial, expected_initial);
 
     let after_pre_optimization =
-        PhysicalExprResolver::new().optimize(plan, &ConfigOptions::new())?;
+        ExecutionTransformationApplier::new().optimize(plan, &ConfigOptions::new())?;
 
     let expected_optimized = [
         "GlobalLimitExec: skip=0, fetch=5",
@@ -257,8 +264,8 @@ fn test_combine_transformers() -> Result<()> {
     let optimized_plan_string = get_plan_string(&after_pre_optimization);
     assert_eq!(optimized_plan_string, expected_optimized);
 
-    let after_post_optimization = PhysicalExprResolver::new_post_optimization()
-        .optimize(after_pre_optimization, &ConfigOptions::new())?;
+    let after_post_optimization =
+        placeholder_resolver().optimize(after_pre_optimization, &ConfigOptions::new())?;
 
     let expected_optimized = [
         "TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1]",

datafusion/core/tests/physical_optimizer/test_utils.rs

@@ -400,7 +400,7 @@ pub fn resolve_placeholders_exec(
     input: Arc<dyn ExecutionPlan>,
 ) -> Arc<dyn ExecutionPlan> {
     Arc::new(
-        TransformPlanExec::try_new(input, vec![Box::new(ResolvePlaceholdersRule::new())])
+        TransformPlanExec::try_new(input, vec![Arc::new(ResolvePlaceholdersRule::new())])
             .unwrap(),
     )
 }

datafusion/datasource/src/values.rs

@@ -437,7 +437,7 @@ mod tests {
         // Should be ValuesSource because of placeholder.
         assert!(values_exec.data_source().as_any().is::<ValuesSource>());
 
-        let rules = vec![Box::new(ResolvePlaceholdersRule::new()) as Box<_>];
+        let rules = vec![Arc::new(ResolvePlaceholdersRule::new()) as _];
         let exec = Arc::new(TransformPlanExec::try_new(values_exec, rules)?);
         let task_ctx = Arc::new(TaskContext::default().with_param_values(
             ParamValues::List(vec![ScalarValue::Int32(Some(10)).into()]),
@@ -470,7 +470,7 @@ mod tests {
         ]];
 
         let values_exec = ValuesSource::try_new_exec(Arc::clone(&schema), data)?;
-        let rules = vec![Box::new(ResolvePlaceholdersRule::new()) as Box<_>];
+        let rules = vec![Arc::new(ResolvePlaceholdersRule::new()) as _];
         let exec = Arc::new(TransformPlanExec::try_new(values_exec, rules)?) as Arc<_>;
 
         let task_ctx = Arc::new(TaskContext::default().with_param_values(
@@ -528,7 +528,7 @@ mod tests {
             vec![vec![lit(10), placeholder("$foo", DataType::Int32)]];
 
         let values_exec = ValuesSource::try_new_exec(Arc::clone(&schema), data)?;
-        let rules = vec![Box::new(ResolvePlaceholdersRule::new()) as Box<_>];
+        let rules = vec![Arc::new(ResolvePlaceholdersRule::new()) as _];
         let exec = Arc::new(TransformPlanExec::try_new(values_exec, rules)?) as Arc<_>;
 
         let task_ctx = Arc::new(TaskContext::default());

…-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,71 +27,72 @@ 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)",
-        }
+        &self.name
     }
 
     fn optimize(
         &self,
         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 +102,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

@@ -30,14 +30,14 @@ pub mod combine_partial_final_agg;
 pub mod enforce_distribution;
 pub mod enforce_sorting;
 pub mod ensure_coop;
+pub mod exec_transform_apply;
 pub mod filter_pushdown;
 pub mod join_selection;
 pub mod limit_pushdown;
 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 projection_pushdown;
 pub use datafusion_pruning as pruning;
 pub mod pushdown_sort;

datafusion/physical-optimizer/src/optimizer.rs

@@ -25,12 +25,12 @@ use crate::combine_partial_final_agg::CombinePartialFinalAggregate;
 use crate::enforce_distribution::EnforceDistribution;
 use crate::enforce_sorting::EnforceSorting;
 use crate::ensure_coop::EnsureCooperative;
+use crate::exec_transform_apply::ExecutionTransformationApplier;
 use crate::filter_pushdown::FilterPushdown;
 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::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.