Move EnforceDistribution into datafusion-physical-optimizer crate

datafusion/core/src/physical_optimizer/mod.rs

@@ -22,7 +22,6 @@
 //!
 //! [`ExecutionPlan`]: crate::physical_plan::ExecutionPlan
 
-pub mod enforce_distribution;
 pub mod enforce_sorting;
 pub mod optimizer;
 pub mod projection_pushdown;

datafusion/physical-optimizer/Cargo.toml

@@ -33,6 +33,7 @@ workspace = true
 
 [features]
 recursive_protection = ["dep:recursive"]
+parquet = ["dep:parquet"]
 
 [dependencies]
 arrow = { workspace = true }
@@ -48,6 +49,7 @@ futures = { workspace = true }
 itertools = { workspace = true }
 log = { workspace = true }
 recursive = { workspace = true, optional = true }
+parquet = { workspace = true, optional = true, default-features = true }
 
 [dev-dependencies]
 datafusion-expr = { workspace = true }

datafusion/core/src/physical_optimizer/enforce_distribution.rs → datafusion/physical-optimizer/src/enforce_distribution.rs

@@ -24,24 +24,30 @@
 use std::fmt::Debug;
 use std::sync::Arc;
 
-use crate::config::ConfigOptions;
-use crate::error::Result;
-use crate::physical_optimizer::utils::{
+use crate::optimizer::PhysicalOptimizerRule;
+use crate::output_requirements::OutputRequirementExec;
+use crate::utils::{
     add_sort_above_with_check, is_coalesce_partitions, is_repartition,
     is_sort_preserving_merge,
 };
-use crate::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy};
-use crate::physical_plan::coalesce_partitions::CoalescePartitionsExec;
-use crate::physical_plan::joins::{
+
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::error::Result;
+
+use datafusion_physical_plan::aggregates::{
+    AggregateExec, AggregateMode, PhysicalGroupBy,
+};
+use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
+use datafusion_physical_plan::joins::{
     CrossJoinExec, HashJoinExec, PartitionMode, SortMergeJoinExec,
 };
-use crate::physical_plan::projection::ProjectionExec;
-use crate::physical_plan::repartition::RepartitionExec;
-use crate::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
-use crate::physical_plan::tree_node::PlanContext;
-use crate::physical_plan::union::{can_interleave, InterleaveExec, UnionExec};
-use crate::physical_plan::windows::WindowAggExec;
-use crate::physical_plan::{Distribution, ExecutionPlan, Partitioning};
+use datafusion_physical_plan::projection::ProjectionExec;
+use datafusion_physical_plan::repartition::RepartitionExec;
+use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
+use datafusion_physical_plan::tree_node::PlanContext;
+use datafusion_physical_plan::union::{can_interleave, InterleaveExec, UnionExec};
+use datafusion_physical_plan::windows::WindowAggExec;
+use datafusion_physical_plan::{Distribution, ExecutionPlan, Partitioning};
 
 use arrow::compute::SortOptions;
 use datafusion_common::stats::Precision;
@@ -53,8 +59,8 @@ use datafusion_physical_expr::{
     physical_exprs_equal, EquivalenceProperties, PhysicalExpr, PhysicalExprRef,
 };
 use datafusion_physical_expr_common::sort_expr::LexOrdering;
-use datafusion_physical_optimizer::output_requirements::OutputRequirementExec;
-use datafusion_physical_optimizer::PhysicalOptimizerRule;
+// use datafusion_physical_optimizer::output_requirements::OutputRequirementExec;
+// use datafusion_physical_optimizer::PhysicalOptimizerRule;
 use datafusion_physical_plan::execution_plan::EmissionType;
 use datafusion_physical_plan::windows::{get_best_fitting_window, BoundedWindowAggExec};
 use datafusion_physical_plan::ExecutionPlanProperties;

datafusion/physical-optimizer/src/lib.rs

@@ -21,6 +21,7 @@
 pub mod aggregate_statistics;
 pub mod coalesce_batches;
 pub mod combine_partial_final_agg;
+pub mod enforce_distribution;
 pub mod join_selection;
 pub mod limit_pushdown;
 pub mod limited_distinct_aggregation;
@@ -31,5 +32,5 @@ pub mod sanity_checker;
 pub mod test_utils;
 pub mod topk_aggregation;
 pub mod update_aggr_exprs;
-
 pub use optimizer::PhysicalOptimizerRule;
+mod utils;

datafusion/physical-optimizer/src/utils.rs

@@ -0,0 +1,110 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Collection of utility functions that are leveraged by the query optimizer rules
+
+use std::sync::Arc;
+
+use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
+use datafusion_physical_plan::repartition::RepartitionExec;
+use datafusion_physical_plan::sorts::sort::SortExec;
+use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
+// use datafusion_physical_plan::union::UnionExec;
+// use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowAggExec};
+use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties};
+
+use datafusion_physical_expr::LexRequirement;
+use datafusion_physical_expr_common::sort_expr::LexOrdering;
+// use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
+use datafusion_physical_plan::tree_node::PlanContext;
+
+/// This utility function adds a `SortExec` above an operator according to the
+/// given ordering requirements while preserving the original partitioning.
+pub fn add_sort_above<T: Clone + Default>(
+    node: PlanContext<T>,
+    sort_requirements: LexRequirement,
+    fetch: Option<usize>,
+) -> PlanContext<T> {
+    let mut sort_expr = LexOrdering::from(sort_requirements);
+    sort_expr.retain(|sort_expr| {
+        !node
+            .plan
+            .equivalence_properties()
+            .is_expr_constant(&sort_expr.expr)
+    });
+    let mut new_sort = SortExec::new(sort_expr, Arc::clone(&node.plan)).with_fetch(fetch);
+    if node.plan.output_partitioning().partition_count() > 1 {
+        new_sort = new_sort.with_preserve_partitioning(true);
+    }
+    PlanContext::new(Arc::new(new_sort), T::default(), vec![node])
+}
+
+/// This utility function adds a `SortExec` above an operator according to the
+/// given ordering requirements while preserving the original partitioning. If
+/// requirement is already satisfied no `SortExec` is added.
+pub fn add_sort_above_with_check<T: Clone + Default>(
+    node: PlanContext<T>,
+    sort_requirements: LexRequirement,
+    fetch: Option<usize>,
+) -> PlanContext<T> {
+    if !node
+        .plan
+        .equivalence_properties()
+        .ordering_satisfy_requirement(&sort_requirements)
+    {
+        add_sort_above(node, sort_requirements, fetch)
+    } else {
+        node
+    }
+}
+
+/// Checks whether the given operator is a limit;
+/// i.e. either a [`LocalLimitExec`] or a [`GlobalLimitExec`].
+// pub fn is_limit(plan: &Arc<dyn ExecutionPlan>) -> bool {
+//     plan.as_any().is::<GlobalLimitExec>() || plan.as_any().is::<LocalLimitExec>()
+// }
+
+/// Checks whether the given operator is a window;
+/// i.e. either a [`WindowAggExec`] or a [`BoundedWindowAggExec`].
+// pub fn is_window(plan: &Arc<dyn ExecutionPlan>) -> bool {
+//     plan.as_any().is::<WindowAggExec>() || plan.as_any().is::<BoundedWindowAggExec>()
+// }
+
+/// Checks whether the given operator is a [`SortExec`].
+// pub fn is_sort(plan: &Arc<dyn ExecutionPlan>) -> bool {
+//     plan.as_any().is::<SortExec>()
+// }
+
+/// Checks whether the given operator is a [`SortPreservingMergeExec`].
+pub fn is_sort_preserving_merge(plan: &Arc<dyn ExecutionPlan>) -> bool {
+    plan.as_any().is::<SortPreservingMergeExec>()
+}
+
+/// Checks whether the given operator is a [`CoalescePartitionsExec`].
+pub fn is_coalesce_partitions(plan: &Arc<dyn ExecutionPlan>) -> bool {
+    plan.as_any().is::<CoalescePartitionsExec>()
+}
+
+/// Checks whether the given operator is a [`UnionExec`].
+// pub fn is_union(plan: &Arc<dyn ExecutionPlan>) -> bool {
+//     plan.as_any().is::<UnionExec>()
+// }
+
+/// Checks whether the given operator is a [`RepartitionExec`].
+pub fn is_repartition(plan: &Arc<dyn ExecutionPlan>) -> bool {
+    plan.as_any().is::<RepartitionExec>()
+}