datafusion_distributed/
worker_resolver.rs

1use crate::DistributedConfig;
2use datafusion::common::{DataFusionError, exec_err, not_impl_err};
3use datafusion::prelude::SessionConfig;
4use std::any::Any;
5use std::sync::Arc;
6use url::Url;
7
8/// Resolves a list of worker URLs in the cluster available for executing parts of the plan.
9pub trait WorkerResolver: Any + Send + Sync {
10    /// Gets all available worker URLs in the cluster. Note how this method is not async, which
11    /// means that any async operation involved in discovering worker URLs must happen on a
12    /// background thread and be retrieved by this method synchronously.
13    ///
14    /// This method will be called in several places during distributed planning:
15    /// - During task count assignation for the different stages, for determining the size of
16    ///   the cluster and limiting the amount of tasks per stage to Vec<Url>.length().
17    /// - Right before execution, for lazily assigning worker URLs to the different tasks in the
18    ///   plan. This is done as close to execution in order to have fresh worker URLs as updated
19    ///   as possible.
20    fn get_urls(&self) -> Result<Vec<Url>, DataFusionError>;
21}
22
23pub(crate) fn set_distributed_worker_resolver(
24    cfg: &mut SessionConfig,
25    worker_resolver: impl WorkerResolver + 'static,
26) {
27    DistributedConfig::ensure_in_config(cfg);
28    cfg.set_extension(Arc::new(WorkerResolverExtension(Arc::new(worker_resolver))));
29}
30
31/// Gets the [WorkerResolver] from the [SessionConfig]'s extensions. Typically called inside
32/// [RouteTasksHandler::route_tasks] to resolve the worker URLs available for distributed tasks.
33pub fn get_distributed_worker_resolver(
34    cfg: &SessionConfig,
35) -> Result<Arc<dyn WorkerResolver>, DataFusionError> {
36    let Some(ext) = cfg.get_extension::<WorkerResolverExtension>() else {
37        return exec_err!("WorkerResolver not present in the session config");
38    };
39    Ok(Arc::clone(&ext.0))
40}
41
42#[derive(Clone)]
43pub(crate) struct WorkerResolverExtension(pub(crate) Arc<dyn WorkerResolver + 'static>);
44
45impl WorkerResolverExtension {
46    pub(crate) fn not_implemented() -> Self {
47        struct NotImplementedWorkerResolver;
48        impl WorkerResolver for NotImplementedWorkerResolver {
49            fn get_urls(&self) -> Result<Vec<Url>, DataFusionError> {
50                not_impl_err!("WorkerResolver::get_urls() not implemented")
51            }
52        }
53        Self(Arc::new(NotImplementedWorkerResolver))
54    }
55
56    pub(crate) fn from_session_config(cfg: &SessionConfig) -> Arc<Self> {
57        cfg.get_extension::<WorkerResolverExtension>()
58            .unwrap_or_else(|| Arc::new(Self::not_implemented()))
59    }
60}
61
62impl WorkerResolver for Arc<dyn WorkerResolver> {
63    fn get_urls(&self) -> Result<Vec<Url>, DataFusionError> {
64        self.as_ref().get_urls()
65    }
66}
datafusion_distributed/worker_resolver.rs

datafusion_distributed/
worker_resolver.rs
