datafusion_distributed/events/
desired_task_count.rs

1use super::TaskCountAnnotation::{Desired, Maximum};
2use super::common::EventHandlerChain;
3use async_trait::async_trait;
4use datafusion::common::Result;
5use datafusion::execution::config::SessionConfig;
6use datafusion::physical_plan::ExecutionPlan;
7use std::sync::Arc;
8
9/// Annotation attached to a single [ExecutionPlan] that determines how many distributed tasks
10/// it should run on.
11#[derive(Debug, Clone, Copy)]
12pub enum TaskCountAnnotation {
13    /// The desired number of distributed tasks for this node. The final task count for the
14    /// annotated node might not be exactly this number, it is more like a hint, so depending
15    /// on the desired task count of adjacent nodes, the final task count might change.
16    Desired(usize),
17    /// Sets a maximum number of distributed tasks for this node. Typically used with the inner
18    /// value of 1, stating that this node cannot be executed in a distributed fashion.
19    Maximum(usize),
20}
21
22/// Information supplied when the planner asks a handler for a node's desired task count.
23///
24/// Handlers may return a [`DesiredTaskCountEventResponse`] for any execution-plan node. The
25/// planner reconciles responses from the nodes in the same stage into its final task count.
26#[derive(Clone, Copy)]
27pub struct DesiredTaskCountEvent<'a> {
28    /// The execution-plan node being evaluated.
29    pub plan: &'a Arc<dyn ExecutionPlan>,
30    /// The session configuration that registered the handlers and holds query options.
31    pub session_config: &'a SessionConfig,
32}
33
34/// Result of running a [TaskEstimator] on a leaf node. It tells the distributed planner hints
35/// about how many tasks should be used in [Stage]s that contain leaf nodes.
36pub struct DesiredTaskCountEventResponse {
37    /// The number of tasks that should be used in the [Stage] containing the leaf node.
38    ///
39    /// Even if implementations get to decide this number, there are situations where it can
40    /// get overridden:
41    /// - If a [Stage] contains multiple leaf nodes, the one that declares the biggest
42    ///   task_count wins.
43    /// - If there are less available workers than this number, the number of available workers
44    ///   is chosen.
45    pub task_count: TaskCountAnnotation,
46}
47
48impl DesiredTaskCountEventResponse {
49    /// Tells the distributed planner that the evaluated stage can have **at maximum** the provided
50    /// number of tasks, setting a hard upper limit.
51    ///
52    /// Returning `DesiredTaskCountEventResponse::maximum(1)` tells the distributed planner that the
53    /// evaluated stage cannot be distributed.
54    ///
55    /// Even if a `DesiredTaskCountEventResponse::maximum(N)` is provided, any other node in the
56    /// same stage providing a value of `DesiredTaskCountEventResponse::maximum(M)` where `M` < `N`
57    /// will have preference.
58    pub fn maximum(value: usize) -> Self {
59        DesiredTaskCountEventResponse {
60            task_count: Maximum(value),
61        }
62    }
63
64    /// Tells the distributed planner that the evaluated can **optimally** have the provided
65    /// number of tasks, setting a soft task count hint that can be overridden by others.
66    ///
67    /// The provided `DesiredTaskCountEventResponse::desired(N)` can be overridden by:
68    /// - Other nodes providing a `DesiredTaskCountEventResponse::desired(M)` where `M` > `N`.
69    /// - Any other node providing a `DesiredTaskCountEventResponse::maximum(M)` where `M` can be
70    ///   anything.
71    pub fn desired(value: usize) -> Self {
72        DesiredTaskCountEventResponse {
73            task_count: Desired(value),
74        }
75    }
76}
77
78#[async_trait]
79pub trait DesiredTaskCountHandler: Send + Sync + 'static {
80    /// Function applied to each node that returns a [DesiredTaskCountEventResponse] hinting how
81    /// many tasks should be used in the [Stage] containing that node, or an error if the hint
82    /// cannot be determined.
83    ///
84    /// Handlers are asynchronous and may await metadata or external services. Handler functions
85    /// return a [`DesiredTaskCountFuture`] so their futures can borrow from the event.
86    ///
87    /// All the [TaskEstimator] registered in the session will be applied to the node
88    /// until one returns an estimation.
89    ///
90    ///
91    /// If no estimation is returned from any of the registered [TaskEstimator]s, then:
92    /// - If the node is a leaf node,`Maximum(1)` is assumed, hinting the distributed planner
93    ///   that the leaf node cannot be distributed across tasks.
94    /// - If the node is a normal node in the plan, then the maximum task count from its children
95    ///   is inherited.
96    async fn handle(
97        &self,
98        ev: DesiredTaskCountEvent<'_>,
99    ) -> Option<Result<DesiredTaskCountEventResponse>>;
100}
101
102impl From<TaskCountAnnotation> for usize {
103    fn from(annotation: TaskCountAnnotation) -> Self {
104        annotation.as_usize()
105    }
106}
107
108impl TaskCountAnnotation {
109    pub fn as_usize(&self) -> usize {
110        match self {
111            Desired(desired) => *desired,
112            Maximum(maximum) => *maximum,
113        }
114    }
115
116    pub(crate) fn limit(self, limit: usize) -> Self {
117        match self {
118            Desired(desired) => Desired(desired.min(limit)),
119            Maximum(maximum) => Maximum(maximum.min(limit)),
120        }
121    }
122
123    pub(crate) fn merge(self, other: TaskCountAnnotation) -> Self {
124        match (self, other) {
125            (Desired(a), Desired(b)) => Desired(std::cmp::max(a, b)),
126            (Desired(_), Maximum(b)) => Maximum(b),
127            (Maximum(a), Desired(_)) => Maximum(a),
128            (Maximum(a), Maximum(b)) => Maximum(std::cmp::min(a, b)),
129        }
130    }
131}
132
133#[async_trait]
134impl<F> DesiredTaskCountHandler for F
135where
136    F: Send + Sync + 'static,
137    F: for<'a> Fn(DesiredTaskCountEvent<'a>) -> Option<Result<DesiredTaskCountEventResponse>>,
138{
139    async fn handle(
140        &self,
141        ev: DesiredTaskCountEvent<'_>,
142    ) -> Option<Result<DesiredTaskCountEventResponse>> {
143        self(ev)
144    }
145}
146
147#[async_trait]
148impl DesiredTaskCountHandler for usize {
149    async fn handle(
150        &self,
151        ev: DesiredTaskCountEvent<'_>,
152    ) -> Option<Result<DesiredTaskCountEventResponse>> {
153        ev.plan
154            .children()
155            .is_empty()
156            .then(|| Ok(DesiredTaskCountEventResponse::desired(*self)))
157    }
158}
159
160#[async_trait]
161impl DesiredTaskCountHandler for Arc<dyn DesiredTaskCountHandler> {
162    async fn handle(
163        &self,
164        ev: DesiredTaskCountEvent<'_>,
165    ) -> Option<Result<DesiredTaskCountEventResponse>> {
166        self.as_ref().handle(ev).await
167    }
168}
169
170pub(crate) type DesiredTaskCountHandlers = EventHandlerChain<dyn DesiredTaskCountHandler>;
171
172impl DesiredTaskCountHandlers {
173    pub(crate) async fn handle(
174        ev: DesiredTaskCountEvent<'_>,
175    ) -> Option<Result<DesiredTaskCountEventResponse>> {
176        let handlers = ev
177            .session_config
178            .get_extension::<DesiredTaskCountHandlers>()?;
179        for handler in handlers.iter() {
180            if let Some(response) = handler.handle(ev).await {
181                return Some(response);
182            }
183        }
184        None
185    }
186}
datafusion_distributed/events/desired_task_count.rs

datafusion_distributed/events/
desired_task_count.rs
