Define worker plan rewrites in terms of events

docs/source/advanced/03-plan-hooks.md

@@ -1,36 +1,49 @@
-# Plan hooks
+# Worker plan rewrite handlers
 
-`Worker::add_on_plan_hook` registers callbacks that run after the worker session
-has been built and the physical plan has been decoded, but before the task plan
-is registered for execution. It's a hook for **worker-local** rewrites of the
-fragment a worker is about to run.
+`DistributedExt::with_distributed_worker_plan_rewrite_handler` registers handlers that run after
+the worker session has been built and the physical plan has been decoded, but before the task plan
+is registered for execution. It is intended for **worker-local** rewrites of the fragment a worker
+is about to run.
 
-Each hook receives the decoded `ExecutionPlan` and the per-query `SessionConfig`,
-and returns the (possibly rewritten) plan:
+Register handlers in the `SessionStateBuilder` used by each worker's `WorkerSessionBuilder`.
+Registering one on the coordinator's session has no effect: handlers are not sent to workers with
+stage plans.
+
+Each handler receives a `WorkerPlanRewriteEvent` and returns the possibly rewritten plan:
 
 ```rust
-worker.add_on_plan_hook(|plan, session_config| {
-    let rule = your_physical_optimizer_rule();
-    rule.optimize(plan, session_config.options())
-});
+# use datafusion::common::Result;
+# use datafusion::execution::SessionState;
+# use datafusion_distributed::{DistributedExt, Worker, WorkerPlanRewriteEvent, WorkerPlanRewriteEventResponse, WorkerQueryContext};
+async fn build_worker_session(ctx: WorkerQueryContext) -> Result<SessionState> {
+    Ok(ctx
+        .builder
+        .with_distributed_worker_plan_rewrite_handler(|event: WorkerPlanRewriteEvent<'_>| {
+            Ok(WorkerPlanRewriteEventResponse::new(event.plan))
+        })
+        .build())
+}
+
+let _worker = Worker::from_session_builder(build_worker_session);
 ```
 
-Hooks run in registration order — each hook sees the plan produced by the
+Handlers run in registration order—each handler sees the plan produced by the
 previous one.
 
-## What hooks may and may not do
+## What handlers may and may not do
+
+Treat handlers as trusted, worker-local rewrites. Instrumentation and semantics-preserving
+physical optimizer rules that retain the stage topology are appropriate uses. A handler must:
 
-Treat hooks as trusted, worker-local rewrites. Transparent instrumentation
-wrappers and semantics-preserving physical optimizer rules are appropriate uses.
+- Keep the same plan topology; do not add or remove nodes or edges.
+- Produce the same rows, including multiplicity, from every output partition.
+- Preserve the head node's output schema, partitioning, ordering, boundedness, and emission type.
 
-The returned plan **must preserve the contract the coordinating context planned
-for**: row semantics, output schema, partitioning, and ordering requirements. Do
-**not** use a hook to add or remove rows or columns, repartition the stage, or
-otherwise re-plan distributed execution. If a hook returns an error, the
-distributed query fails when the coordinating context tries to execute that task.
+Intermediate nodes may change their output schema. If a handler returns an error, the distributed
+query fails when the coordinating context tries to execute that task.
 
 ```{note}
-A plan hook acts on a single worker's copy of a stage, after distribution has
+A worker plan rewrite handler acts on a single worker's copy of a stage, after distribution has
 already been decided. It is **not** a way to change how a query is split across
 the cluster — for that, see
 [Building Custom Distributed Plans](05-custom-distributed-plans.md).

docs/source/user-guide/03-worker.md

@@ -109,7 +109,7 @@ object-store registrations (pass them via `Worker::from_session_builder` /
 
 A few more worker capabilities have their own pages:
 
-- [Plan hooks](../advanced/03-plan-hooks.md) — run worker-local rewrites on each decoded plan
+- [Worker plan rewrite handlers](../advanced/03-plan-hooks.md) — run worker-local rewrites on each decoded plan
   before it executes.
 - [Worker versioning](../advanced/07-worker-versioning.md) — tag workers with a version and
   route queries only to compatible workers during rolling deployments.

src/distributed_ext.rs

@@ -4,7 +4,8 @@ use crate::config_extension_ext::{
 };
 use crate::events::{
     DesiredTaskCountHandler, DesiredTaskCountHandlers, RouteTasksHandler, RouteTasksHandlers,
-    ScaleUpLeafNodeHandler, ScaleUpLeafNodeHandlers,
+    ScaleUpLeafNodeHandler, ScaleUpLeafNodeHandlers, WorkerPlanRewriteHandler,
+    WorkerPlanRewriteHandlers,
 };
 use crate::passthrough_headers::set_passthrough_headers;
 use crate::protocol::set_distributed_channel_resolver;
@@ -690,6 +691,52 @@ pub trait DistributedExt: Sized {
     /// Same as [DistributedExt::with_distributed_route_tasks_handler] but with an in-place
     /// mutation.
     fn set_distributed_route_tasks_handler<T: RouteTasksHandler>(&mut self, handler: T);
+
+    /// Registers a handler that rewrites a decoded worker stage plan before it is executed.
+    ///
+    /// Handlers registered with this method need to meet the following criteria:
+    ///
+    /// - The topology of the plan must remain the same: no nodes or edges may be added or removed.
+    /// - Each output partition must produce the same rows, including their multiplicity.
+    /// - The output schema, partitioning, ordering, boundedness, and emission type of the head
+    ///   node must remain unchanged. Intermediate nodes may change their output schema.
+    ///
+    /// Handlers run in registration order, each receiving the plan returned by the preceding
+    /// handler. Returning an error aborts worker plan registration.
+    ///
+    /// Register this handler on the [`SessionStateBuilder`] used by each
+    /// [`crate::WorkerSessionBuilder`]. Registering it on the coordinating session has no effect:
+    /// the handler is not sent to workers with the stage plan.
+    ///
+    /// ```rust
+    /// # use datafusion::common::Result;
+    /// # use datafusion::execution::SessionState;
+    /// # use datafusion_distributed::{DistributedExt, Worker, WorkerPlanRewriteEvent, WorkerPlanRewriteEventResponse, WorkerQueryContext};
+    ///
+    /// async fn build_worker_session(ctx: WorkerQueryContext) -> Result<SessionState> {
+    ///     Ok(ctx
+    ///         .builder
+    ///         .with_distributed_worker_plan_rewrite_handler(
+    ///             |event: WorkerPlanRewriteEvent<'_>| {
+    ///                 Ok(WorkerPlanRewriteEventResponse::new(event.plan))
+    ///             },
+    ///         )
+    ///         .build())
+    /// }
+    ///
+    /// let _worker = Worker::from_session_builder(build_worker_session);
+    /// ```
+    fn with_distributed_worker_plan_rewrite_handler<T: WorkerPlanRewriteHandler>(
+        self,
+        handler: T,
+    ) -> Self;
+
+    /// Same as [DistributedExt::with_distributed_worker_plan_rewrite_handler] but with an
+    /// in-place mutation.
+    fn set_distributed_worker_plan_rewrite_handler<T: WorkerPlanRewriteHandler>(
+        &mut self,
+        handler: T,
+    );
 }
 
 /// Trait to have a unified interface for getting structs & properties from SessionConfig that are used in distributed context.
@@ -866,6 +913,10 @@ impl DistributedExt for SessionConfig {
         RouteTasksHandlers::push_custom(self, Arc::new(h));
     }
 
+    fn set_distributed_worker_plan_rewrite_handler<H: WorkerPlanRewriteHandler>(&mut self, h: H) {
+        WorkerPlanRewriteHandlers::push_custom(self, Arc::new(h));
+    }
+
     delegate! {
         to self {
             #[call(set_distributed_option_extension)]
@@ -969,6 +1020,10 @@ impl DistributedExt for SessionConfig {
             #[call(set_distributed_route_tasks_handler)]
             #[expr($;self)]
             fn with_distributed_route_tasks_handler<H: RouteTasksHandler>(mut self, h: H) -> Self;
+
+            #[call(set_distributed_worker_plan_rewrite_handler)]
+            #[expr($;self)]
+            fn with_distributed_worker_plan_rewrite_handler<H: WorkerPlanRewriteHandler>(mut self, h: H) -> Self;
         }
     }
 }
@@ -1112,6 +1167,11 @@ impl DistributedExt for SessionStateBuilder {
             #[call(set_distributed_route_tasks_handler)]
             #[expr($;self)]
             fn with_distributed_route_tasks_handler<H: RouteTasksHandler>(mut self, h: H) -> Self;
+
+            fn set_distributed_worker_plan_rewrite_handler<H: WorkerPlanRewriteHandler>(&mut self, h: H);
+            #[call(set_distributed_worker_plan_rewrite_handler)]
+            #[expr($;self)]
+            fn with_distributed_worker_plan_rewrite_handler<H: WorkerPlanRewriteHandler>(mut self, h: H) -> Self;
         }
     }
 }
@@ -1257,6 +1317,11 @@ impl DistributedExt for SessionState {
             #[call(set_distributed_route_tasks_handler)]
             #[expr($;self)]
             fn with_distributed_route_tasks_handler<H: RouteTasksHandler>(mut self, h: H) -> Self;
+
+            fn set_distributed_worker_plan_rewrite_handler<H: WorkerPlanRewriteHandler>(&mut self, h: H);
+            #[call(set_distributed_worker_plan_rewrite_handler)]
+            #[expr($;self)]
+            fn with_distributed_worker_plan_rewrite_handler<H: WorkerPlanRewriteHandler>(mut self, h: H) -> Self;
         }
     }
 }
@@ -1395,6 +1460,11 @@ impl DistributedExt for SessionContext {
             #[call(set_distributed_route_tasks_handler)]
             #[expr($;self)]
             fn with_distributed_route_tasks_handler<H: RouteTasksHandler>(self, h: H) -> Self;
+
+            fn set_distributed_worker_plan_rewrite_handler<H: WorkerPlanRewriteHandler>(&mut self, h: H);
+            #[call(set_distributed_worker_plan_rewrite_handler)]
+            #[expr($;self)]
+            fn with_distributed_worker_plan_rewrite_handler<H: WorkerPlanRewriteHandler>(self, h: H) -> Self;
         }
     }
 }

src/events/common.rs

@@ -33,6 +33,17 @@ impl<H: ?Sized> EventHandlerChain<H> {
         // If no user handler handled the event, use the built ins.
         self.builtin.iter().find_map(|handler| f(handler.as_ref()))
     }
+
+    pub(super) fn try_fold<T, E>(
+        &self,
+        mut value: T,
+        mut f: impl FnMut(T, &H) -> Result<T, E>,
+    ) -> Result<T, E> {
+        for handler in self.custom.iter().chain(&self.builtin) {
+            value = f(value, handler.as_ref())?;
+        }
+        Ok(value)
+    }
 }
 
 impl<H: ?Sized + Send + Sync + 'static> EventHandlerChain<H> {

src/events/mod.rs

@@ -3,6 +3,7 @@ mod defaults;
 mod desired_task_count;
 mod route_tasks;
 mod scale_up_leaf_node;
+mod worker_plan_rewrite;
 
 pub(crate) use defaults::{
     file_scan_config_desired_task_count, file_scan_config_scale_up_leaf_node, random_routing,
@@ -19,3 +20,7 @@ pub(crate) use scale_up_leaf_node::ScaleUpLeafNodeHandlers;
 pub use scale_up_leaf_node::{
     ScaleUpLeafNodeEvent, ScaleUpLeafNodeEventResponse, ScaleUpLeafNodeHandler,
 };
+pub(crate) use worker_plan_rewrite::WorkerPlanRewriteHandlers;
+pub use worker_plan_rewrite::{
+    WorkerPlanRewriteEvent, WorkerPlanRewriteEventResponse, WorkerPlanRewriteHandler,
+};

src/events/worker_plan_rewrite.rs

@@ -0,0 +1,74 @@
+use super::common::EventHandlerChain;
+use datafusion::error::Result;
+use datafusion::execution::config::SessionConfig;
+use datafusion::physical_plan::ExecutionPlan;
+use std::sync::Arc;
+
+/// Information supplied while rewriting a decoded worker stage plan before registration.
+pub struct WorkerPlanRewriteEvent<'a> {
+    /// The worker-local plan. Each handler receives the plan returned by the previous handler.
+    pub plan: Arc<dyn ExecutionPlan>,
+    /// The configuration of the worker session that will execute the plan.
+    pub session_config: &'a SessionConfig,
+}
+
+/// The worker-local plan produced by a [`WorkerPlanRewriteHandler`].
+pub struct WorkerPlanRewriteEventResponse {
+    /// The plan to pass to the next handler or register for execution.
+    pub plan: Arc<dyn ExecutionPlan>,
+}
+
+impl WorkerPlanRewriteEventResponse {
+    /// Returns a response containing the rewritten worker-local plan.
+    pub fn new(plan: Arc<dyn ExecutionPlan>) -> Self {
+        Self { plan }
+    }
+}
+
+/// Rewrites a decoded worker-local plan before it is registered for execution.
+///
+/// Every registered handler runs in registration order and receives the plan returned by the
+/// previous handler. Returning an error aborts plan registration.
+pub trait WorkerPlanRewriteHandler: Send + Sync + 'static {
+    /// Returns the plan to pass to the next handler.
+    fn rewrite_worker_plan(
+        &self,
+        ev: WorkerPlanRewriteEvent,
+    ) -> Result<WorkerPlanRewriteEventResponse>;
+}
+
+impl<F> WorkerPlanRewriteHandler for F
+where
+    F: Send + Sync + 'static,
+    F: for<'a> Fn(WorkerPlanRewriteEvent<'a>) -> Result<WorkerPlanRewriteEventResponse>,
+{
+    fn rewrite_worker_plan(
+        &self,
+        ev: WorkerPlanRewriteEvent,
+    ) -> Result<WorkerPlanRewriteEventResponse> {
+        self(ev)
+    }
+}
+
+pub(crate) type WorkerPlanRewriteHandlers = EventHandlerChain<dyn WorkerPlanRewriteHandler>;
+
+impl WorkerPlanRewriteHandlers {
+    pub(crate) fn handle(ev: WorkerPlanRewriteEvent) -> Result<WorkerPlanRewriteEventResponse> {
+        let WorkerPlanRewriteEvent {
+            plan,
+            session_config,
+        } = ev;
+        let plan = match session_config.get_extension::<WorkerPlanRewriteHandlers>() {
+            Some(handlers) => handlers.try_fold(plan, |plan, handler| {
+                handler
+                    .rewrite_worker_plan(WorkerPlanRewriteEvent {
+                        plan,
+                        session_config,
+                    })
+                    .map(|response| response.plan)
+            })?,
+            None => plan,
+        };
+        Ok(WorkerPlanRewriteEventResponse::new(plan))
+    }
+}

src/lib.rs

@@ -25,7 +25,8 @@ pub use distributed_planner::{
 pub use events::{
     DesiredTaskCountEvent, DesiredTaskCountEventResponse, DesiredTaskCountHandler, RouteTasksEvent,
     RouteTasksEventResponse, RouteTasksHandler, ScaleUpLeafNodeEvent, ScaleUpLeafNodeEventResponse,
-    ScaleUpLeafNodeHandler, TaskCountAnnotation,
+    ScaleUpLeafNodeHandler, TaskCountAnnotation, WorkerPlanRewriteEvent,
+    WorkerPlanRewriteEventResponse, WorkerPlanRewriteHandler,
 };
 pub use execution_plans::{
     BroadcastExec, DistributedLeafExec, NetworkBroadcastExec, NetworkCoalesceExec,

src/worker/impl_coordinator_channel.rs

@@ -1,4 +1,5 @@
 use crate::common::TreeNodeExt;
+use crate::events::{WorkerPlanRewriteEvent, WorkerPlanRewriteHandlers};
 use crate::execution_plans::SamplerExec;
 use crate::work_unit_feed::{RemoteWorkUnitFeedRegistry, set_work_unit_received_time};
 use crate::worker::LocalWorkerContext;
@@ -86,11 +87,11 @@ impl Worker {
             let codec = DistributedCodec::new_combined_with_user(session_state.config());
             let task_ctx = session_state.task_ctx();
             let proto_node = PhysicalPlanNode::try_decode(request.plan_proto.as_ref())?;
-            let mut plan = proto_node.try_into_physical_plan(&task_ctx, &codec)?;
-
-            for hook in self.hooks.on_plan.iter() {
-                plan = hook(plan, session_state.config())?;
-            }
+            let ev = WorkerPlanRewriteEvent {
+                plan: proto_node.try_into_physical_plan(&task_ctx, &codec)?,
+                session_config: session_state.config(),
+            };
+            let plan = WorkerPlanRewriteHandlers::handle(ev)?.plan;
             load_info_rxs =
                 SamplerExec::kick_off_first_sampler(Arc::clone(&plan), Arc::clone(&task_ctx))?;