Move EnforceDistribution into datafusion-physical-optimizer crate (#14190)

* moving enforce_distribution

* formatting fix

* Update Cargo.toml

Fix circular dependency

* pulll source

* move_tests to core integration tests

* remove tests from enforce_distribution.rs

* passes lint

* forgotten license header

* move enforce_sorting

* import order

* merge fixes + formatting

* fix: forgotten license

* fix: cargo fmt

* fix tests

* fix: ci tests

* fix: Cargo.toml formatting

* further removals

* Further migrations and simplificaitons

* Fix failing tests

* Final pass

* Update datafusion-testing

* Update test_utils.rs

* fix the dep

* Update Cargo.toml

---------

Co-authored-by: buraksenn <[email protected]>
Co-authored-by: berkaysynnada <[email protected]>

Author: 3 people

datafusion-cli/Cargo.lock

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

datafusion/core/src/physical_optimizer/enforce_distribution.rs

@@ -44,11 +44,7 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 use arrow::record_batch::RecordBatch;
 use datafusion_catalog::Session;
 use datafusion_common::TableReference;
-use datafusion_expr::utils::COUNT_STAR_EXPANSION;
 use datafusion_expr::{CreateExternalTable, Expr, SortExpr, TableType};
-use datafusion_functions_aggregate::count::count_udaf;
-use datafusion_physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr};
-use datafusion_physical_expr::{expressions, PhysicalExpr};
 
 use async_trait::async_trait;
 use futures::Stream;
@@ -276,57 +272,3 @@ pub fn bounded_stream(batch: RecordBatch, limit: usize) -> SendableRecordBatchSt
         batch,
     })
 }
-
-/// Describe the type of aggregate being tested
-pub enum TestAggregate {
-    /// Testing COUNT(*) type aggregates
-    CountStar,
-
-    /// Testing for COUNT(column) aggregate
-    ColumnA(Arc<Schema>),
-}
-
-impl TestAggregate {
-    /// Create a new COUNT(*) aggregate
-    pub fn new_count_star() -> Self {
-        Self::CountStar
-    }
-
-    /// Create a new COUNT(column) aggregate
-    pub fn new_count_column(schema: &Arc<Schema>) -> Self {
-        Self::ColumnA(Arc::clone(schema))
-    }
-
-    /// Return appropriate expr depending if COUNT is for col or table (*)
-    pub fn count_expr(&self, schema: &Schema) -> AggregateFunctionExpr {
-        AggregateExprBuilder::new(count_udaf(), vec![self.column()])
-            .schema(Arc::new(schema.clone()))
-            .alias(self.column_name())
-            .build()
-            .unwrap()
-    }
-
-    /// what argument would this aggregate need in the plan?
-    fn column(&self) -> Arc<dyn PhysicalExpr> {
-        match self {
-            Self::CountStar => expressions::lit(COUNT_STAR_EXPANSION),
-            Self::ColumnA(s) => expressions::col("a", s).unwrap(),
-        }
-    }
-
-    /// What name would this aggregate produce in a plan?
-    pub fn column_name(&self) -> &'static str {
-        match self {
-            Self::CountStar => "COUNT(*)",
-            Self::ColumnA(_) => "COUNT(a)",
-        }
-    }
-
-    /// What is the expected count?
-    pub fn expected_count(&self) -> i64 {
-        match self {
-            TestAggregate::CountStar => 3,
-            TestAggregate::ColumnA(_) => 2,
-        }
-    }
-}

datafusion/core/src/physical_optimizer/mod.rs

@@ -42,7 +42,6 @@ mod custom_sources_cases;
 /// Run all tests that are found in the `optimizer` directory
 mod optimizer;
 
-/// Run all tests that are found in the `physical_optimizer` directory
 mod physical_optimizer;
 
 #[cfg(test)]

datafusion/core/src/test_util/mod.rs

@@ -19,14 +19,15 @@
 //!
 //! Note these tests are not in the same module as the optimizer pass because
 //! they rely on `ParquetExec` which is in the core crate.
+
 use std::sync::Arc;
 
+use crate::physical_optimizer::parquet_exec;
+
 use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
-use datafusion::datasource::listing::PartitionedFile;
-use datafusion::datasource::physical_plan::{FileScanConfig, ParquetExec};
 use datafusion::physical_optimizer::combine_partial_final_agg::CombinePartialFinalAggregate;
+use datafusion::physical_optimizer::test_utils::trim_plan_display;
 use datafusion_common::config::ConfigOptions;
-use datafusion_execution::object_store::ObjectStoreUrl;
 use datafusion_functions_aggregate::count::count_udaf;
 use datafusion_functions_aggregate::sum::sum_udaf;
 use datafusion_physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr};
@@ -62,13 +63,6 @@ macro_rules! assert_optimized {
     };
 }
 
-fn trim_plan_display(plan: &str) -> Vec<&str> {
-    plan.split('\n')
-        .map(|s| s.trim())
-        .filter(|s| !s.is_empty())
-        .collect()
-}
-
 fn schema() -> SchemaRef {
     Arc::new(Schema::new(vec![
         Field::new("a", DataType::Int64, true),
@@ -77,14 +71,6 @@ fn schema() -> SchemaRef {
     ]))
 }
 
-fn parquet_exec(schema: &SchemaRef) -> Arc<ParquetExec> {
-    ParquetExec::builder(
-        FileScanConfig::new(ObjectStoreUrl::parse("test:///").unwrap(), schema.clone())
-            .with_file(PartitionedFile::new("x".to_string(), 100)),
-    )
-    .build_arc()
-}
-
 fn partial_aggregate_exec(
     input: Arc<dyn ExecutionPlan>,
     group_by: PhysicalGroupBy,