Use files_per_task and cardinality_effect_factor for choosing the right tasks for the different stages

Author: gabotechs

benchmarks/src/tpch/run.rs

@@ -45,7 +45,6 @@ use datafusion::prelude::*;
 use datafusion_distributed::test_utils::localhost::{
     LocalHostChannelResolver, spawn_flight_service,
 };
-use datafusion_distributed::test_utils::test_task_estimator::FixedDataSourceExecTaskEstimator;
 use datafusion_distributed::{
     DistributedExt, DistributedSessionBuilder, DistributedSessionBuilderContext, NetworkBoundaryExt,
 };
@@ -101,10 +100,6 @@ pub struct RunOpt {
     #[structopt(short = "t", long = "sorted")]
     sorted: bool,
 
-    /// Sets how many tasks are used per stage.
-    #[structopt(long)]
-    stage_tasks: Option<usize>,
-
     /// Spawns a worker in the specified port.
     #[structopt(long)]
     spawn: Option<u16>,
@@ -116,6 +111,14 @@ pub struct RunOpt {
     /// Number of physical threads per worker.
     #[structopt(long)]
     threads: Option<usize>,
+
+    /// Number of files per each distributed task.
+    #[structopt(long)]
+    files_per_task: Option<usize>,
+
+    /// Task count scale factor for when nodes in stages change the cardinality of the data
+    #[structopt(long)]
+    cardinality_task_sf: Option<f64>,
 }
 
 #[async_trait]
@@ -136,16 +139,17 @@ impl DistributedSessionBuilder for RunOpt {
             .with_config(config)
             .with_distributed_user_codec(InMemoryCacheExecCodec)
             .with_distributed_execution(LocalHostChannelResolver::new(self.workers.clone()))
-            .with_distributed_option_extension_from_headers::<WarmingUpMarker>(&ctx.headers)?;
+            .with_distributed_option_extension_from_headers::<WarmingUpMarker>(&ctx.headers)?
+            .with_distributed_files_per_task(
+                self.files_per_task.unwrap_or(get_available_parallelism()),
+            )?
+            .with_distributed_cardinality_effect_task_scale_factor(
+                self.cardinality_task_sf.unwrap_or(1.0),
+            )?;
 
         if self.mem_table {
             builder = builder.with_physical_optimizer_rule(Arc::new(InMemoryDataSourceRule));
         }
-        if !self.workers.is_empty() {
-            builder = builder.with_distributed_task_estimator(FixedDataSourceExecTaskEstimator(
-                self.stage_tasks.unwrap_or(self.workers.len()),
-            ));
-        }
 
         Ok(builder.build())
     }

src/distributed_ext.rs

@@ -4,7 +4,7 @@ use crate::config_extension_ext::{
 };
 use crate::distributed_planner::set_distributed_task_estimator;
 use crate::protobuf::{set_distributed_user_codec, set_distributed_user_codec_arc};
-use crate::{ChannelResolver, DistributedPhysicalOptimizerRule, TaskEstimator};
+use crate::{ChannelResolver, DistributedConfig, DistributedPhysicalOptimizerRule, TaskEstimator};
 use datafusion::common::DataFusionError;
 use datafusion::config::ConfigExtension;
 use datafusion::execution::SessionStateBuilder;
@@ -250,6 +250,45 @@ pub trait DistributedExt: Sized {
         &mut self,
         estimator: T,
     );
+
+    /// Sets the maximum number of files each task in a stage with a FileScanConfig node will
+    /// handle. Reducing this number will increment the amount of tasks. By default, this
+    /// is close to the number of cores in the machine.
+    fn with_distributed_files_per_task(
+        self,
+        files_per_task: usize,
+    ) -> Result<Self, DataFusionError>;
+
+    /// Same as [DistributedExt::with_distributed_files_per_task] but with an in-place mutation.
+    fn set_distributed_files_per_task(
+        &mut self,
+        files_per_task: usize,
+    ) -> Result<(), DataFusionError>;
+
+    /// The number of tasks in each stage is calculated in a bottom-to-top fashion.
+    ///
+    /// Bottom stages containing leaf nodes will provide an estimation of the amount of tasks
+    /// for those stages, but upper stages might see a reduction (or increment) in the amount
+    /// of tasks based on the cardinality effect bottom stages have in the data.
+    ///
+    /// For example: If there are two stages, and the leaf stage is estimated to use 10 tasks,
+    ///  the upper stage might use less (e.g. 5) if it sees that the leaf stage is returning
+    ///  less data because of filters or aggregations.
+    ///
+    /// This function sets the scale factor for when encountering these nodes that change the
+    /// cardinality of the data. For example, if a stage with 10 tasks contains an AggregateExec
+    /// node, and the scale factor is 2.0, the following stage will use  10 / 2.0 = 5 tasks.
+    fn with_distributed_cardinality_effect_task_scale_factor(
+        self,
+        factor: f64,
+    ) -> Result<Self, DataFusionError>;
+
+    /// Same as [DistributedExt::with_distributed_cardinality_effect_task_scale_factor] but with
+    /// an in-place mutation.
+    fn set_distributed_cardinality_effect_task_scale_factor(
+        &mut self,
+        factor: f64,
+    ) -> Result<(), DataFusionError>;
 }
 
 impl DistributedExt for SessionStateBuilder {
@@ -295,6 +334,26 @@ impl DistributedExt for SessionStateBuilder {
         set_distributed_task_estimator(self.config().get_or_insert_default(), estimator)
     }
 
+    fn set_distributed_files_per_task(
+        &mut self,
+        files_per_task: usize,
+    ) -> Result<(), DataFusionError> {
+        let cfg = self.config().get_or_insert_default();
+        let d_cfg = DistributedConfig::from_config_options_mut(cfg.options_mut())?;
+        d_cfg.files_per_task = files_per_task;
+        Ok(())
+    }
+
+    fn set_distributed_cardinality_effect_task_scale_factor(
+        &mut self,
+        factor: f64,
+    ) -> Result<(), DataFusionError> {
+        let cfg = self.config().get_or_insert_default();
+        let d_cfg = DistributedConfig::from_config_options_mut(cfg.options_mut())?;
+        d_cfg.cardinality_task_count_factor = factor;
+        Ok(())
+    }
+
     delegate! {
         to self {
             #[call(set_distributed_option_extension)]
@@ -320,6 +379,14 @@ impl DistributedExt for SessionStateBuilder {
             #[call(set_distributed_task_estimator)]
             #[expr($;self)]
             fn with_distributed_task_estimator<T: TaskEstimator + Send + Sync + 'static>(mut self, estimator: T) -> Self;
+
+            #[call(set_distributed_files_per_task)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_files_per_task(mut self, files_per_task: usize) -> Result<Self, DataFusionError>;
+
+            #[call(set_distributed_cardinality_effect_task_scale_factor)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_cardinality_effect_task_scale_factor(mut self, factor: f64) -> Result<Self, DataFusionError>;
         }
     }
 }

src/distributed_planner/distributed_config.rs

@@ -3,6 +3,7 @@ use crate::distributed_planner::task_estimator::CombinedTaskEstimator;
 use crate::{BoxCloneSyncChannel, ChannelResolver, TaskEstimator};
 use arrow_flight::flight_service_client::FlightServiceClient;
 use async_trait::async_trait;
+use datafusion::common::utils::get_available_parallelism;
 use datafusion::common::{DataFusionError, extensions_options, not_impl_err, plan_err};
 use datafusion::config::{ConfigExtension, ConfigField, ConfigOptions, Visit};
 use std::fmt::{Debug, Formatter};
@@ -12,6 +13,9 @@ use url::Url;
 extensions_options! {
     /// Configuration for the distributed planner.
     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()
         /// 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.
@@ -48,6 +52,14 @@ impl DistributedConfig {
         };
         Ok(distributed_cfg)
     }
+
+    /// Gets the [DistributedConfig] from the [ConfigOptions]'s extensions.
+    pub fn from_config_options_mut(cfg: &mut ConfigOptions) -> Result<&mut Self, DataFusionError> {
+        let Some(distributed_cfg) = cfg.extensions.get_mut::<DistributedConfig>() else {
+            return plan_err!("DistributedConfig is not in ConfigOptions.extensions");
+        };
+        Ok(distributed_cfg)
+    }
 }
 
 impl ConfigExtension for DistributedConfig {

src/distributed_planner/task_estimator.rs

@@ -1,9 +1,10 @@
-use crate::{DistributedConfig, PartitionIsolatorExec};
+use crate::{ChannelResolver, DistributedConfig, PartitionIsolatorExec};
 use datafusion::catalog::memory::DataSourceExec;
 use datafusion::config::ConfigOptions;
 use datafusion::datasource::physical_plan::FileScanConfig;
-use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties};
 use datafusion::prelude::SessionConfig;
+use std::collections::HashSet;
 use std::fmt::Debug;
 use std::sync::Arc;
 
@@ -96,14 +97,43 @@ impl TaskEstimator for FileScanConfigTaskEstimator {
     fn estimate_tasks(
         &self,
         plan: &Arc<dyn ExecutionPlan>,
-        _: &ConfigOptions,
+        cfg: &ConfigOptions,
     ) -> Option<TaskEstimation> {
+        let d_cfg = cfg.extensions.get::<DistributedConfig>()?;
         let dse: &DataSourceExec = plan.as_any().downcast_ref()?;
-        let _config: &FileScanConfig = dse.data_source().as_any().downcast_ref()?;
+        let file_scan: &FileScanConfig = dse.data_source().as_any().downcast_ref()?;
+
+        // Count how many distinct files we have in the FileScanConfig.
+        let mut distinct_files = HashSet::new();
+        for file_group in &file_scan.file_groups {
+            for file in file_group.iter() {
+                distinct_files.insert(file.object_meta.location.clone());
+            }
+        }
+        let distinct_files = distinct_files.len();
+
+        // Based on the user-provided files_per_task configuration, do the math to calculate
+        // how many tasks should be used, without surpassing the number of available workers.
+        let mut task_count = distinct_files.div_ceil(d_cfg.files_per_task);
+        let workers = match d_cfg.__private_channel_resolver.0.get_urls() {
+            Ok(urls) => urls.len(),
+            Err(_) => 1,
+        };
+        task_count = task_count.max(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
+        // partitions, but as we are going wrap it with PartitionIsolatorExec that's fine.
+        let scaled_partitions = task_count * plan.output_partitioning().partition_count();
+        let mut plan = Arc::clone(plan);
+        if let Ok(Some(repartitioned)) = plan.repartitioned(scaled_partitions, cfg) {
+            plan = repartitioned;
+        }
+        plan = Arc::new(PartitionIsolatorExec::new(plan));
 
         Some(TaskEstimation {
-            task_count: 1, // TODO
-            new_plan: Some(Arc::new(PartitionIsolatorExec::new(Arc::clone(plan)))),
+            task_count,
+            new_plan: Some(plan),
         })
     }
 }