Rework task assignation mechanism (#216)

* Introduce with_distributed_execution method and remove DistributedExt from anything that's not a SessionStateBuilder

* Rework how tasks are assigned to stages

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

* Update all tests to use the new task estimation mechanism

* Fix when worker urls is zero or one and add tests

* Move back to `with_distributed_channel_resolver` and let users inject either `DistributedPhysicalOptimizerRule` or their own rules

* Add files_per_task and cardinality_task_sf to examples

* Improve docs

* add impl TaskEstimator for Arc<&dyn TaskEstimator>

* Add task estimator tests

* Explain why we need to count distinct files

* Return an error if task_count == 0

* Factor out apply_scale_factor

* Add task estimator doc comments

* Add more docs to apply_network_boundaries

* Extend files_per_task docs

* Rollback DistributedExt de-implementation from SessionConfig, SessionState and SessionContext

Author: gabotechs

benchmarks/src/tpch/run.rs

@@ -101,28 +101,6 @@ pub struct RunOpt {
     #[structopt(short = "t", long = "sorted")]
     sorted: bool,
 
-    /// Upon shuffling data, this defines how many tasks are employed into performing the shuffling.
-    /// ```text
-    ///  ( task 1 )  ( task 2 ) ( task 3 )
-    ///      ▲           ▲          ▲
-    ///      └────┬──────┴─────┬────┘
-    ///       ( task 1 )  ( task 2 )       N tasks
-    /// ```
-    /// This parameter defines N
-    #[structopt(long)]
-    shuffle_tasks: Option<usize>,
-
-    /// Upon merging multiple tasks into one, this defines how many tasks are merged.
-    /// ```text
-    ///              ( task 1 )
-    ///                  ▲
-    ///      ┌───────────┴──────────┐
-    ///  ( task 1 )  ( task 2 ) ( task 3 )  N tasks
-    /// ```
-    /// This parameter defines N
-    #[structopt(long)]
-    coalesce_tasks: Option<usize>,
-
     /// Spawns a worker in the specified port.
     #[structopt(long)]
     spawn: Option<u16>,
@@ -134,6 +112,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]
@@ -142,37 +128,32 @@ impl DistributedSessionBuilder for RunOpt {
         &self,
         ctx: DistributedSessionBuilderContext,
     ) -> Result<SessionState, DataFusionError> {
-        let mut builder = SessionStateBuilder::new().with_default_features();
-
+        let rt_builder = self.common.runtime_env_builder()?;
         let config = self
             .common
             .config()?
-            .with_collect_statistics(!self.disable_statistics)
+            .with_target_partitions(self.partitions())
+            .with_collect_statistics(!self.disable_statistics);
+        let mut builder = SessionStateBuilder::new()
+            .with_runtime_env(rt_builder.build_arc()?)
+            .with_default_features()
+            .with_config(config)
             .with_distributed_user_codec(InMemoryCacheExecCodec)
             .with_distributed_channel_resolver(LocalHostChannelResolver::new(self.workers.clone()))
+            .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule))
             .with_distributed_option_extension_from_headers::<WarmingUpMarker>(&ctx.headers)?
-            .with_target_partitions(self.partitions());
-
-        let rt_builder = self.common.runtime_env_builder()?;
+            .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_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule))
-                .with_distributed_network_coalesce_tasks(
-                    self.coalesce_tasks.unwrap_or(self.workers.len()),
-                )
-                .with_distributed_network_shuffle_tasks(
-                    self.shuffle_tasks.unwrap_or(self.workers.len()),
-                );
-        }
 
-        Ok(builder
-            .with_config(config)
-            .with_runtime_env(rt_builder.build_arc()?)
-            .build())
+        Ok(builder.build())
     }
 }
 
@@ -196,7 +177,12 @@ impl RunOpt {
     }
 
     async fn run_local(mut self) -> Result<()> {
-        let state = self.build_session_state(Default::default()).await?;
+        let mut state = self.build_session_state(Default::default()).await?;
+        if self.mem_table {
+            state = SessionStateBuilder::from(state)
+                .with_distributed_option_extension(WarmingUpMarker::warming_up())?
+                .build();
+        }
         let ctx = SessionContext::new_with_state(state);
         self.register_tables(&ctx).await?;
 
@@ -218,9 +204,6 @@ impl RunOpt {
             for query_id in query_range.clone() {
                 // put the WarmingUpMarker in the context, otherwise, queries will fail as the
                 // InMemoryCacheExec node will think they should already be warmed up.
-                let ctx = ctx
-                    .clone()
-                    .with_distributed_option_extension(WarmingUpMarker::warming_up())?;
                 for query in get_query_sql(query_id)? {
                     self.execute_query(&ctx, &query).await?;
                 }

examples/in_memory_cluster.rs

@@ -2,6 +2,7 @@ use arrow::util::pretty::pretty_format_batches;
 use arrow_flight::flight_service_client::FlightServiceClient;
 use async_trait::async_trait;
 use datafusion::common::DataFusionError;
+use datafusion::common::utils::get_available_parallelism;
 use datafusion::execution::SessionStateBuilder;
 use datafusion::physical_plan::displayable;
 use datafusion::prelude::{ParquetReadOptions, SessionContext};
@@ -28,11 +29,11 @@ struct Args {
     #[structopt(long)]
     explain: bool,
 
-    #[structopt(long, default_value = "3")]
-    network_shuffle_tasks: usize,
+    #[structopt(long)]
+    files_per_task: Option<usize>,
 
-    #[structopt(long, default_value = "3")]
-    network_coalesce_tasks: usize,
+    #[structopt(long)]
+    cardinality_task_sf: Option<f64>,
 }
 
 #[tokio::main]
@@ -43,8 +44,12 @@ async fn main() -> Result<(), Box<dyn Error>> {
         .with_default_features()
         .with_distributed_channel_resolver(InMemoryChannelResolver::new())
         .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule))
-        .with_distributed_network_coalesce_tasks(args.network_shuffle_tasks)
-        .with_distributed_network_shuffle_tasks(args.network_coalesce_tasks)
+        .with_distributed_files_per_task(
+            args.files_per_task.unwrap_or(get_available_parallelism()),
+        )?
+        .with_distributed_cardinality_effect_task_scale_factor(
+            args.cardinality_task_sf.unwrap_or(1.),
+        )?
         .build();
 
     let ctx = SessionContext::from(state);

examples/localhost_run.rs

@@ -3,6 +3,7 @@ use arrow_flight::flight_service_client::FlightServiceClient;
 use async_trait::async_trait;
 use dashmap::{DashMap, Entry};
 use datafusion::common::DataFusionError;
+use datafusion::common::utils::get_available_parallelism;
 use datafusion::execution::SessionStateBuilder;
 use datafusion::physical_plan::displayable;
 use datafusion::prelude::{ParquetReadOptions, SessionContext};
@@ -29,11 +30,11 @@ struct Args {
     #[structopt(long)]
     explain: bool,
 
-    #[structopt(long, default_value = "3")]
-    network_shuffle_tasks: usize,
+    #[structopt(long)]
+    files_per_task: Option<usize>,
 
-    #[structopt(long, default_value = "3")]
-    network_coalesce_tasks: usize,
+    #[structopt(long)]
+    cardinality_task_sf: Option<f64>,
 }
 
 #[tokio::main]
@@ -49,8 +50,12 @@ async fn main() -> Result<(), Box<dyn Error>> {
         .with_default_features()
         .with_distributed_channel_resolver(localhost_resolver)
         .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule))
-        .with_distributed_network_coalesce_tasks(args.network_coalesce_tasks)
-        .with_distributed_network_shuffle_tasks(args.network_shuffle_tasks)
+        .with_distributed_files_per_task(
+            args.files_per_task.unwrap_or(get_available_parallelism()),
+        )?
+        .with_distributed_cardinality_effect_task_scale_factor(
+            args.cardinality_task_sf.unwrap_or(1.),
+        )?
         .build();
 
     let ctx = SessionContext::from(state);

src/channel_resolver_ext.rs

@@ -1,6 +1,7 @@
+use crate::DistributedConfig;
 use arrow_flight::flight_service_client::FlightServiceClient;
 use async_trait::async_trait;
-use datafusion::common::exec_datafusion_err;
+use datafusion::common::exec_err;
 use datafusion::error::DataFusionError;
 use datafusion::prelude::SessionConfig;
 use std::sync::Arc;
@@ -11,21 +12,30 @@ pub(crate) fn set_distributed_channel_resolver(
     cfg: &mut SessionConfig,
     channel_resolver: impl ChannelResolver + Send + Sync + 'static,
 ) {
-    cfg.set_extension(Arc::new(ChannelResolverExtension(Arc::new(
-        channel_resolver,
-    ))));
+    let opts = cfg.options_mut();
+    let channel_resolver_ext = ChannelResolverExtension(Arc::new(channel_resolver));
+    if let Some(distributed_cfg) = opts.extensions.get_mut::<DistributedConfig>() {
+        distributed_cfg.__private_channel_resolver = channel_resolver_ext;
+    } else {
+        opts.extensions.insert(DistributedConfig {
+            __private_channel_resolver: channel_resolver_ext,
+            ..Default::default()
+        });
+    }
 }
 
 pub(crate) fn get_distributed_channel_resolver(
     cfg: &SessionConfig,
 ) -> Result<Arc<dyn ChannelResolver + Send + Sync>, DataFusionError> {
-    cfg.get_extension::<ChannelResolverExtension>()
-        .map(|cm| cm.0.clone())
-        .ok_or_else(|| exec_datafusion_err!("ChannelResolver not present in the session config"))
+    let opts = cfg.options();
+    let Some(distributed_cfg) = opts.extensions.get::<DistributedConfig>() else {
+        return exec_err!("ChannelResolver not present in the session config");
+    };
+    Ok(Arc::clone(&distributed_cfg.__private_channel_resolver.0))
 }
 
 #[derive(Clone)]
-struct ChannelResolverExtension(Arc<dyn ChannelResolver + Send + Sync>);
+pub(crate) struct ChannelResolverExtension(pub(crate) Arc<dyn ChannelResolver + Send + Sync>);
 
 pub type BoxCloneSyncChannel = tower::util::BoxCloneSyncService<
     http::Request<Body>,