Add dynamic stage built events

src/coordinator/prepare_dynamic_plan.rs

@@ -5,11 +5,13 @@ use crate::distributed_planner::{
     InjectNetworkBoundaryContext, NetworkBoundaryBuilderResult, ProducerHead, calculate_cost,
     inject_network_boundaries,
 };
+use crate::events::DynamicStageBuiltHandlers;
 use crate::events::TaskCountAnnotation::{Desired, Maximum};
 use crate::execution_plans::SamplerExec;
 use crate::stage::{LocalStage, RemoteStage};
 use crate::{
-    BytesCounterMetric, LoadInfo, MaxGaugeMetric, NetworkBoundaryExt, NetworkCoalesceExec, Stage,
+    BytesCounterMetric, DynamicStageBuiltEvent, LoadInfo, MaxGaugeMetric, NetworkBoundaryExt,
+    NetworkCoalesceExec, Stage,
 };
 use dashmap::DashMap;
 use datafusion::common::stats::Precision;
@@ -69,6 +71,15 @@ pub(super) async fn prepare_dynamic_plan(
                 .task_count(&input_stage.plan)?
                 .merge(Desired(compute_based_task_count));
 
+            let ev = DynamicStageBuiltEvent {
+                session_config: nb_ctx.cfg,
+                cost,
+                plan: &input_stage.plan,
+            };
+            if let Some(response) = DynamicStageBuiltHandlers::handle(ev).transpose()? {
+                input_stage.plan = response.plan
+            };
+
             // Propagate the final task_count inferred based on runtime statistics and compute cost.
             // Here is where leaf nodes are scaled up by ScaleUpLeafNodeHandler, and the
             // plan is finally left ready for distribution.

src/distributed_ext.rs

@@ -3,17 +3,17 @@ use crate::config_extension_ext::{
     set_distributed_option_extension, set_distributed_option_extension_from_headers,
 };
 use crate::events::{
-    DesiredTaskCountHandler, DesiredTaskCountHandlers, RouteTasksHandler, RouteTasksHandlers,
-    ScaleUpLeafNodeHandler, ScaleUpLeafNodeHandlers, WorkerPlanRewriteHandler,
-    WorkerPlanRewriteHandlers,
+    DesiredTaskCountHandler, DesiredTaskCountHandlers, DynamicStageBuiltHandlers,
+    RouteTasksHandler, RouteTasksHandlers, ScaleUpLeafNodeHandler, ScaleUpLeafNodeHandlers,
+    WorkerPlanRewriteHandler, WorkerPlanRewriteHandlers,
 };
 use crate::passthrough_headers::set_passthrough_headers;
 use crate::protocol::set_distributed_channel_resolver;
 use crate::work_unit_feed::set_distributed_work_unit_feed;
 use crate::worker_resolver::set_distributed_worker_resolver;
 use crate::{
-    ChannelResolver, DistributedConfig, LocalWorkerContext, WorkUnitFeed, WorkUnitFeedProvider,
-    WorkerResolver, get_distributed_worker_resolver,
+    ChannelResolver, DistributedConfig, DynamicStageBuiltHandler, LocalWorkerContext, WorkUnitFeed,
+    WorkUnitFeedProvider, WorkerResolver, get_distributed_worker_resolver,
 };
 use datafusion::common::DataFusionError;
 use datafusion::config::ConfigExtension;
@@ -737,6 +737,39 @@ pub trait DistributedExt: Sized {
         &mut self,
         handler: T,
     );
+
+    /// Register an event handler that fires every time a new stage is built during dynamic
+    /// planning.
+    ///
+    /// In this handler, users can inspect the plan that is about to be sent to a remote worker,
+    /// optimize it, or short circuit execution.
+    ///
+    /// ```rust
+    /// # use datafusion::common::{exec_err, Result};
+    /// # use datafusion::execution::SessionStateBuilder;
+    /// # use datafusion_distributed::{DistributedExt, DynamicStageBuiltEvent, DynamicStageBuiltEventResponse};
+    ///
+    /// fn handle_dynamic_stage_built(event: DynamicStageBuiltEvent) -> Option<Result<DynamicStageBuiltEventResponse>> {
+    ///     if *event.cost.cpu.get_value().unwrap_or(&0) > 1024 * 1024 * 1024 {
+    ///         return Some(exec_err!("Plan is too expensive to execute"))
+    ///     }
+    ///     None
+    /// }
+    ///
+    /// SessionStateBuilder::new()
+    ///     .with_distributed_dynamic_stage_built_handler(handle_dynamic_stage_built);
+    /// ```
+    fn with_distributed_dynamic_stage_built_handler<T: DynamicStageBuiltHandler>(
+        self,
+        handler: T,
+    ) -> Self;
+
+    /// Same as [DistributedExt::with_distributed_dynamic_stage_built_handler] but with an
+    /// in-place mutation.
+    fn set_distributed_dynamic_stage_built_handler<T: DynamicStageBuiltHandler>(
+        &mut self,
+        handler: T,
+    );
 }
 
 /// Trait to have a unified interface for getting structs & properties from SessionConfig that are used in distributed context.
@@ -917,6 +950,10 @@ impl DistributedExt for SessionConfig {
         WorkerPlanRewriteHandlers::push_custom(self, Arc::new(h));
     }
 
+    fn set_distributed_dynamic_stage_built_handler<T: DynamicStageBuiltHandler>(&mut self, h: T) {
+        DynamicStageBuiltHandlers::push_custom(self, Arc::new(h));
+    }
+
     delegate! {
         to self {
             #[call(set_distributed_option_extension)]
@@ -1024,6 +1061,10 @@ impl DistributedExt for SessionConfig {
             #[call(set_distributed_worker_plan_rewrite_handler)]
             #[expr($;self)]
             fn with_distributed_worker_plan_rewrite_handler<H: WorkerPlanRewriteHandler>(mut self, h: H) -> Self;
+
+            #[call(set_distributed_dynamic_stage_built_handler)]
+            #[expr($;self)]
+            fn with_distributed_dynamic_stage_built_handler<H: DynamicStageBuiltHandler>(mut self, h: H) -> Self;
         }
     }
 }
@@ -1172,6 +1213,11 @@ impl DistributedExt for SessionStateBuilder {
             #[call(set_distributed_worker_plan_rewrite_handler)]
             #[expr($;self)]
             fn with_distributed_worker_plan_rewrite_handler<H: WorkerPlanRewriteHandler>(mut self, h: H) -> Self;
+
+            fn set_distributed_dynamic_stage_built_handler<H: DynamicStageBuiltHandler>(&mut self, h: H);
+            #[call(set_distributed_dynamic_stage_built_handler)]
+            #[expr($;self)]
+            fn with_distributed_dynamic_stage_built_handler<H: DynamicStageBuiltHandler>(mut self, h: H) -> Self;
         }
     }
 }
@@ -1322,6 +1368,11 @@ impl DistributedExt for SessionState {
             #[call(set_distributed_worker_plan_rewrite_handler)]
             #[expr($;self)]
             fn with_distributed_worker_plan_rewrite_handler<H: WorkerPlanRewriteHandler>(mut self, h: H) -> Self;
+
+            fn set_distributed_dynamic_stage_built_handler<H: DynamicStageBuiltHandler>(&mut self, h: H);
+            #[call(set_distributed_dynamic_stage_built_handler)]
+            #[expr($;self)]
+            fn with_distributed_dynamic_stage_built_handler<H: DynamicStageBuiltHandler>(mut self, h: H) -> Self;
         }
     }
 }
@@ -1465,6 +1516,11 @@ impl DistributedExt for SessionContext {
             #[call(set_distributed_worker_plan_rewrite_handler)]
             #[expr($;self)]
             fn with_distributed_worker_plan_rewrite_handler<H: WorkerPlanRewriteHandler>(self, h: H) -> Self;
+
+            fn set_distributed_dynamic_stage_built_handler<H: DynamicStageBuiltHandler>(&mut self, h: H);
+            #[call(set_distributed_dynamic_stage_built_handler)]
+            #[expr($;self)]
+            fn with_distributed_dynamic_stage_built_handler<H: DynamicStageBuiltHandler>(self, h: H) -> Self;
         }
     }
 }

src/distributed_planner/inject_network_boundaries.rs

@@ -158,8 +158,8 @@ pub(crate) async fn inject_network_boundaries(
 #[derive(Clone)]
 pub(crate) struct InjectNetworkBoundaryContext<'a> {
     pub(crate) d_cfg: &'a DistributedConfig,
+    pub(crate) cfg: &'a SessionConfig,
 
-    cfg: &'a SessionConfig,
     worker_resolver: Arc<WorkerResolverExtension>,
     nb_builder: &'a (dyn NetworkBoundaryBuilder + Send + Sync),
     task_counts: &'a Mutex<HashMap<usize, TaskCountAnnotation>>,

src/distributed_planner/statistics/cost.rs

@@ -1,3 +1,4 @@
+use crate::Cost;
 use crate::distributed_planner::statistics::complexity_cpu::complexity_cpu;
 use crate::distributed_planner::statistics::complexity_memory::complexity_memory;
 use crate::distributed_planner::statistics::complexity_network::complexity_network;
@@ -8,13 +9,6 @@ use datafusion::physical_plan::{ExecutionPlan, Statistics};
 use std::ops::AddAssign;
 use std::sync::Arc;
 
-#[derive(Default, Debug)]
-pub(crate) struct Cost {
-    pub(crate) cpu: Precision<usize>,
-    pub(crate) memory: Precision<usize>,
-    pub(crate) network: Precision<usize>,
-}
-
 impl AddAssign for Cost {
     fn add_assign(&mut self, rhs: Self) {
         self.cpu = sum_precision(self.cpu, rhs.cpu);

src/distributed_planner/statistics/mod.rs

@@ -6,5 +6,4 @@ mod cost;
 mod default_bytes_for_datatype;
 mod plan_statistics;
 
-#[allow(unused)] // will be used in a follow-up PR.
 pub(crate) use cost::calculate_cost;

src/events/dynamic_stage_built.rs

@@ -0,0 +1,67 @@
+use crate::events::common::EventHandlerChain;
+use datafusion::common::{Result, stats::Precision};
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::prelude::SessionConfig;
+use std::sync::Arc;
+
+/// Estimated cost associated to the fragment of the plan. Calculated based on the amount of data
+/// that will flow through each node and the static cpu, memory and network complexity of the
+/// different nodes based on the algorithms that will run inside them.
+#[derive(Default, Debug, Clone, Copy)]
+pub struct Cost {
+    /// Estimated cost of CPU measured in bytes.
+    pub cpu: Precision<usize>,
+    /// Estimated amount of memory consumed measured in bytes.
+    pub memory: Precision<usize>,
+    /// Estimated amount of data that will need to be transferred over the wire in bytes.
+    pub network: Precision<usize>,
+}
+
+/// Event definition fired for each individual stage during dynamic planning.
+#[derive(Copy, Clone)]
+pub struct DynamicStageBuiltEvent<'a> {
+    /// Cost associated to the plan in the field below and all its children recursively until
+    /// network boundaries.
+    pub cost: Cost,
+    /// Plan that is about to be sent to a remote worker for runtime sampling.
+    pub plan: &'a Arc<dyn ExecutionPlan>,
+    /// SessionConfig in scope at the moment of building the stage.
+    pub session_config: &'a SessionConfig,
+}
+
+pub struct DynamicStageBuiltEventResponse {
+    /// A potentially modified plan (e.g., optimizations applied).
+    pub plan: Arc<dyn ExecutionPlan>,
+}
+
+impl DynamicStageBuiltEventResponse {
+    pub fn new(plan: Arc<dyn ExecutionPlan>) -> Self {
+        Self { plan }
+    }
+}
+
+pub trait DynamicStageBuiltHandler: Send + Sync + 'static {
+    fn handle(&self, ev: DynamicStageBuiltEvent) -> Option<Result<DynamicStageBuiltEventResponse>>;
+}
+
+impl<F> DynamicStageBuiltHandler for F
+where
+    F: Send + Sync + 'static,
+    F: for<'a> Fn(DynamicStageBuiltEvent<'a>) -> Option<Result<DynamicStageBuiltEventResponse>>,
+{
+    fn handle(&self, ev: DynamicStageBuiltEvent) -> Option<Result<DynamicStageBuiltEventResponse>> {
+        self(ev)
+    }
+}
+
+pub(crate) type DynamicStageBuiltHandlers = EventHandlerChain<dyn DynamicStageBuiltHandler>;
+
+impl DynamicStageBuiltHandlers {
+    pub(crate) fn handle(
+        ev: DynamicStageBuiltEvent,
+    ) -> Option<Result<DynamicStageBuiltEventResponse>> {
+        ev.session_config
+            .get_extension::<DynamicStageBuiltHandlers>()?
+            .find_map(|handler| handler.handle(ev))
+    }
+}

src/events/mod.rs

@@ -1,6 +1,7 @@
 mod common;
 mod defaults;
 mod desired_task_count;
+mod dynamic_stage_built;
 mod route_tasks;
 mod scale_up_leaf_node;
 mod worker_plan_rewrite;
@@ -14,6 +15,10 @@ pub use desired_task_count::{
     DesiredTaskCountEvent, DesiredTaskCountEventResponse, DesiredTaskCountHandler,
     TaskCountAnnotation,
 };
+pub(crate) use dynamic_stage_built::DynamicStageBuiltHandlers;
+pub use dynamic_stage_built::{
+    Cost, DynamicStageBuiltEvent, DynamicStageBuiltEventResponse, DynamicStageBuiltHandler,
+};
 pub(crate) use route_tasks::RouteTasksHandlers;
 pub use route_tasks::{RouteTasksEvent, RouteTasksEventResponse, RouteTasksHandler};
 pub(crate) use scale_up_leaf_node::ScaleUpLeafNodeHandlers;

src/lib.rs

@@ -23,10 +23,11 @@ pub use distributed_planner::{
     DistributedConfig, NetworkBoundary, NetworkBoundaryExt, SessionStateBuilderExt,
 };
 pub use events::{
-    DesiredTaskCountEvent, DesiredTaskCountEventResponse, DesiredTaskCountHandler, RouteTasksEvent,
-    RouteTasksEventResponse, RouteTasksHandler, ScaleUpLeafNodeEvent, ScaleUpLeafNodeEventResponse,
-    ScaleUpLeafNodeHandler, TaskCountAnnotation, WorkerPlanRewriteEvent,
-    WorkerPlanRewriteEventResponse, WorkerPlanRewriteHandler,
+    Cost, DesiredTaskCountEvent, DesiredTaskCountEventResponse, DesiredTaskCountHandler,
+    DynamicStageBuiltEvent, DynamicStageBuiltEventResponse, DynamicStageBuiltHandler,
+    RouteTasksEvent, RouteTasksEventResponse, RouteTasksHandler, ScaleUpLeafNodeEvent,
+    ScaleUpLeafNodeEventResponse, ScaleUpLeafNodeHandler, TaskCountAnnotation,
+    WorkerPlanRewriteEvent, WorkerPlanRewriteEventResponse, WorkerPlanRewriteHandler,
 };
 pub use execution_plans::{
     BroadcastExec, DistributedLeafExec, NetworkBroadcastExec, NetworkCoalesceExec,