Update all tests to use the new task estimation mechanism

gabotechs · gabotechs · commit e932754111f8 · 2025-11-04T15:49:57.000-05:00
diff --git a/examples/in_memory_cluster.rs b/examples/in_memory_cluster.rs
@@ -5,7 +5,6 @@ use datafusion::common::DataFusionError;
 use datafusion::execution::SessionStateBuilder;
 use datafusion::physical_plan::displayable;
 use datafusion::prelude::{ParquetReadOptions, SessionContext};
-use datafusion_distributed::test_utils::test_task_estimator::FixedDataSourceExecTaskEstimator;
 use datafusion_distributed::{
     ArrowFlightEndpoint, BoxCloneSyncChannel, ChannelResolver, DistributedExt,
     DistributedSessionBuilderContext, create_flight_client,
@@ -27,9 +26,6 @@ struct Args {
 
     #[structopt(long)]
     explain: bool,
-
-    #[structopt(long, default_value = "3")]
-    stage_tasks: usize,
 }
 
 #[tokio::main]
@@ -39,7 +35,6 @@ async fn main() -> Result<(), Box<dyn Error>> {
     let state = SessionStateBuilder::new()
         .with_default_features()
         .with_distributed_execution(InMemoryChannelResolver::new())
-        .with_distributed_task_estimator(FixedDataSourceExecTaskEstimator(args.stage_tasks))
         .build();
 
     let ctx = SessionContext::from(state);
diff --git a/examples/localhost_run.rs b/examples/localhost_run.rs
@@ -6,7 +6,6 @@ use datafusion::common::DataFusionError;
 use datafusion::execution::SessionStateBuilder;
 use datafusion::physical_plan::displayable;
 use datafusion::prelude::{ParquetReadOptions, SessionContext};
-use datafusion_distributed::test_utils::test_task_estimator::FixedDataSourceExecTaskEstimator;
 use datafusion_distributed::{BoxCloneSyncChannel, ChannelResolver, DistributedExt};
 use futures::TryStreamExt;
 use std::error::Error;
@@ -26,9 +25,6 @@ struct Args {
 
     #[structopt(long)]
     explain: bool,
-
-    #[structopt(long, default_value = "3")]
-    stage_tasks: usize,
 }
 
 #[tokio::main]
@@ -43,7 +39,6 @@ async fn main() -> Result<(), Box<dyn Error>> {
     let state = SessionStateBuilder::new()
         .with_default_features()
         .with_distributed_execution(localhost_resolver)
-        .with_distributed_task_estimator(FixedDataSourceExecTaskEstimator(args.stage_tasks))
         .build();
 
     let ctx = SessionContext::from(state);
diff --git a/src/distributed_planner/distributed_config.rs b/src/distributed_planner/distributed_config.rs
@@ -15,13 +15,13 @@ extensions_options! {
     pub struct DistributedConfig {
         /// Sets the maximum amount of files that will be assigned to each task. Reducing this
         /// number will spawn more tasks for the same number of files.
-        pub files_per_task: usize, default = get_available_parallelism()
+        pub files_per_task: usize, default = files_per_task_default()
         /// Task multiplying factor for when a node declares that it changes the cardinality
         /// of the data:
         /// - If a node is increasing the cardinality of the data, this factor will increase.
         /// - If a node reduces the cardinality of the data, this factor will decrease.
         /// - In any other situation, this factor is left intact.
-        pub cardinality_task_count_factor: f64, default = 1.0
+        pub cardinality_task_count_factor: f64, default = cardinality_task_count_factor_default()
         /// Collection of [TaskEstimator]s that will be applied to leaf nodes in order to
         /// estimate how many tasks should be spawned for the [Stage] containing the leaf node.
         pub(crate) __private_task_estimator: CombinedTaskEstimator, default = CombinedTaskEstimator::default()
@@ -31,6 +31,22 @@ extensions_options! {
     }
 }
 
+fn files_per_task_default() -> usize {
+    if cfg!(test) || cfg!(feature = "integration") {
+        1
+    } else {
+        get_available_parallelism()
+    }
+}
+
+fn cardinality_task_count_factor_default() -> f64 {
+    if cfg!(test) || cfg!(feature = "integration") {
+        1.5
+    } else {
+        1.0
+    }
+}
+
 impl DistributedConfig {
     /// Appends a [TaskEstimator] to the list. [TaskEstimator] will be executed sequentially in
     /// order on leaf nodes, and the first one to provide a value is the one that gets to decide
diff --git a/src/distributed_planner/distributed_physical_optimizer_rule.rs b/src/distributed_planner/distributed_physical_optimizer_rule.rs
@@ -401,7 +401,6 @@ mod tests {
     use crate::DistributedExt;
     use crate::test_utils::in_memory_channel_resolver::InMemoryChannelResolver;
     use crate::test_utils::parquet::register_parquet_tables;
-    use crate::test_utils::test_task_estimator::FixedDataSourceExecTaskEstimator;
     use crate::{assert_snapshot, display_plan_ascii};
     use datafusion::error::DataFusionError;
     use datafusion::execution::SessionStateBuilder;
@@ -458,13 +457,13 @@ mod tests {
           │   ProjectionExec: expr=[count(Int64(1))@1 as count(*), RainToday@0 as RainToday, count(Int64(1))@1 as count(Int64(1))]
           │     AggregateExec: mode=FinalPartitioned, gby=[RainToday@0 as RainToday], aggr=[count(Int64(1))]
           │       CoalesceBatchesExec: target_batch_size=8192
-          │         [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=2
+          │         [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=3
           └──────────────────────────────────────────────────
-            ┌───── Stage 1 ── Tasks: t0:[p0..p7] t1:[p0..p7] 
+            ┌───── Stage 1 ── Tasks: t0:[p0..p7] t1:[p0..p7] t2:[p0..p7] 
             │ RepartitionExec: partitioning=Hash([RainToday@0], 8), input_partitions=4
-            │   RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=2
+            │   RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
             │     AggregateExec: mode=Partial, gby=[RainToday@0 as RainToday], aggr=[count(Int64(1))]
-            │       PartitionIsolatorExec: t0:[p0,p1,__] t1:[__,__,p0] 
+            │       PartitionIsolatorExec: t0:[p0,__,__] t1:[__,p0,__] t2:[__,__,p0] 
             │         DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet], [/testdata/weather/result-000001.parquet], [/testdata/weather/result-000002.parquet]]}, projection=[RainToday], file_type=parquet
             └──────────────────────────────────────────────────
         ");
@@ -487,13 +486,13 @@ mod tests {
           │   ProjectionExec: expr=[count(Int64(1))@1 as count(*), RainToday@0 as RainToday, count(Int64(1))@1 as count(Int64(1))]
           │     AggregateExec: mode=FinalPartitioned, gby=[RainToday@0 as RainToday], aggr=[count(Int64(1))]
           │       CoalesceBatchesExec: target_batch_size=8192
-          │         [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=2
+          │         [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=3
           └──────────────────────────────────────────────────
-            ┌───── Stage 1 ── Tasks: t0:[p0..p7] t1:[p0..p7] 
+            ┌───── Stage 1 ── Tasks: t0:[p0..p7] t1:[p0..p7] t2:[p0..p7] 
             │ RepartitionExec: partitioning=Hash([RainToday@0], 8), input_partitions=4
-            │   RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=2
+            │   RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
             │     AggregateExec: mode=Partial, gby=[RainToday@0 as RainToday], aggr=[count(Int64(1))]
-            │       PartitionIsolatorExec: t0:[p0,p1,__] t1:[__,__,p0] 
+            │       PartitionIsolatorExec: t0:[p0,__,__] t1:[__,p0,__] t2:[__,__,p0] 
             │         DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet], [/testdata/weather/result-000001.parquet], [/testdata/weather/result-000002.parquet]]}, projection=[RainToday], file_type=parquet
             └──────────────────────────────────────────────────
         ");
@@ -550,30 +549,30 @@ mod tests {
         │       ProjectionExec: expr=[avg(weather.MaxTemp)@1 as MaxTemp, RainTomorrow@0 as RainTomorrow]
         │         AggregateExec: mode=FinalPartitioned, gby=[RainTomorrow@0 as RainTomorrow], aggr=[avg(weather.MaxTemp)]
         │           CoalesceBatchesExec: target_batch_size=8192
-        │             [Stage 3] => NetworkShuffleExec: output_partitions=4, input_tasks=2
+        │             [Stage 3] => NetworkShuffleExec: output_partitions=4, input_tasks=3
         └──────────────────────────────────────────────────
           ┌───── Stage 2 ── Tasks: t0:[p0..p3] t1:[p0..p3] 
           │ ProjectionExec: expr=[avg(weather.MinTemp)@1 as MinTemp, RainTomorrow@0 as RainTomorrow]
           │   AggregateExec: mode=FinalPartitioned, gby=[RainTomorrow@0 as RainTomorrow], aggr=[avg(weather.MinTemp)]
           │     CoalesceBatchesExec: target_batch_size=8192
-          │       [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=2
+          │       [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=3
           └──────────────────────────────────────────────────
-            ┌───── Stage 1 ── Tasks: t0:[p0..p7] t1:[p0..p7] 
+            ┌───── Stage 1 ── Tasks: t0:[p0..p7] t1:[p0..p7] t2:[p0..p7] 
             │ RepartitionExec: partitioning=Hash([RainTomorrow@0], 8), input_partitions=4
             │   AggregateExec: mode=Partial, gby=[RainTomorrow@1 as RainTomorrow], aggr=[avg(weather.MinTemp)]
             │     CoalesceBatchesExec: target_batch_size=8192
             │       FilterExec: RainToday@1 = yes, projection=[MinTemp@0, RainTomorrow@2]
-            │         RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=2
-            │           PartitionIsolatorExec: t0:[p0,p1,__] t1:[__,__,p0] 
+            │         RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+            │           PartitionIsolatorExec: t0:[p0,__,__] t1:[__,p0,__] t2:[__,__,p0] 
             │             DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet], [/testdata/weather/result-000001.parquet], [/testdata/weather/result-000002.parquet]]}, projection=[MinTemp, RainToday, RainTomorrow], file_type=parquet, predicate=RainToday@1 = yes, pruning_predicate=RainToday_null_count@2 != row_count@3 AND RainToday_min@0 <= yes AND yes <= RainToday_max@1, required_guarantees=[RainToday in (yes)]
             └──────────────────────────────────────────────────
-          ┌───── Stage 3 ── Tasks: t0:[p0..p3] t1:[p0..p3] 
+          ┌───── Stage 3 ── Tasks: t0:[p0..p3] t1:[p0..p3] t2:[p0..p3] 
           │ RepartitionExec: partitioning=Hash([RainTomorrow@0], 4), input_partitions=4
           │   AggregateExec: mode=Partial, gby=[RainTomorrow@1 as RainTomorrow], aggr=[avg(weather.MaxTemp)]
           │     CoalesceBatchesExec: target_batch_size=8192
           │       FilterExec: RainToday@1 = no, projection=[MaxTemp@0, RainTomorrow@2]
-          │         RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=2
-          │           PartitionIsolatorExec: t0:[p0,p1,__] t1:[__,__,p0] 
+          │         RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+          │           PartitionIsolatorExec: t0:[p0,__,__] t1:[__,p0,__] t2:[__,__,p0] 
           │             DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet], [/testdata/weather/result-000001.parquet], [/testdata/weather/result-000002.parquet]]}, projection=[MaxTemp, RainToday, RainTomorrow], file_type=parquet, predicate=RainToday@1 = no, pruning_predicate=RainToday_null_count@2 != row_count@3 AND RainToday_min@0 <= no AND no <= RainToday_max@1, required_guarantees=[RainToday in (no)]
           └──────────────────────────────────────────────────
         ");
@@ -588,11 +587,11 @@ mod tests {
         assert_snapshot!(plan, @r"
         ┌───── DistributedExec ── Tasks: t0:[p0] 
         │ SortPreservingMergeExec: [MinTemp@0 DESC]
-        │   [Stage 1] => NetworkCoalesceExec: output_partitions=4, input_tasks=2
+        │   [Stage 1] => NetworkCoalesceExec: output_partitions=3, input_tasks=3
         └──────────────────────────────────────────────────
-          ┌───── Stage 1 ── Tasks: t0:[p0..p1] t1:[p2..p3] 
+          ┌───── Stage 1 ── Tasks: t0:[p0] t1:[p1] t2:[p2] 
           │ SortExec: expr=[MinTemp@0 DESC], preserve_partitioning=[true]
-          │   PartitionIsolatorExec: t0:[p0,p1,__] t1:[__,__,p0] 
+          │   PartitionIsolatorExec: t0:[p0,__,__] t1:[__,p0,__] t2:[__,__,p0] 
           │     DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet], [/testdata/weather/result-000001.parquet], [/testdata/weather/result-000002.parquet]]}, projection=[MinTemp, MaxTemp, Rainfall, Evaporation, Sunshine, WindGustDir, WindGustSpeed, WindDir9am, WindDir3pm, WindSpeed9am, WindSpeed3pm, Humidity9am, Humidity3pm, Pressure9am, Pressure3pm, Cloud9am, Cloud3pm, Temp9am, Temp3pm, RainToday, RISK_MM, RainTomorrow], file_type=parquet
           └──────────────────────────────────────────────────
         ");
@@ -612,13 +611,13 @@ mod tests {
           ┌───── Stage 2 ── Tasks: t0:[p0..p3] t1:[p0..p3] 
           │ AggregateExec: mode=FinalPartitioned, gby=[RainToday@0 as RainToday, WindGustDir@1 as WindGustDir], aggr=[]
           │   CoalesceBatchesExec: target_batch_size=8192
-          │     [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=2
+          │     [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=3
           └──────────────────────────────────────────────────
-            ┌───── Stage 1 ── Tasks: t0:[p0..p7] t1:[p0..p7] 
+            ┌───── Stage 1 ── Tasks: t0:[p0..p7] t1:[p0..p7] t2:[p0..p7] 
             │ RepartitionExec: partitioning=Hash([RainToday@0, WindGustDir@1], 8), input_partitions=4
-            │   RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=2
+            │   RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
             │     AggregateExec: mode=Partial, gby=[RainToday@0 as RainToday, WindGustDir@1 as WindGustDir], aggr=[]
-            │       PartitionIsolatorExec: t0:[p0,p1,__] t1:[__,__,p0] 
+            │       PartitionIsolatorExec: t0:[p0,__,__] t1:[__,p0,__] t2:[__,__,p0] 
             │         DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet], [/testdata/weather/result-000001.parquet], [/testdata/weather/result-000002.parquet]]}, projection=[RainToday, WindGustDir], file_type=parquet
             └──────────────────────────────────────────────────
         ");
@@ -649,7 +648,6 @@ mod tests {
             .with_default_features()
             .with_config(config)
             .with_distributed_execution(InMemoryChannelResolver::new())
-            .with_distributed_task_estimator(FixedDataSourceExecTaskEstimator(2))
             .build();
 
         let ctx = SessionContext::new_with_state(state);
diff --git a/src/distributed_planner/task_estimator.rs b/src/distributed_planner/task_estimator.rs
@@ -119,7 +119,7 @@ impl TaskEstimator for FileScanConfigTaskEstimator {
             Ok(urls) => urls.len(),
             Err(_) => 1,
         };
-        task_count = task_count.max(workers);
+        task_count = task_count.min(workers);
 
         // Based on the task count, attempt to scale up the partitions in the DataSourceExec by
         // repartitioning it. This will result in a DataSourceExec with potentially a lot of
diff --git a/src/metrics/task_metrics_collector.rs b/src/metrics/task_metrics_collector.rs
@@ -133,7 +133,6 @@ mod tests {
     use crate::test_utils::in_memory_channel_resolver::InMemoryChannelResolver;
     use crate::test_utils::plans::{count_plan_nodes, get_stages_and_stage_keys};
     use crate::test_utils::session_context::register_temp_parquet_table;
-    use crate::test_utils::test_task_estimator::FixedDataSourceExecTaskEstimator;
     use datafusion::execution::{SessionStateBuilder, context::SessionContext};
     use datafusion::prelude::SessionConfig;
     use datafusion::{
@@ -153,7 +152,7 @@ mod tests {
             .with_default_features()
             .with_config(config)
             .with_distributed_execution(InMemoryChannelResolver::new())
-            .with_distributed_task_estimator(FixedDataSourceExecTaskEstimator(2))
+            .with_distributed_task_estimator(2)
             .build();
 
         let ctx = SessionContext::from(state);
diff --git a/src/metrics/task_metrics_rewriter.rs b/src/metrics/task_metrics_rewriter.rs
@@ -219,7 +219,6 @@ mod tests {
 
     use crate::DistributedExt;
     use crate::metrics::task_metrics_rewriter::MetricsWrapperExec;
-    use crate::test_utils::test_task_estimator::FixedDataSourceExecTaskEstimator;
     use datafusion::physical_plan::empty::EmptyExec;
     use datafusion::physical_plan::metrics::MetricsSet;
     use datafusion::prelude::SessionConfig;
@@ -246,7 +245,7 @@ mod tests {
         if distributed {
             builder = builder
                 .with_distributed_execution(InMemoryChannelResolver::new())
-                .with_distributed_task_estimator(FixedDataSourceExecTaskEstimator(2))
+                .with_distributed_task_estimator(2)
         }
 
         let state = builder.build();
diff --git a/src/test_utils/mod.rs b/src/test_utils/mod.rs
@@ -6,5 +6,4 @@ pub mod mock_exec;
 pub mod parquet;
 pub mod plans;
 pub mod session_context;
-pub mod test_task_estimator;
 pub mod tpch;
diff --git a/src/test_utils/test_task_estimator.rs b/src/test_utils/test_task_estimator.rs
diff --git a/tests/distributed_aggregation.rs b/tests/distributed_aggregation.rs
diff --git a/tests/tpch_validation_test.rs b/tests/tpch_validation_test.rs
