migrate event handlers to one generic handler

benchmarks/cdk/bin/worker.rs

@@ -12,7 +12,8 @@ use datafusion::execution::runtime_env::RuntimeEnv;
 use datafusion::physical_plan::execute_stream;
 use datafusion::prelude::SessionContext;
 use datafusion_distributed::test_utils::work_unit_file_scan::{
-    WorkUnitFileScanCodec, WorkUnitFileScanConfig, WorkUnitFileScanTaskEstimator,
+    WorkUnitFileScanCodec, WorkUnitFileScanConfig, work_unit_file_scan_desired_task_count,
+    work_unit_file_scan_scale_up_leaf_node,
 };
 use datafusion_distributed::{
     ChannelResolver, DistributedExt, DistributedMetricsFormat, NetworkBoundaryExt,
@@ -121,7 +122,8 @@ async fn main() -> Result<(), Box<dyn Error>> {
         // Uncomment for enabling WorkUnitFileScans.
         // .with_physical_optimizer_rule(Arc::new(WorkUnitFileScanRule))
         .with_distributed_user_codec(WorkUnitFileScanCodec)
-        .with_distributed_task_estimator(WorkUnitFileScanTaskEstimator)
+        .with_distributed_event_handler(work_unit_file_scan_desired_task_count)
+        .with_distributed_event_handler(work_unit_file_scan_scale_up_leaf_node)
         .with_distributed_work_unit_feed(|dse: &DataSourceExec| {
             dse.data_source()
                 .downcast_ref::<WorkUnitFileScanConfig>()

benchmarks/src/run.rs

@@ -30,7 +30,7 @@ use datafusion::prelude::*;
 use datafusion_distributed::test_utils::localhost::LocalHostWorkerResolver;
 use datafusion_distributed::test_utils::work_unit_file_scan::{
     WorkUnitFileScanCodec, WorkUnitFileScanConfig, WorkUnitFileScanRule,
-    WorkUnitFileScanTaskEstimator,
+    work_unit_file_scan_desired_task_count, work_unit_file_scan_scale_up_leaf_node,
 };
 use datafusion_distributed::{
     DistributedExt, DistributedMetricsFormat, NetworkBoundaryExt, SessionStateBuilderExt, Worker,
@@ -216,7 +216,8 @@ impl RunOpt {
             )?
             .with_distributed_max_tasks_per_stage(self.max_tasks_per_stage)?
             .with_distributed_user_codec(WorkUnitFileScanCodec)
-            .with_distributed_task_estimator(WorkUnitFileScanTaskEstimator)
+            .with_distributed_event_handler(work_unit_file_scan_desired_task_count)
+            .with_distributed_event_handler(work_unit_file_scan_scale_up_leaf_node)
             .with_distributed_work_unit_feed(|dse: &DataSourceExec| {
                 dse.data_source()
                     .downcast_ref::<WorkUnitFileScanConfig>()

examples/custom_execution_plan.rs

@@ -39,8 +39,9 @@ use datafusion_distributed::test_utils::in_memory_channel_resolver::{
     InMemoryChannelResolver, InMemoryWorkerResolver,
 };
 use datafusion_distributed::{
-    DistributedExt, DistributedTaskContext, SessionStateBuilderExt, TaskEstimation, TaskEstimator,
-    WorkerQueryContext, display_plan_ascii,
+    DesiredTaskCountEvent, DesiredTaskCountEventResponse, DistributedExt, DistributedTaskContext,
+    ScaleUpLeafNodeEvent, ScaleUpLeafNodeEventResponse, SessionStateBuilderExt, WorkerQueryContext,
+    display_plan_ascii,
 };
 use datafusion_proto::physical_plan::PhysicalExtensionCodec;
 use datafusion_proto::protobuf;
@@ -297,47 +298,36 @@ impl ConfigExtension for NumbersConfig {
     const PREFIX: &'static str = "numbers";
 }
 
-/// Custom TaskEstimator that tells the planner how to distribute NumbersExec.
-#[derive(Debug)]
-struct NumbersTaskEstimator;
-
-impl TaskEstimator for NumbersTaskEstimator {
-    fn task_estimation(
-        &self,
-        plan: &Arc<dyn ExecutionPlan>,
-        cfg: &datafusion::config::ConfigOptions,
-    ) -> Option<TaskEstimation> {
-        let plan = plan.downcast_ref::<NumbersExec>()?;
-        let cfg: &NumbersConfig = cfg.extensions.get()?;
-        let task_count = (plan.ranges_per_task[0].end - plan.ranges_per_task[0].start) as f64
-            / cfg.numbers_per_task as f64;
-
-        Some(TaskEstimation::desired(task_count.ceil() as usize))
-    }
-
-    fn scale_up_leaf_node(
-        &self,
-        plan: &Arc<dyn ExecutionPlan>,
-        task_count: usize,
-        _cfg: &datafusion::config::ConfigOptions,
-    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
-        let Some(plan) = plan.downcast_ref::<NumbersExec>() else {
-            return Ok(None);
-        };
-        let range = &plan.ranges_per_task[0];
-        let chunk_size = ((range.end - range.start) as f64 / task_count as f64).ceil() as i64;
-
-        let ranges_per_task = (0..task_count).map(|i| {
-            let start = range.start + (i as i64 * chunk_size);
-            let end = (start + chunk_size).min(range.end);
-            start..end
-        });
+fn numbers_desired_task_count_handler(
+    ev: DesiredTaskCountEvent,
+) -> Option<DesiredTaskCountEventResponse> {
+    let cfg = ev.session_config;
+    let plan = ev.plan.downcast_ref::<NumbersExec>()?;
+    let cfg: &NumbersConfig = cfg.options().extensions.get()?;
+    let task_count = (plan.ranges_per_task[0].end - plan.ranges_per_task[0].start) as f64
+        / cfg.numbers_per_task as f64;
+
+    Some(DesiredTaskCountEventResponse::desired(
+        task_count.ceil() as usize
+    ))
+}
 
-        Ok(Some(Arc::new(NumbersExec::new(
-            ranges_per_task,
-            plan.schema(),
-        ))))
-    }
+fn numbers_scale_up_leaf_node_handler(
+    ev: ScaleUpLeafNodeEvent,
+) -> Option<Result<ScaleUpLeafNodeEventResponse>> {
+    let plan = ev.plan.downcast_ref::<NumbersExec>()?;
+    let range = &plan.ranges_per_task[0];
+    let chunk_size = ((range.end - range.start) as f64 / ev.task_count as f64).ceil() as i64;
+
+    let ranges_per_task = (0..ev.task_count).map(|i| {
+        let start = range.start + (i as i64 * chunk_size);
+        let end = (start + chunk_size).min(range.end);
+        start..end
+    });
+
+    Some(Ok(ScaleUpLeafNodeEventResponse::new(Arc::new(
+        NumbersExec::new(ranges_per_task, plan.schema()),
+    ))))
 }
 
 #[derive(StructOpt)]
@@ -381,7 +371,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         .with_distributed_channel_resolver(channel_resolver)
         .with_distributed_planner()
         .with_distributed_user_codec(NumbersExecCodec)
-        .with_distributed_task_estimator(NumbersTaskEstimator)
+        .with_distributed_event_handler(numbers_desired_task_count_handler)
+        .with_distributed_event_handler(numbers_scale_up_leaf_node_handler)
         .build();
 
     let ctx = SessionContext::from(state);

examples/custom_worker_url_routing.rs

@@ -6,9 +6,9 @@
 //! without changes to the table registration.
 //!
 //! Routing is a two-step pipeline:
-//! - [`CachedFileScanConfigTaskEstimator::scale_up_leaf_node`] assigns each file to a task slot by
+//! - `cached_file_scan_scale_up_leaf_node_handler` assigns each file to a task slot by
 //!   hashing its path (mod task_count), so the same file always lands in the same slot.
-//! - [`CachedFileScanConfigTaskEstimator::route_tasks`] maps slot `i` to `sorted_urls[i % n]`,
+//! - `cached_file_scan_route_tasks_handler` maps slot `i` to `sorted_urls[i % n]`,
 //!   so each slot always reaches the same worker URL.
 //!
 //! Together these guarantee that each worker consistently reads the same set of files and its
@@ -41,8 +41,9 @@ use datafusion_distributed::test_utils::localhost::{
     LocalHostWorkerResolver, spawn_worker_service,
 };
 use datafusion_distributed::{
-    DistributedExt, DistributedGetterExt, DistributedLeafExec, SessionStateBuilderExt,
-    TaskEstimation, TaskEstimator, TaskRoutingContext, WorkerQueryContext, display_plan_ascii,
+    DesiredTaskCountEvent, DesiredTaskCountEventResponse, DistributedExt, DistributedGetterExt,
+    DistributedLeafExec, RouteTasksEvent, RouteTasksEventResponse, ScaleUpLeafNodeEvent,
+    ScaleUpLeafNodeEventResponse, SessionStateBuilderExt, WorkerQueryContext, display_plan_ascii,
 };
 use datafusion_proto::physical_plan::PhysicalExtensionCodec;
 use datafusion_proto::protobuf;
@@ -55,7 +56,6 @@ use std::sync::Arc;
 use std::time::{Duration, Instant};
 use structopt::StructOpt;
 use tokio::net::TcpListener;
-use url::Url;
 
 /// Worker-level cache shared across all task invocations on the same worker, keyed by a stable
 /// hash of the file group being scanned.
@@ -157,68 +157,56 @@ fn hash_key(file_group: &FileGroup) -> usize {
     hasher.finish() as usize
 }
 
-/// Assigns each parquet file to a task slot (by hashing its path) and pins each slot to a worker.
-#[derive(Debug)]
-struct CachedFileScanConfigTaskEstimator;
+fn cached_file_scan_desired_task_count_handler(
+    ev: DesiredTaskCountEvent,
+) -> Option<DesiredTaskCountEventResponse> {
+    ev.plan.downcast_ref::<CacheExec>()?;
+    Some(DesiredTaskCountEventResponse::desired(usize::MAX))
+}
 
-impl TaskEstimator for CachedFileScanConfigTaskEstimator {
-    fn task_estimation(
-        &self,
-        plan: &Arc<dyn ExecutionPlan>,
-        _: &ConfigOptions,
-    ) -> Option<TaskEstimation> {
-        plan.downcast_ref::<CacheExec>()?;
-        Some(TaskEstimation::desired(usize::MAX))
+fn cached_file_scan_scale_up_leaf_node_handler(
+    ev: ScaleUpLeafNodeEvent,
+) -> Option<Result<ScaleUpLeafNodeEventResponse>> {
+    let cfg = ev.session_config;
+    let cache_exec = ev.plan.downcast_ref::<CacheExec>()?;
+    let dse = cache_exec.child.downcast_ref::<DataSourceExec>()?;
+    let fsc = dse.data_source().downcast_ref::<FileScanConfig>()?;
+
+    // Hash each file to a slot so that the same file always lands in the same variant
+    // regardless of the original file_groups layout.
+    let mut per_task_files: Vec<Vec<_>> = vec![vec![]; ev.task_count];
+    for file in fsc.file_groups.iter().flat_map(|g| g.files()) {
+        let idx = hash_key(&FileGroup::new(vec![file.clone()])) % ev.task_count;
+        per_task_files[idx].push(file.clone());
     }
 
-    fn scale_up_leaf_node(
-        &self,
-        plan: &Arc<dyn ExecutionPlan>,
-        task_count: usize,
-        cfg: &ConfigOptions,
-    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
-        let Some(cache_exec) = plan.downcast_ref::<CacheExec>() else {
-            return Ok(None);
-        };
-        let Some(dse) = cache_exec.child.downcast_ref::<DataSourceExec>() else {
-            return Ok(None);
-        };
-        let Some(fsc) = dse.data_source().downcast_ref::<FileScanConfig>() else {
-            return Ok(None);
-        };
-
-        // Hash each file to a slot so that the same file always lands in the same variant
-        // regardless of the original file_groups layout.
-        let mut per_task_files: Vec<Vec<_>> = vec![vec![]; task_count];
-        for file in fsc.file_groups.iter().flat_map(|g| g.files()) {
-            let idx = hash_key(&FileGroup::new(vec![file.clone()])) % task_count;
-            per_task_files[idx].push(file.clone());
-        }
+    let target_partitions = cfg.target_partitions();
+    let variants = (0..ev.task_count)
+        .map(|i| {
+            let files = std::mem::take(&mut per_task_files[i]);
+            // Spread files across up to `target_partitions` FileGroups so DataFusion can
+            // read them in parallel within the task.
+            let n_groups = files.len().clamp(1, target_partitions);
+            let mut groups: Vec<Vec<_>> = vec![vec![]; n_groups];
+            for (j, file) in files.into_iter().enumerate() {
+                groups[j % n_groups].push(file);
+            }
+            let mut new_fsc = fsc.clone();
+            new_fsc.file_groups = groups.into_iter().map(FileGroup::new).collect();
+            CacheExec::new(DataSourceExec::from_data_source(new_fsc)) as Arc<dyn ExecutionPlan>
+        })
+        .collect::<Vec<_>>();
 
-        let target_partitions = cfg.execution.target_partitions;
-        let variants = (0..task_count)
-            .map(|i| {
-                let files = std::mem::take(&mut per_task_files[i]);
-                // Spread files across up to `target_partitions` FileGroups so DataFusion can
-                // read them in parallel within the task.
-                let n_groups = files.len().clamp(1, target_partitions);
-                let mut groups: Vec<Vec<_>> = vec![vec![]; n_groups];
-                for (j, file) in files.into_iter().enumerate() {
-                    groups[j % n_groups].push(file);
-                }
-                let mut new_fsc = fsc.clone();
-                new_fsc.file_groups = groups.into_iter().map(FileGroup::new).collect();
-                CacheExec::new(DataSourceExec::from_data_source(new_fsc)) as Arc<dyn ExecutionPlan>
-            })
-            .collect::<Vec<_>>();
-
-        Ok(Some(Arc::new(DistributedLeafExec::try_new(
-            Arc::clone(plan),
-            variants,
-        )?)))
-    }
+    Some(
+        DistributedLeafExec::try_new(Arc::clone(ev.plan), variants)
+            .map(|plan| ScaleUpLeafNodeEventResponse::new(Arc::new(plan))),
+    )
+}
 
-    fn route_tasks(&self, ctx: &TaskRoutingContext<'_>) -> Result<Option<Vec<Url>>> {
+fn cached_file_scan_route_tasks_handler(
+    ctx: RouteTasksEvent,
+) -> Option<Result<RouteTasksEventResponse>> {
+    (|| -> Result<Option<RouteTasksEventResponse>> {
         let available_urls = ctx
             .task_ctx
             .session_config()
@@ -242,8 +230,9 @@ impl TaskEstimator for CachedFileScanConfigTaskEstimator {
             }
             Ok(TreeNodeRecursion::Continue)
         })?;
-        Ok(routed)
-    }
+        Ok(routed.map(RouteTasksEventResponse::new))
+    })()
+    .transpose()
 }
 
 /// Codec for [`CacheExec`]. The child (`DataSourceExec(FileScanConfig)`) is encoded by the
@@ -354,7 +343,9 @@ async fn main() -> Result<(), Box<dyn Error>> {
         .with_distributed_worker_resolver(LocalHostWorkerResolver::new(ports))
         .with_distributed_planner()
         .with_distributed_user_codec(CachedFileScanCodec)
-        .with_distributed_task_estimator(CachedFileScanConfigTaskEstimator)
+        .with_distributed_event_handler(cached_file_scan_desired_task_count_handler)
+        .with_distributed_event_handler(cached_file_scan_scale_up_leaf_node_handler)
+        .with_distributed_event_handler(cached_file_scan_route_tasks_handler)
         .build();
     state
         .config_mut()

examples/work_unit_feed.rs

@@ -30,8 +30,9 @@ use datafusion_distributed::test_utils::in_memory_channel_resolver::{
     InMemoryChannelResolver, InMemoryWorkerResolver,
 };
 use datafusion_distributed::{
-    DistributedExt, DistributedTaskContext, SessionStateBuilderExt, TaskEstimation, TaskEstimator,
-    WorkUnitFeed, WorkUnitFeedProto, WorkUnitFeedProvider, WorkerQueryContext, display_plan_ascii,
+    DesiredTaskCountEvent, DesiredTaskCountEventResponse, DistributedExt, DistributedTaskContext,
+    ScaleUpLeafNodeEvent, ScaleUpLeafNodeEventResponse, SessionStateBuilderExt, WorkUnitFeed,
+    WorkUnitFeedProto, WorkUnitFeedProvider, WorkerQueryContext, display_plan_ascii,
 };
 use datafusion_proto::physical_plan::PhysicalExtensionCodec;
 use datafusion_proto::protobuf::proto_error;
@@ -242,41 +243,30 @@ impl PhysicalExtensionCodec for RemoteScanExecCodec {
     }
 }
 
-/// Tells the planner how many tasks the leaf stage gets, and rebuilds the leaf so each task
-/// advertises its share of the partitions.
-#[derive(Debug)]
-struct RemoteScanTaskEstimator;
-
-impl TaskEstimator for RemoteScanTaskEstimator {
-    fn task_estimation(
-        &self,
-        plan: &Arc<dyn ExecutionPlan>,
-        _: &datafusion::config::ConfigOptions,
-    ) -> Option<TaskEstimation> {
-        let task_count = plan
-            .downcast_ref::<RemoteScanExec>()?
-            .feed
-            .inner()?
-            .task_count;
-        Some(TaskEstimation::desired(task_count))
-    }
+fn remote_scan_desired_task_count_handler(
+    ev: DesiredTaskCountEvent,
+) -> Option<DesiredTaskCountEventResponse> {
+    let task_count = ev
+        .plan
+        .downcast_ref::<RemoteScanExec>()?
+        .feed
+        .inner()?
+        .task_count;
+    Some(DesiredTaskCountEventResponse::desired(task_count))
+}
 
-    fn scale_up_leaf_node(
-        &self,
-        plan: &Arc<dyn ExecutionPlan>,
-        task_count: usize,
-        _: &datafusion::config::ConfigOptions,
-    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
-        let Some(exec) = plan.downcast_ref::<RemoteScanExec>() else {
-            return Ok(None);
-        };
-        let partitions_per_task = exec.feed.try_inner()?.per_partition_chunks.len() / task_count;
-        Ok(Some(Arc::new(RemoteScanExec::new(
+fn remote_scan_scale_up_leaf_node_handler(
+    ev: ScaleUpLeafNodeEvent,
+) -> Option<Result<ScaleUpLeafNodeEventResponse>> {
+    let exec = ev.plan.downcast_ref::<RemoteScanExec>()?;
+    Some(exec.feed.try_inner().map(|feed| {
+        let partitions_per_task = feed.per_partition_chunks.len() / ev.task_count;
+        ScaleUpLeafNodeEventResponse::new(Arc::new(RemoteScanExec::new(
             exec.feed.clone(),
             partitions_per_task,
             exec.projection.clone(),
-        ))))
-    }
+        )))
+    }))
 }
 
 /// `scan(task_count, 'chunks_p0', 'chunks_p1', ...)` — `task_count` tasks, with one
@@ -392,7 +382,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         .with_distributed_channel_resolver(channel_resolver)
         .with_distributed_planner()
         .with_distributed_user_codec(RemoteScanExecCodec)
-        .with_distributed_task_estimator(RemoteScanTaskEstimator)
+        .with_distributed_event_handler(remote_scan_desired_task_count_handler)
+        .with_distributed_event_handler(remote_scan_scale_up_leaf_node_handler)
         // For every `RemoteScanExec`, hand the planner the feed it must drive from the coordinator.
         .with_distributed_work_unit_feed(|exec: &RemoteScanExec| Some(&exec.feed))
         .build();

src/coordinator/distributed.rs

@@ -20,7 +20,7 @@ use std::sync::{Arc, Mutex};
 /// [ExecutionPlan] that executes the inner plan in distributed mode.
 /// Before executing it, two modifications are lazily performed on the plan:
 /// 1. Assigns worker URLs to all the stages. Unless explicitly set in
-///    [crate::TaskEstimator::route_tasks], a random set of URLs are sampled from the
+///    [`crate::EventHandler<crate::RouteTasksHandler>`], a random set of URLs are sampled from the
 ///    channel resolver and assigned to each task in each stage.
 /// 2. Encodes all the plans in protobuf format so that network boundary nodes can send them
 ///    over the wire.