Apply TaskEstimator on all nodes (#271)

gabotechs · web-flow · commit 84243d3c92e5 · 2025-12-29T14:48:28.000+01:00
* Split channel resolver in two

* Simplify WorkerResolverExtension and ChannelResolverExtension

* Add default builder to ArrowFlightEndpoint

* Add some docs

* Listen to clippy

* Split get_flight_client_for_url in two

* Fix conflicts

* Remove unnecessary channel resolver

* Improve WorkerResolver docs

* Use one ChannelResolver per runtime

* Improve error reporting on client connection failure

* Add a from_session_builder method for constructing an InMemoryChannelResolver

* Add ChannelResolver and WorkerResolver default implementations for Arcs

* Make TPC-DS tests use DataFusion test dataset

* Remove non-working in-memory option from benchmarks

* Remove unnecessary utils folder

* Refactor benchmark folder

* Rename to prepare_tpch.rs

* Adapt benchmarks for TPC-DS

* Update benchmarks README.md

* Fix conflicts

* Use default session state builder

* Add ChildrenIsolatorUnionExec

* Add proto serde for ChildrenIsolatorUnionExec

* Wire up ChildrenIsolatorUnionExec to planner

* Add integration tests for distributed UNIONs

* Skip query 72 in TPC-DS benchmarks

* Allow setting children isolator unions

* Allow passing multiple queries in benchmarks

* Extend TaskEstimator API to also be applicable to intermediate nodes
diff --git a/src/distributed_planner/plan_annotator.rs b/src/distributed_planner/plan_annotator.rs
@@ -138,6 +138,7 @@ fn _annotate_plan(
 ) -> Result<AnnotatedPlan, DataFusionError> {
     use TaskCountAnnotation::*;
     let d_cfg = DistributedConfig::from_config_options(cfg)?;
+    let estimator = &d_cfg.__private_task_estimator;
     let n_workers = d_cfg.__private_worker_resolver.0.get_urls()?.len().max(1);
 
     let annotated_children = plan
@@ -150,8 +151,7 @@ fn _annotate_plan(
         // This is a leaf node, maybe a DataSourceExec, or maybe something else custom from the
         // user. We need to estimate how many tasks are needed for this leaf node, and we'll take
         // this decision into account when deciding how many tasks will be actually used.
-        let estimator = &d_cfg.__private_task_estimator;
-        if let Some(estimate) = estimator.tasks_for_leaf_node(&plan, cfg) {
+        if let Some(estimate) = estimator.task_estimation(&plan, cfg) {
             return Ok(AnnotatedPlan {
                 plan,
                 children: Vec::new(),
@@ -170,7 +170,9 @@ fn _annotate_plan(
         }
     }
 
-    let mut task_count = Desired(1);
+    let mut task_count = estimator
+        .task_estimation(&plan, cfg)
+        .map_or(Desired(1), |v| v.task_count);
     if d_cfg.children_isolator_unions && plan.as_any().is::<UnionExec>() {
         // Unions have the chance to decide how many tasks they should run on. If there's a union
         // with a bunch of children, the user might want to increase parallelism and increase the
@@ -341,11 +343,11 @@ mod tests {
     use super::*;
     use crate::test_utils::in_memory_channel_resolver::InMemoryWorkerResolver;
     use crate::test_utils::parquet::register_parquet_tables;
-    use crate::{DistributedExt, assert_snapshot};
+    use crate::{DistributedExt, TaskEstimation, assert_snapshot};
     use datafusion::execution::SessionStateBuilder;
+    use datafusion::physical_plan::filter::FilterExec;
     use datafusion::prelude::{SessionConfig, SessionContext};
     use itertools::Itertools;
-
     /* schema for the "weather" table
 
      MinTemp [type=DOUBLE] [repetitiontype=OPTIONAL]
@@ -584,16 +586,116 @@ mod tests {
         ")
     }
 
+    #[tokio::test]
+    async fn test_intermediate_task_estimator() {
+        let query = r#"
+        SELECT DISTINCT "RainToday" FROM weather
+        "#;
+        let annotated = sql_to_annotated_with_estimator(query, |_: &RepartitionExec| {
+            Some(TaskEstimation::maximum(1))
+        })
+        .await;
+        assert_snapshot!(annotated, @r"
+        AggregateExec: task_count=Desired(1)
+          CoalesceBatchesExec: task_count=Desired(1), required_network_boundary=Shuffle
+            RepartitionExec: task_count=Maximum(1)
+              RepartitionExec: task_count=Maximum(1)
+                AggregateExec: task_count=Maximum(1)
+                  DataSourceExec: task_count=Maximum(1)
+        ")
+    }
+
+    #[tokio::test]
+    async fn test_union_all_limited_by_intermediate_estimator() {
+        let query = r#"
+        SELECT "MinTemp" FROM weather WHERE "RainToday" = 'yes'
+        UNION ALL
+        SELECT "MaxTemp" FROM weather WHERE "RainToday" = 'no'
+        "#;
+        let annotated = sql_to_annotated_with_estimator(query, |_: &FilterExec| {
+            Some(TaskEstimation::maximum(1))
+        })
+        .await;
+        assert_snapshot!(annotated, @r"
+        ChildrenIsolatorUnionExec: task_count=Desired(2)
+          CoalesceBatchesExec: task_count=Maximum(1)
+            FilterExec: task_count=Maximum(1)
+              RepartitionExec: task_count=Maximum(1)
+                DataSourceExec: task_count=Maximum(1)
+          ProjectionExec: task_count=Maximum(1)
+            CoalesceBatchesExec: task_count=Maximum(1)
+              FilterExec: task_count=Maximum(1)
+                RepartitionExec: task_count=Maximum(1)
+                  DataSourceExec: task_count=Maximum(1)
+        ")
+    }
+
+    #[allow(clippy::type_complexity)]
+    struct CallbackEstimator {
+        f: Arc<dyn Fn(&(dyn ExecutionPlan)) -> Option<TaskEstimation> + Send + Sync>,
+    }
+
+    impl CallbackEstimator {
+        fn new<T: ExecutionPlan + 'static>(
+            f: impl Fn(&T) -> Option<TaskEstimation> + Send + Sync + 'static,
+        ) -> Self {
+            let f = Arc::new(move |plan: &dyn ExecutionPlan| -> Option<TaskEstimation> {
+                if let Some(plan) = plan.as_any().downcast_ref::<T>() {
+                    f(plan)
+                } else {
+                    None
+                }
+            });
+            Self { f }
+        }
+    }
+
+    impl TaskEstimator for CallbackEstimator {
+        fn task_estimation(
+            &self,
+            plan: &Arc<dyn ExecutionPlan>,
+            _: &ConfigOptions,
+        ) -> Option<TaskEstimation> {
+            (self.f)(plan.as_ref())
+        }
+
+        fn scale_up_leaf_node(
+            &self,
+            _: &Arc<dyn ExecutionPlan>,
+            _: usize,
+            _: &ConfigOptions,
+        ) -> Option<Arc<dyn ExecutionPlan>> {
+            None
+        }
+    }
+
     async fn sql_to_annotated(query: &str) -> String {
+        sql_to_annotated_with_options(query, move |b| b).await
+    }
+
+    async fn sql_to_annotated_with_estimator<T: ExecutionPlan + Send + Sync + 'static>(
+        query: &str,
+        estimator: impl Fn(&T) -> Option<TaskEstimation> + Send + Sync + 'static,
+    ) -> String {
+        sql_to_annotated_with_options(query, move |b| {
+            b.with_distributed_task_estimator(CallbackEstimator::new(estimator))
+        })
+        .await
+    }
+
+    async fn sql_to_annotated_with_options(
+        query: &str,
+        f: impl FnOnce(SessionStateBuilder) -> SessionStateBuilder,
+    ) -> String {
         let config = SessionConfig::new()
             .with_target_partitions(4)
             .with_information_schema(true);
 
-        let state = SessionStateBuilder::new()
+        let state = f(SessionStateBuilder::new()
             .with_default_features()
             .with_config(config)
-            .with_distributed_worker_resolver(InMemoryWorkerResolver::new(4))
-            .build();
+            .with_distributed_worker_resolver(InMemoryWorkerResolver::new(4)))
+        .build();
 
         let ctx = SessionContext::new_with_state(state);
         let mut queries = query.split(";").collect_vec();
diff --git a/src/distributed_planner/task_estimator.rs b/src/distributed_planner/task_estimator.rs
@@ -5,6 +5,7 @@ use datafusion::config::ConfigOptions;
 use datafusion::datasource::physical_plan::FileScanConfig;
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion::prelude::SessionConfig;
+use delegate::delegate;
 use std::collections::HashSet;
 use std::fmt::Debug;
 use std::sync::Arc;
@@ -58,6 +59,34 @@ pub struct TaskEstimation {
     pub task_count: TaskCountAnnotation,
 }
 
+impl TaskEstimation {
+    /// Tells the distributed planner that the evaluated stage can have **at maximum** the provided
+    /// number of tasks, setting a hard upper limit.
+    ///
+    /// Returning `TaskEstimation::maximum(1)` tells the distributed planner that the evaluated
+    /// stage cannot be distributed.
+    ///
+    /// Even if a `TaskEstimation::maximum(N)` is provided, any other node in the same stage
+    /// providing a value of `TaskEstimation::maximum(M)` where `M` < `N` will have preference.
+    pub fn maximum(value: usize) -> Self {
+        TaskEstimation {
+            task_count: TaskCountAnnotation::Maximum(value),
+        }
+    }
+
+    /// Tells the distributed planner that the evaluated can **optimally** have the provided
+    /// number of tasks, setting a soft task count hint that can be overridden by others.
+    ///
+    /// The provided `TaskEstimation::desired(N)` can be overridden by:
+    /// - Other nodes providing a `TaskEstimation::desired(M)` where `M` > `N`.
+    /// - Any other node providing a `TaskEstimation::maximum(M)` where `M` can be anything.
+    pub fn desired(value: usize) -> Self {
+        TaskEstimation {
+            task_count: TaskCountAnnotation::Desired(value),
+        }
+    }
+}
+
 /// Given a leaf node, provides an estimation about how many tasks should be used in the
 /// stage containing it, and if the leaf node should be replaced by some other.
 ///
@@ -66,14 +95,19 @@ pub struct TaskEstimation {
 /// count calculated based on whether lower stages are reducing the cardinality of the data
 /// or increasing it.
 pub trait TaskEstimator {
-    /// Function applied to leaf nodes that returns a [TaskEstimation] hinting how many
-    /// tasks should be used in the [Stage] containing that leaf node.
+    /// Function applied to each node that returns a [TaskEstimation] hinting how many
+    /// tasks should be used in the [Stage] containing that node.
+    ///
+    /// All the [TaskEstimator] registered in the session will be applied to the node
+    /// until one returns an estimation.
+    ///
     ///
-    /// All the [TaskEstimator] registered in the session will be applied to the leaf node
-    /// until one returns an estimation. If no estimation is return from any of the
-    /// [TaskEstimator]s, then `Maximum(1)` is returned, hinting the distributed planner to not
-    /// distribute the stage containing that node.
-    fn tasks_for_leaf_node(
+    /// If no estimation is returned from any of the registered [TaskEstimator]s, then:
+    /// - If the node is a leaf node,`Maximum(1)` is assumed, hinting the distributed planner
+    ///   that the leaf node cannot be distributed across tasks.
+    /// - If the node is a normal node in the plan, then the maximum task count from its children
+    ///   is inherited.
+    fn task_estimation(
         &self,
         plan: &Arc<dyn ExecutionPlan>,
         cfg: &ConfigOptions,
@@ -91,14 +125,18 @@ pub trait TaskEstimator {
 }
 
 impl TaskEstimator for usize {
-    fn tasks_for_leaf_node(
+    fn task_estimation(
         &self,
-        _: &Arc<dyn ExecutionPlan>,
+        inputs: &Arc<dyn ExecutionPlan>,
         _: &ConfigOptions,
     ) -> Option<TaskEstimation> {
-        Some(TaskEstimation {
-            task_count: TaskCountAnnotation::Desired(*self),
-        })
+        if inputs.children().is_empty() {
+            Some(TaskEstimation {
+                task_count: TaskCountAnnotation::Desired(*self),
+            })
+        } else {
+            None
+        }
     }
 
     fn scale_up_leaf_node(
@@ -112,40 +150,20 @@ impl TaskEstimator for usize {
 }
 
 impl TaskEstimator for Arc<dyn TaskEstimator> {
-    fn tasks_for_leaf_node(
-        &self,
-        plan: &Arc<dyn ExecutionPlan>,
-        cfg: &ConfigOptions,
-    ) -> Option<TaskEstimation> {
-        self.as_ref().tasks_for_leaf_node(plan, cfg)
-    }
-
-    fn scale_up_leaf_node(
-        &self,
-        plan: &Arc<dyn ExecutionPlan>,
-        task_count: usize,
-        cfg: &ConfigOptions,
-    ) -> Option<Arc<dyn ExecutionPlan>> {
-        self.as_ref().scale_up_leaf_node(plan, task_count, cfg)
+    delegate! {
+        to self.as_ref() {
+            fn task_estimation(&self, plan: &Arc<dyn ExecutionPlan>, cfg: &ConfigOptions) -> Option<TaskEstimation>;
+            fn scale_up_leaf_node(&self, plan: &Arc<dyn ExecutionPlan>, task_count: usize, cfg: &ConfigOptions) -> Option<Arc<dyn ExecutionPlan>>;
+        }
     }
 }
 
 impl TaskEstimator for Arc<dyn TaskEstimator + Send + Sync> {
-    fn tasks_for_leaf_node(
-        &self,
-        plan: &Arc<dyn ExecutionPlan>,
-        cfg: &ConfigOptions,
-    ) -> Option<TaskEstimation> {
-        self.as_ref().tasks_for_leaf_node(plan, cfg)
-    }
-
-    fn scale_up_leaf_node(
-        &self,
-        plan: &Arc<dyn ExecutionPlan>,
-        task_count: usize,
-        cfg: &ConfigOptions,
-    ) -> Option<Arc<dyn ExecutionPlan>> {
-        self.as_ref().scale_up_leaf_node(plan, task_count, cfg)
+    delegate! {
+        to self.as_ref() {
+            fn task_estimation(&self, plan: &Arc<dyn ExecutionPlan>, cfg: &ConfigOptions) -> Option<TaskEstimation>;
+            fn scale_up_leaf_node(&self, plan: &Arc<dyn ExecutionPlan>, task_count: usize, cfg: &ConfigOptions) -> Option<Arc<dyn ExecutionPlan>>;
+        }
     }
 }
 
@@ -177,15 +195,16 @@ pub(crate) fn set_distributed_task_estimator(
 struct FileScanConfigTaskEstimator;
 
 impl TaskEstimator for FileScanConfigTaskEstimator {
-    fn tasks_for_leaf_node(
+    fn task_estimation(
         &self,
         plan: &Arc<dyn ExecutionPlan>,
         cfg: &ConfigOptions,
     ) -> Option<TaskEstimation> {
-        let d_cfg = cfg.extensions.get::<DistributedConfig>()?;
         let dse: &DataSourceExec = plan.as_any().downcast_ref()?;
         let file_scan: &FileScanConfig = dse.data_source().as_any().downcast_ref()?;
 
+        let d_cfg = cfg.extensions.get::<DistributedConfig>()?;
+
         // Count how many distinct files we have in the FileScanConfig. Each file in each
         // file group is a PartitionedFile rather than a full file, so it's possible that
         // many entries refer to different chunks of the same physical file. By keeping a
@@ -244,21 +263,21 @@ pub(crate) struct CombinedTaskEstimator {
 }
 
 impl TaskEstimator for CombinedTaskEstimator {
-    fn tasks_for_leaf_node(
+    fn task_estimation(
         &self,
         plan: &Arc<dyn ExecutionPlan>,
         cfg: &ConfigOptions,
     ) -> Option<TaskEstimation> {
         for estimator in &self.user_provided {
-            if let Some(result) = estimator.tasks_for_leaf_node(plan, cfg) {
+            if let Some(result) = estimator.task_estimation(plan, cfg) {
                 return Some(result);
             }
         }
         // We want to execute the default estimators last so that the user-provided ones have
         // a chance of providing an estimation.
         // If none of the user-provided returned an estimation, the default ones are used.
         for default_estimator in [&FileScanConfigTaskEstimator as &dyn TaskEstimator] {
-            if let Some(result) = default_estimator.tasks_for_leaf_node(plan, cfg) {
+            if let Some(result) = default_estimator.task_estimation(plan, cfg) {
                 return Some(result);
             }
         }
@@ -348,7 +367,7 @@ mod tests {
                 ..Default::default()
             };
             cfg.extensions.insert(f(d_cfg));
-            self.tasks_for_leaf_node(&node, &cfg)
+            self.task_estimation(&node, &cfg)
                 .unwrap()
                 .task_count
                 .as_usize()
@@ -370,7 +389,7 @@ mod tests {
     }
 
     impl<F: Fn(&Arc<dyn ExecutionPlan>, &ConfigOptions) -> Option<TaskEstimation>> TaskEstimator for F {
-        fn tasks_for_leaf_node(
+        fn task_estimation(
             &self,
             plan: &Arc<dyn ExecutionPlan>,
             cfg: &ConfigOptions,
