datafusion_distributed/events/
scale_up_leaf_node.rs

1use super::common::EventHandlerChain;
2use datafusion::error::Result;
3use datafusion::execution::config::SessionConfig;
4use datafusion::physical_plan::ExecutionPlan;
5use std::sync::Arc;
6
7/// Information supplied when a leaf has been assigned its final stage task count.
8///
9/// This event runs after desired task counts from all of the stage's leaves have been reconciled,
10/// so [`Self::handle`] is the count that the stage will actually use. A handler can use it to
11/// replace a leaf with a task-specialized plan, such as a [`crate::DistributedLeafExec`] that
12/// selects a different input for every task.
13#[derive(Clone, Copy)]
14pub struct ScaleUpLeafNodeEvent<'a> {
15    /// The leaf execution plan to replace, if this handler recognizes it.
16    pub plan: &'a Arc<dyn ExecutionPlan>,
17    /// The final number of tasks that will execute the leaf's stage.
18    pub task_count: usize,
19    /// The session configuration that registered the event handlers and holds query options.
20    pub session_config: &'a SessionConfig,
21}
22
23/// A replacement plan returned by a [`ScaleUpLeafNodeHandler`].
24///
25/// The distributed planner annotates every node in this plan with the final stage task count, so
26/// the replacement may be a small subtree rather than only a single leaf node.
27pub struct ScaleUpLeafNodeEventResponse {
28    /// The replacement for [`ScaleUpLeafNodeEvent::plan`].
29    pub plan: Arc<dyn ExecutionPlan>,
30}
31
32impl ScaleUpLeafNodeEventResponse {
33    /// Returns a response that replaces the event's leaf with `plan`.
34    pub fn new(plan: Arc<dyn ExecutionPlan>) -> Self {
35        Self { plan }
36    }
37}
38
39/// Handles optional leaf rewrites after a stage's task count is final.
40///
41/// Handlers are evaluated in reverse registration order. Return `Ok(Some(_))` to select a
42/// replacement and stop dispatch, or `Ok(None)` to let earlier handlers try the same leaf. If all
43/// handlers return `None`, the original leaf is left unchanged. Returning an error aborts planning.
44pub trait ScaleUpLeafNodeHandler: Send + Sync + 'static {
45    /// Optionally replaces the leaf described by `ev`.
46    ///
47    /// `ev.task_count` already accounts for the constraints and desired counts of every leaf in
48    /// the stage. Implementations should use it when creating the per-task variants of the
49    /// replacement plan.
50    fn handle(&self, ev: ScaleUpLeafNodeEvent) -> Option<Result<ScaleUpLeafNodeEventResponse>>;
51}
52
53impl<F> ScaleUpLeafNodeHandler for F
54where
55    F: Send + Sync + 'static,
56    F: for<'a> Fn(ScaleUpLeafNodeEvent<'a>) -> Option<Result<ScaleUpLeafNodeEventResponse>>,
57{
58    fn handle(&self, ev: ScaleUpLeafNodeEvent) -> Option<Result<ScaleUpLeafNodeEventResponse>> {
59        self(ev)
60    }
61}
62
63impl ScaleUpLeafNodeHandler for Arc<dyn ScaleUpLeafNodeHandler> {
64    fn handle(&self, ev: ScaleUpLeafNodeEvent) -> Option<Result<ScaleUpLeafNodeEventResponse>> {
65        self.as_ref().handle(ev)
66    }
67}
68
69pub(crate) type ScaleUpLeafNodeHandlers = EventHandlerChain<dyn ScaleUpLeafNodeHandler>;
70
71impl ScaleUpLeafNodeHandlers {
72    pub(crate) fn handle(ev: ScaleUpLeafNodeEvent) -> Option<Result<ScaleUpLeafNodeEventResponse>> {
73        ev.session_config
74            .get_extension::<ScaleUpLeafNodeHandlers>()?
75            .find_map(|handler| handler.handle(ev))
76    }
77}
datafusion_distributed/events/scale_up_leaf_node.rs

datafusion_distributed/events/
scale_up_leaf_node.rs
