consolidate tests and utils (#14298)

Author: berkaysynnada

datafusion-cli/Cargo.lock

@@ -15,31 +15,30 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use datafusion_common::config::ConfigOptions;
-use datafusion_execution::TaskContext;
-use datafusion_physical_optimizer::aggregate_statistics::AggregateStatistics;
-use datafusion_physical_optimizer::PhysicalOptimizerRule;
-use datafusion_physical_plan::aggregates::AggregateExec;
-use datafusion_physical_plan::projection::ProjectionExec;
-use datafusion_physical_plan::ExecutionPlan;
 use std::sync::Arc;
 
-use datafusion_common::Result;
-
-use datafusion_physical_plan::aggregates::PhysicalGroupBy;
-use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
-use datafusion_physical_plan::common;
-use datafusion_physical_plan::filter::FilterExec;
-use datafusion_physical_plan::memory::MemoryExec;
+use crate::physical_optimizer::test_utils::TestAggregate;
 
 use arrow::array::Int32Array;
 use arrow::datatypes::{DataType, Field, Schema};
 use arrow::record_batch::RecordBatch;
 use datafusion_common::cast::as_int64_array;
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::Result;
+use datafusion_execution::TaskContext;
 use datafusion_expr::Operator;
 use datafusion_physical_expr::expressions::{self, cast};
-use datafusion_physical_optimizer::test_utils::TestAggregate;
+use datafusion_physical_optimizer::aggregate_statistics::AggregateStatistics;
+use datafusion_physical_optimizer::PhysicalOptimizerRule;
+use datafusion_physical_plan::aggregates::AggregateExec;
 use datafusion_physical_plan::aggregates::AggregateMode;
+use datafusion_physical_plan::aggregates::PhysicalGroupBy;
+use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
+use datafusion_physical_plan::common;
+use datafusion_physical_plan::filter::FilterExec;
+use datafusion_physical_plan::memory::MemoryExec;
+use datafusion_physical_plan::projection::ProjectionExec;
+use datafusion_physical_plan::ExecutionPlan;
 
 /// Mock data using a MemoryExec which has an exact count statistic
 fn mock_data() -> Result<Arc<MemoryExec>> {

datafusion/core/tests/physical_optimizer/aggregate_statistics.rs

@@ -22,11 +22,10 @@
 
 use std::sync::Arc;
 
-use crate::physical_optimizer::parquet_exec;
+use crate::physical_optimizer::test_utils::{parquet_exec, trim_plan_display};
 
 use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 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_functions_aggregate::count::count_udaf;
 use datafusion_functions_aggregate::sum::sum_udaf;

datafusion/core/tests/physical_optimizer/combine_partial_final_agg.rs

@@ -19,7 +19,11 @@ use std::fmt::Debug;
 use std::ops::Deref;
 use std::sync::Arc;
 
-use crate::physical_optimizer::parquet_exec_with_sort;
+use crate::physical_optimizer::test_utils::{
+    check_integrity, coalesce_partitions_exec, repartition_exec, schema,
+    sort_merge_join_exec, sort_preserving_merge_exec,
+};
+use crate::physical_optimizer::test_utils::{parquet_exec_with_sort, trim_plan_display};
 
 use arrow::compute::SortOptions;
 use datafusion::config::ConfigOptions;
@@ -40,11 +44,6 @@ use datafusion_physical_expr_common::sort_expr::LexRequirement;
 use datafusion_physical_optimizer::enforce_distribution::*;
 use datafusion_physical_optimizer::enforce_sorting::EnforceSorting;
 use datafusion_physical_optimizer::output_requirements::OutputRequirements;
-use datafusion_physical_optimizer::test_utils::trim_plan_display;
-use datafusion_physical_optimizer::test_utils::{
-    check_integrity, coalesce_partitions_exec, repartition_exec, schema,
-    sort_merge_join_exec, sort_preserving_merge_exec,
-};
 use datafusion_physical_optimizer::PhysicalOptimizerRule;
 use datafusion_physical_plan::aggregates::{
     AggregateExec, AggregateMode, PhysicalGroupBy,
@@ -293,7 +292,7 @@ fn hash_join_exec(
     join_on: &JoinOn,
     join_type: &JoinType,
 ) -> Arc<dyn ExecutionPlan> {
-    datafusion_physical_optimizer::test_utils::hash_join_exec(
+    crate::physical_optimizer::test_utils::hash_join_exec(
         left,
         right,
         join_on.clone(),

datafusion/core/tests/physical_optimizer/enforce_distribution.rs

@@ -17,7 +17,14 @@
 
 use std::sync::Arc;
 
-use crate::physical_optimizer::parquet_exec;
+use crate::physical_optimizer::test_utils::{
+    aggregate_exec, bounded_window_exec, check_integrity, coalesce_batches_exec,
+    coalesce_partitions_exec, create_test_schema, create_test_schema2,
+    create_test_schema3, filter_exec, global_limit_exec, hash_join_exec, limit_exec,
+    local_limit_exec, memory_exec, parquet_exec, repartition_exec, sort_exec, sort_expr,
+    sort_expr_options, sort_merge_join_exec, sort_preserving_merge_exec,
+    spr_repartition_exec, stream_exec_ordered, union_exec, RequirementsTestExec,
+};
 
 use datafusion_physical_plan::displayable;
 use arrow::compute::SortOptions;
@@ -37,15 +44,14 @@ use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeE
 use datafusion_physical_plan::{get_plan_string, ExecutionPlan};
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::tree_node::{TreeNode, TransformedResult};
-use datafusion_physical_optimizer::test_utils::{check_integrity, bounded_window_exec, coalesce_partitions_exec, create_test_schema, create_test_schema2, create_test_schema3, filter_exec, global_limit_exec, hash_join_exec, limit_exec, local_limit_exec, memory_exec, repartition_exec, sort_exec, sort_expr, sort_expr_options, sort_merge_join_exec, sort_preserving_merge_exec, spr_repartition_exec, stream_exec_ordered, union_exec, coalesce_batches_exec, aggregate_exec, RequirementsTestExec};
 use datafusion::datasource::physical_plan::{CsvExec, FileScanConfig, ParquetExec};
 use datafusion_execution::object_store::ObjectStoreUrl;
 use datafusion::datasource::listing::PartitionedFile;
 use datafusion::datasource::file_format::file_compression_type::FileCompressionType;
-
 use datafusion_physical_optimizer::enforce_distribution::EnforceDistribution;
 use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
 use datafusion_physical_plan::sorts::sort::SortExec;
+
 use rstest::rstest;
 
 /// Create a csv exec for tests