From 010971694c5666d80871d94a7177e5fa1de35c0a Mon Sep 17 00:00:00 2001 From: Gabriel Date: Mon, 27 Jul 2026 09:38:24 +0200 Subject: [PATCH 1/2] Make desired_task_count.rs event handler async --- .../inject_network_boundaries.rs | 12 ++++--- src/events/common.rs | 8 +++++ src/events/defaults/file_scan_config.rs | 2 ++ src/events/desired_task_count.rs | 34 ++++++++++++++----- 4 files changed, 42 insertions(+), 14 deletions(-) diff --git a/src/distributed_planner/inject_network_boundaries.rs b/src/distributed_planner/inject_network_boundaries.rs index 6e433ec6..4c3b8e72 100644 --- a/src/distributed_planner/inject_network_boundaries.rs +++ b/src/distributed_planner/inject_network_boundaries.rs @@ -237,7 +237,7 @@ async fn _inject_network_boundaries( plan: &plan, session_config: nb_ctx.cfg, }; - return if let Some(estimate) = DesiredTaskCountHandlers::handle(ev) { + return if let Some(estimate) = DesiredTaskCountHandlers::handle(ev).await { Ok(nb_ctx.plan_with_task_count(plan, estimate.task_count.limit(nb_ctx.max_tasks()?))) } else { // We could not determine how many tasks this leaf node should run on, so @@ -261,10 +261,12 @@ async fn _inject_network_boundaries( plan: &plan, session_config: nb_ctx.cfg, }; - let mut task_count = DesiredTaskCountHandlers::handle(ev).map_or(Desired(1), |v| v.task_count); - if plan.is::() { - // Isolating unions have the chance to decide how many tasks they should run on. If there - // is a union with a bunch of children, the user might want to increase parallelism and the + let mut task_count = DesiredTaskCountHandlers::handle(ev) + .await + .map_or(Desired(1), |v| v.task_count); + if nb_ctx.d_cfg.children_isolator_unions && plan.is::() { + // 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 // task count for the stage running that. let mut count = 0; for processed_child in processed_children.iter() { diff --git a/src/events/common.rs b/src/events/common.rs index 7f715059..fc357e9e 100644 --- a/src/events/common.rs +++ b/src/events/common.rs @@ -25,6 +25,14 @@ impl Clone for EventHandlerChain { } impl EventHandlerChain { + pub(super) fn iter(&self) -> impl Iterator { + self.custom + .iter() + .rev() + .chain(&self.builtin) + .map(AsRef::as_ref) + } + pub(super) fn find_map(&self, mut f: impl FnMut(&H) -> Option) -> Option { // Give priority to custom handlers registered by users. if let Some(res) = self.custom.iter().find_map(|handler| f(handler.as_ref())) { diff --git a/src/events/defaults/file_scan_config.rs b/src/events/defaults/file_scan_config.rs index 32aef96c..f11451b8 100644 --- a/src/events/defaults/file_scan_config.rs +++ b/src/events/defaults/file_scan_config.rs @@ -115,6 +115,7 @@ mod tests { plan: &plan, session_config: &cfg, }) + .await .expect("a handler should respond"); assert_eq!(response.task_count.as_usize(), 10); Ok(()) @@ -131,6 +132,7 @@ mod tests { plan: &plan, session_config: &cfg, }) + .await .expect("a handler should respond"); assert_eq!(response.task_count.as_usize(), 30); Ok(()) diff --git a/src/events/desired_task_count.rs b/src/events/desired_task_count.rs index 85981b79..cd6cf28e 100644 --- a/src/events/desired_task_count.rs +++ b/src/events/desired_task_count.rs @@ -1,5 +1,6 @@ use super::TaskCountAnnotation::{Desired, Maximum}; use super::common::EventHandlerChain; +use async_trait::async_trait; use datafusion::execution::config::SessionConfig; use datafusion::physical_plan::ExecutionPlan; use std::sync::Arc; @@ -73,10 +74,14 @@ impl DesiredTaskCountEventResponse { } } +#[async_trait] pub trait DesiredTaskCountHandler: Send + Sync + 'static { /// Function applied to each node that returns a [DesiredTaskCountEventResponse] hinting how /// many tasks should be used in the [Stage] containing that node. /// + /// Handlers are asynchronous and may await metadata or external services. Handler functions + /// return a [`DesiredTaskCountFuture`] so their futures can borrow from the event. + /// /// All the [TaskEstimator] registered in the session will be applied to the node /// until one returns an estimation. /// @@ -86,7 +91,7 @@ pub trait DesiredTaskCountHandler: Send + Sync + 'static { /// 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 handle(&self, ev: DesiredTaskCountEvent) -> Option; + async fn handle(&self, ev: DesiredTaskCountEvent<'_>) -> Option; } impl From for usize { @@ -120,18 +125,20 @@ impl TaskCountAnnotation { } } +#[async_trait] impl DesiredTaskCountHandler for F where F: Send + Sync + 'static, F: for<'a> Fn(DesiredTaskCountEvent<'a>) -> Option, { - fn handle(&self, ev: DesiredTaskCountEvent) -> Option { + async fn handle(&self, ev: DesiredTaskCountEvent<'_>) -> Option { self(ev) } } +#[async_trait] impl DesiredTaskCountHandler for usize { - fn handle(&self, ev: DesiredTaskCountEvent) -> Option { + async fn handle(&self, ev: DesiredTaskCountEvent<'_>) -> Option { ev.plan .children() .is_empty() @@ -139,18 +146,27 @@ impl DesiredTaskCountHandler for usize { } } +#[async_trait] impl DesiredTaskCountHandler for Arc { - fn handle(&self, ev: DesiredTaskCountEvent) -> Option { - self.as_ref().handle(ev) + async fn handle(&self, ev: DesiredTaskCountEvent<'_>) -> Option { + self.as_ref().handle(ev).await } } pub(crate) type DesiredTaskCountHandlers = EventHandlerChain; impl DesiredTaskCountHandlers { - pub(crate) fn handle(ev: DesiredTaskCountEvent) -> Option { - ev.session_config - .get_extension::()? - .find_map(|handler| handler.handle(ev)) + pub(crate) async fn handle( + ev: DesiredTaskCountEvent<'_>, + ) -> Option { + let handlers = ev + .session_config + .get_extension::()?; + for handler in handlers.iter() { + if let Some(response) = handler.handle(ev).await { + return Some(response); + } + } + None } } From a5e0b27f5efe2edada4d2475d5251902a758fb74 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Mon, 27 Jul 2026 10:49:04 +0200 Subject: [PATCH 2/2] Run desired task-count handlers in registration order --- src/distributed_planner/inject_network_boundaries.rs | 6 +++--- src/events/common.rs | 6 +----- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/src/distributed_planner/inject_network_boundaries.rs b/src/distributed_planner/inject_network_boundaries.rs index 4c3b8e72..75e6813d 100644 --- a/src/distributed_planner/inject_network_boundaries.rs +++ b/src/distributed_planner/inject_network_boundaries.rs @@ -264,9 +264,9 @@ async fn _inject_network_boundaries( let mut task_count = DesiredTaskCountHandlers::handle(ev) .await .map_or(Desired(1), |v| v.task_count); - if nb_ctx.d_cfg.children_isolator_unions && plan.is::() { - // 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 + if plan.is::() { + // Isolating unions have the chance to decide how many tasks they should run on. If there + // is a union with a bunch of children, the user might want to increase parallelism and the // task count for the stage running that. let mut count = 0; for processed_child in processed_children.iter() { diff --git a/src/events/common.rs b/src/events/common.rs index fc357e9e..fd5c02e4 100644 --- a/src/events/common.rs +++ b/src/events/common.rs @@ -26,11 +26,7 @@ impl Clone for EventHandlerChain { impl EventHandlerChain { pub(super) fn iter(&self) -> impl Iterator { - self.custom - .iter() - .rev() - .chain(&self.builtin) - .map(AsRef::as_ref) + self.custom.iter().chain(&self.builtin).map(AsRef::as_ref) } pub(super) fn find_map(&self, mut f: impl FnMut(&H) -> Option) -> Option {