diff --git a/examples/custom_worker_url_routing.rs b/examples/custom_worker_url_routing.rs index 4ae9735f..fed8e7b5 100644 --- a/examples/custom_worker_url_routing.rs +++ b/examples/custom_worker_url_routing.rs @@ -41,8 +41,8 @@ use datafusion_distributed::test_utils::localhost::{ LocalHostWorkerResolver, spawn_worker_service, }; use datafusion_distributed::{ - DistributedConfig, DistributedExt, DistributedLeafExec, SessionStateBuilderExt, TaskEstimation, - TaskEstimator, TaskRoutingContext, WorkerQueryContext, display_plan_ascii, + DistributedExt, DistributedGetterExt, DistributedLeafExec, SessionStateBuilderExt, + TaskEstimation, TaskEstimator, TaskRoutingContext, WorkerQueryContext, display_plan_ascii, }; use datafusion_proto::physical_plan::PhysicalExtensionCodec; use datafusion_proto::protobuf; @@ -219,8 +219,11 @@ impl TaskEstimator for CachedFileScanConfigTaskEstimator { } fn route_tasks(&self, ctx: &TaskRoutingContext<'_>) -> Result>> { - let d_cfg = DistributedConfig::from_task_context(&ctx.task_ctx)?; - let available_urls = d_cfg.worker_resolver().get_urls()?; + let available_urls = ctx + .task_ctx + .session_config() + .get_distributed_worker_resolver()? + .get_urls()?; let mut routed = None; ctx.plan.apply(|node| { diff --git a/src/coordinator/prepare_dynamic_plan.rs b/src/coordinator/prepare_dynamic_plan.rs index 463cc58b..dfc062b2 100644 --- a/src/coordinator/prepare_dynamic_plan.rs +++ b/src/coordinator/prepare_dynamic_plan.rs @@ -140,7 +140,7 @@ pub(super) async fn prepare_dynamic_plan( }) }) }, - query_coordinator.session_config().options(), + query_coordinator.session_config(), ) .await?; diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index 887a109f..05ba741b 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -2,16 +2,18 @@ use crate::common::{TreeNodeExt, now_ns, task_ctx_with_extension}; use crate::config_extension_ext::get_config_extension_propagation_headers; use crate::coordinator::MetricsStore; use crate::coordinator::latency_metric::LatencyMetric; +use crate::distributed_planner::CombinedTaskEstimator; use crate::execution_plans::{ChildrenIsolatorUnionExec, DistributedLeafExec}; use crate::passthrough_headers::get_passthrough_headers; use crate::stage::LocalStage; +use crate::work_unit_feed::WorkUnitFeedRegistry; use crate::work_unit_feed::{build_work_unit_batch_msg, set_work_unit_send_time}; use crate::{ BytesCounterMetric, BytesMetricExt, CoordinatorToWorkerMsg, - DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedCodec, DistributedConfig, - DistributedTaskContext, DistributedWorkUnitFeedContext, LoadInfo, NetworkBoundaryExt, - SetPlanRequest, Stage, TaskEstimator, TaskKey, TaskRoutingContext, WorkUnitFeedDeclaration, - WorkerToCoordinatorMsg, get_distributed_channel_resolver, get_distributed_worker_resolver, + DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedCodec, DistributedTaskContext, + DistributedWorkUnitFeedContext, LoadInfo, NetworkBoundaryExt, SetPlanRequest, Stage, + TaskEstimator, TaskKey, TaskRoutingContext, WorkUnitFeedDeclaration, WorkerToCoordinatorMsg, + get_distributed_channel_resolver, get_distributed_worker_resolver, }; use datafusion::common::DataFusionError; use datafusion::common::instant::Instant; @@ -261,8 +263,9 @@ impl<'a> StageCoordinator<'a> { tx: UnboundedSender, ) -> Result<()> { let session_config = self.task_ctx.session_config(); - let d_cfg = DistributedConfig::from_config_options(session_config.options())?; - let wuf_registry = &d_cfg.__private_work_unit_feed_registry; + let wuf_registry = session_config + .get_extension::() + .unwrap_or_default(); let d_ctx = DistributedTaskContext { task_index: task_i, @@ -333,8 +336,9 @@ impl<'a> StageCoordinator<'a> { task_i: usize, ) -> Result<(Arc, Vec)> { let session_config = self.task_ctx.session_config(); - let d_cfg = DistributedConfig::from_config_options(session_config.options())?; - let wuf_registry = &d_cfg.__private_work_unit_feed_registry; + let wuf_registry = session_config + .get_extension::() + .unwrap_or_default(); let mut work_unit_feed_declarations = vec![]; let d_ctx = DistributedTaskContext { @@ -373,9 +377,8 @@ impl<'a> StageCoordinator<'a> { /// [TaskEstimator::route_tasks] method. pub(super) fn routed_urls(&self) -> Result> { let session_config = self.task_ctx.session_config(); - let d_cfg = DistributedConfig::from_config_options(session_config.options())?; let worker_resolver = get_distributed_worker_resolver(session_config)?; - let task_estimator = &d_cfg.__private_task_estimator; + let task_estimator = CombinedTaskEstimator::from_session_config(session_config); let routed_urls = match task_estimator.route_tasks(&TaskRoutingContext { task_ctx: Arc::clone(self.task_ctx), diff --git a/src/distributed_ext.rs b/src/distributed_ext.rs index 9efdb420..a6f9e032 100644 --- a/src/distributed_ext.rs +++ b/src/distributed_ext.rs @@ -9,7 +9,7 @@ use crate::work_unit_feed::set_distributed_work_unit_feed; use crate::worker_resolver::set_distributed_worker_resolver; use crate::{ ChannelResolver, DistributedConfig, TaskEstimator, WorkUnitFeed, WorkUnitFeedProvider, - WorkerResolver, + WorkerResolver, get_distributed_worker_resolver, }; use datafusion::common::DataFusionError; use datafusion::config::ConfigExtension; @@ -602,6 +602,13 @@ pub trait DistributedExt: Sized { ) -> Result<(), DataFusionError>; } +/// Trait to have a unified interface for getting structs & properties from SessionConfig that are used in distributed context. +pub trait DistributedGetterExt: Sized { + /// Gets the [WorkerResolver] from this session's config. Typically called inside + /// [TaskEstimator::route_tasks] to resolve available worker URLs. + fn get_distributed_worker_resolver(&self) -> Result, DataFusionError>; +} + impl DistributedExt for SessionConfig { fn set_distributed_option_extension(&mut self, t: T) { set_distributed_option_extension(self, t) @@ -855,6 +862,11 @@ impl DistributedExt for SessionConfig { } } } +impl DistributedGetterExt for SessionConfig { + fn get_distributed_worker_resolver(&self) -> Result, DataFusionError> { + get_distributed_worker_resolver(self) + } +} impl DistributedExt for SessionStateBuilder { delegate! { @@ -978,6 +990,13 @@ impl DistributedExt for SessionStateBuilder { } } } +impl DistributedGetterExt for SessionState { + delegate! { + to self.config() { + fn get_distributed_worker_resolver(&self) -> Result, DataFusionError>; + } + } +} impl DistributedExt for SessionState { delegate! { @@ -1224,3 +1243,11 @@ impl DistributedExt for SessionContext { } } } + +impl DistributedGetterExt for SessionContext { + delegate! { + to self.state_ref().read().config() { + fn get_distributed_worker_resolver(&self) -> Result, DataFusionError>; + } + } +} diff --git a/src/distributed_planner/distributed_config.rs b/src/distributed_planner/distributed_config.rs index 92cc9352..03dd156b 100644 --- a/src/distributed_planner/distributed_config.rs +++ b/src/distributed_planner/distributed_config.rs @@ -1,13 +1,8 @@ -use crate::distributed_planner::task_estimator::CombinedTaskEstimator; -use crate::protocol::ChannelResolverExtension; -use crate::work_unit_feed::WorkUnitFeedRegistry; -use crate::worker_resolver::WorkerResolverExtension; -use crate::{TaskEstimator, WorkerResolver}; -use datafusion::common::{DataFusionError, extensions_options, not_impl_err, plan_err}; -use datafusion::config::{ConfigExtension, ConfigField, ConfigOptions, Visit}; +use crate::config_extension_ext::set_distributed_option_extension; +use datafusion::common::{DataFusionError, extensions_options, plan_err}; +use datafusion::config::{ConfigExtension, ConfigOptions}; use datafusion::execution::TaskContext; use datafusion::prelude::SessionConfig; -use std::fmt::{Debug, Formatter}; use std::sync::Arc; extensions_options! { @@ -76,18 +71,6 @@ extensions_options! { /// If `dynamic_task_count` is enabled, this value is the amount of bytes/second each /// partition is expected to handle. Lower values will result in greater parallelism. pub bytes_per_partition_per_second: usize, default = 16 * 1024 * 1024 - /// Collection of [TaskEstimator]s that will be applied to leaf nodes in order to - /// estimate how many tasks should be spawned for the [Stage] containing the leaf node. - pub(crate) __private_task_estimator: CombinedTaskEstimator, default = CombinedTaskEstimator::default() - /// [ChannelResolver] implementation that tells the distributed planner information about - /// the available workers ready to execute distributed tasks. - pub(crate) __private_channel_resolver: ChannelResolverExtension, default = ChannelResolverExtension::default() - /// [WorkerResolver] implementation that tells the distributed planner information about - /// the available workers ready to execute distributed tasks. - pub(crate) __private_worker_resolver: WorkerResolverExtension, default = WorkerResolverExtension::not_implemented() - /// [WorkUnitFeedRegistry] that contains a set of getters that, applied to each node in a - /// plan, will return the [crate::WorkUnitFeed]s present in all nodes. - pub(crate) __private_work_unit_feed_registry: WorkUnitFeedRegistry, default = WorkUnitFeedRegistry::default() } } @@ -100,19 +83,6 @@ fn cardinality_task_count_factor_default() -> f64 { } impl DistributedConfig { - /// Appends a [TaskEstimator] to the list. [TaskEstimator] will be executed sequentially in - /// order on leaf nodes, and the first one to provide a value is the one that gets to decide - /// how many tasks are used for that [Stage]. - pub fn with_task_estimator( - mut self, - task_estimator: impl TaskEstimator + Send + Sync + 'static, - ) -> Self { - self.__private_task_estimator - .user_provided - .push(Arc::new(task_estimator)); - self - } - /// Gets the [DistributedConfig] from the [ConfigOptions]'s extensions. pub fn from_config_options(cfg: &ConfigOptions) -> Result<&Self, DataFusionError> { let Some(distributed_cfg) = cfg.extensions.get::() else { @@ -120,7 +90,6 @@ impl DistributedConfig { }; Ok(distributed_cfg) } - /// Gets the [DistributedConfig] from the [ConfigOptions]'s extensions. pub fn from_config_options_mut(cfg: &mut ConfigOptions) -> Result<&mut Self, DataFusionError> { let Some(distributed_cfg) = cfg.extensions.get_mut::() else { @@ -139,84 +108,20 @@ impl DistributedConfig { Self::from_session_config(ctx.session_config()) } - /// Returns the [WorkerResolver] currently in scope for this [DistributedConfig]. - pub fn worker_resolver(&self) -> &Arc { - &self.__private_worker_resolver.0 + /// Ensures that the [DistributedConfig] is present in the [SessionConfig]'s [ConfigOptions]. + /// If not, it will insert a default [DistributedConfig] into the [SessionConfig]'s [ConfigOptions]. + pub(crate) fn ensure_in_config(cfg: &mut SessionConfig) { + if cfg + .options() + .extensions + .get::() + .is_none() + { + set_distributed_option_extension(cfg, DistributedConfig::default()) + } } } impl ConfigExtension for DistributedConfig { const PREFIX: &'static str = "distributed"; } - -// FIXME: Ideally, both ChannelResolverExtension and TaskEstimators would be passed as -// extensions in SessionConfig's AnyMap instead of the ConfigOptions. However, we need -// to pass this as ConfigOptions as we need these two fields to be present during -// planning in the DistributedQueryPlanner, and the signature of the create_physical_plan() -// method there accepts a SessionState which only provides ConfigOptions. -// The following PR addresses this: https://github.com/apache/datafusion/pull/18168 -// but it still has not been accepted or merged. -// Because of this, all the boilerplate trait implementations below are needed. -impl ConfigField for ChannelResolverExtension { - fn visit(&self, _: &mut V, _: &str, _: &'static str) { - // nothing to do. - } - - fn set(&mut self, _: &str, _: &str) -> datafusion::common::Result<()> { - not_impl_err!("Not implemented") - } -} - -impl Debug for ChannelResolverExtension { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!(f, "ChannelResolverExtension") - } -} - -impl ConfigField for WorkerResolverExtension { - fn visit(&self, _: &mut V, _: &str, _: &'static str) { - // nothing to do. - } - - fn set(&mut self, _: &str, _: &str) -> datafusion::common::Result<()> { - not_impl_err!("Not implemented") - } -} - -impl Debug for WorkerResolverExtension { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!(f, "WorkerResolverExtension") - } -} - -impl ConfigField for CombinedTaskEstimator { - fn visit(&self, _: &mut V, _: &str, _: &'static str) { - //nothing to do. - } - - fn set(&mut self, _: &str, _: &str) -> Result<(), DataFusionError> { - not_impl_err!("not implemented") - } -} - -impl Debug for CombinedTaskEstimator { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!(f, "TaskEstimators") - } -} - -impl ConfigField for WorkUnitFeedRegistry { - fn visit(&self, _: &mut V, _: &str, _: &'static str) { - //nothing to do. - } - - fn set(&mut self, _: &str, _: &str) -> Result<(), DataFusionError> { - not_impl_err!("not implemented") - } -} - -impl Debug for WorkUnitFeedRegistry { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!(f, "WorkUnitFeedRegistry") - } -} diff --git a/src/distributed_planner/distributed_query_planner.rs b/src/distributed_planner/distributed_query_planner.rs index 7dcaf55e..4b6fa46d 100644 --- a/src/distributed_planner/distributed_query_planner.rs +++ b/src/distributed_planner/distributed_query_planner.rs @@ -1,4 +1,5 @@ use crate::common::TreeNodeExt; +use crate::distributed_planner::CombinedTaskEstimator; use crate::distributed_planner::inject_network_boundaries::{ CardinalityBasedNetworkBoundaryBuilder, inject_network_boundaries, }; @@ -70,7 +71,8 @@ impl QueryPlanner for DistributedQueryPlanner { return Ok(original_plan); } - let cfg = session_state.config_options().as_ref(); + let session_cfg = session_state.config(); + let cfg = session_cfg.options(); let d_cfg = DistributedConfig::from_config_options(cfg)?; // The plan already contains network boundaries set by the user. Just ensure they have nice @@ -81,7 +83,7 @@ impl QueryPlanner for DistributedQueryPlanner { if !plan.children().is_empty() { return Ok(Transformed::no(plan)); } - let task_estimator = &d_cfg.__private_task_estimator; + let task_estimator = CombinedTaskEstimator::from_session_config(session_cfg); match task_estimator.scale_up_leaf_node(&plan, task_count, cfg)? { None => Ok(Transformed::no(plan)), Some(scaled) => Ok(Transformed::yes(scaled)), @@ -116,7 +118,8 @@ impl QueryPlanner for DistributedQueryPlanner { } // Compute per-node task counts and inject `Network*Exec` nodes at the stage boundaries. - plan = inject_network_boundaries(plan, CardinalityBasedNetworkBoundaryBuilder, cfg).await?; + plan = inject_network_boundaries(plan, CardinalityBasedNetworkBoundaryBuilder, session_cfg) + .await?; plan = prepare_network_boundaries(plan)?; if !plan.exists(|plan| Ok(plan.is_network_boundary()))? { diff --git a/src/distributed_planner/inject_network_boundaries.rs b/src/distributed_planner/inject_network_boundaries.rs index 51bc4ab5..17a69a45 100644 --- a/src/distributed_planner/inject_network_boundaries.rs +++ b/src/distributed_planner/inject_network_boundaries.rs @@ -1,9 +1,11 @@ use crate::TaskCountAnnotation::{Desired, Maximum}; +use crate::distributed_planner::{CombinedTaskEstimator, TaskEstimator}; use crate::execution_plans::{ChildWeight, ChildrenIsolatorUnionExec}; use crate::stage::LocalStage; +use crate::worker_resolver::WorkerResolverExtension; use crate::{ BroadcastExec, DistributedConfig, NetworkBoundaryExt, NetworkBroadcastExec, - NetworkCoalesceExec, NetworkShuffleExec, Stage, TaskCountAnnotation, TaskEstimator, + NetworkCoalesceExec, NetworkShuffleExec, Stage, TaskCountAnnotation, }; use async_trait::async_trait; use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; @@ -17,6 +19,7 @@ use datafusion::physical_plan::repartition::RepartitionExec; use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion::physical_plan::union::UnionExec; use datafusion::physical_plan::{ExecutionPlan, PlanProperties}; +use datafusion::prelude::SessionConfig; use std::any::TypeId; use std::sync::Arc; use std::sync::Mutex; @@ -136,11 +139,14 @@ use uuid::Uuid; pub(crate) async fn inject_network_boundaries( plan: Arc, nb_builder: impl NetworkBoundaryBuilder + Send + Sync, - cfg: &ConfigOptions, + session_cfg: &SessionConfig, ) -> Result> { + let cfg = session_cfg.options(); let ctx = InjectNetworkBoundaryContext { cfg, d_cfg: DistributedConfig::from_config_options(cfg)?, + worker_resolver: WorkerResolverExtension::from_session_config(session_cfg), + task_estimator: CombinedTaskEstimator::from_session_config(session_cfg), nb_builder: &nb_builder, task_counts: &Mutex::new(HashMap::new()), query_id: Uuid::new_v4(), @@ -155,6 +161,8 @@ pub(crate) struct InjectNetworkBoundaryContext<'a> { pub(crate) d_cfg: &'a DistributedConfig, cfg: &'a ConfigOptions, + worker_resolver: Arc, + task_estimator: Arc, nb_builder: &'a (dyn NetworkBoundaryBuilder + Send + Sync), task_counts: &'a Mutex>, query_id: Uuid, @@ -164,13 +172,7 @@ pub(crate) struct InjectNetworkBoundaryContext<'a> { impl<'a> InjectNetworkBoundaryContext<'a> { pub(crate) fn max_tasks(&self) -> Result { Ok(match self.d_cfg.max_tasks_per_stage { - 0 => self - .d_cfg - .__private_worker_resolver - .0 - .get_urls()? - .len() - .max(1), + 0 => self.worker_resolver.0.get_urls()?.len().max(1), v => v, }) } @@ -228,7 +230,7 @@ async fn _inject_network_boundaries( nb_ctx: &InjectNetworkBoundaryContext<'_>, ) -> Result> { let broadcast_joins_enabled = nb_ctx.d_cfg.broadcast_joins; - let estimator = &nb_ctx.d_cfg.__private_task_estimator; + let estimator = nb_ctx.task_estimator.as_ref(); if plan.children().is_empty() { // This is a leaf node, maybe a DataSourceExec, or maybe something else custom from the @@ -416,7 +418,7 @@ impl InjectNetworkBoundaryContext<'_> { ) -> Result> { // Handle leaf nodes. if plan.children().is_empty() { - let scaled_up = self.d_cfg.__private_task_estimator.scale_up_leaf_node( + let scaled_up = self.task_estimator.as_ref().scale_up_leaf_node( plan, task_count.as_usize(), self.cfg, @@ -1267,6 +1269,8 @@ mod tests { let network_boundaries_ctx = InjectNetworkBoundaryContext { cfg: session_config.options(), d_cfg: DistributedConfig::from_config_options(session_config.options()).unwrap(), + worker_resolver: WorkerResolverExtension::from_session_config(&session_config), + task_estimator: CombinedTaskEstimator::from_session_config(&session_config), task_counts: &Mutex::new(HashMap::new()), query_id: Uuid::new_v4(), stage_id: &AtomicUsize::new(1), diff --git a/src/distributed_planner/mod.rs b/src/distributed_planner/mod.rs index 18cfe6e0..96baae96 100644 --- a/src/distributed_planner/mod.rs +++ b/src/distributed_planner/mod.rs @@ -18,5 +18,5 @@ pub(crate) use network_boundary::ProducerHead; pub use network_boundary::{NetworkBoundary, NetworkBoundaryExt}; pub use session_state_builder_ext::SessionStateBuilderExt; pub(crate) use statistics::calculate_cost; -pub(crate) use task_estimator::set_distributed_task_estimator; +pub(crate) use task_estimator::{CombinedTaskEstimator, set_distributed_task_estimator}; pub use task_estimator::{TaskCountAnnotation, TaskEstimation, TaskEstimator, TaskRoutingContext}; diff --git a/src/distributed_planner/session_state_builder_ext.rs b/src/distributed_planner/session_state_builder_ext.rs index 7d1efe24..174bf7f8 100644 --- a/src/distributed_planner/session_state_builder_ext.rs +++ b/src/distributed_planner/session_state_builder_ext.rs @@ -1,3 +1,4 @@ +use crate::distributed_planner::DistributedConfig; use crate::distributed_planner::distributed_query_planner::DistributedQueryPlanner; use datafusion::execution::SessionStateBuilder; use std::sync::Arc; @@ -16,6 +17,7 @@ pub trait SessionStateBuilderExt { impl SessionStateBuilderExt for SessionStateBuilder { fn with_distributed_planner(mut self) -> Self { + DistributedConfig::ensure_in_config(self.config().get_or_insert_default()); self.config() .get_or_insert_default() .options_mut() diff --git a/src/distributed_planner/task_estimator.rs b/src/distributed_planner/task_estimator.rs index b915e6c1..7f6d0319 100644 --- a/src/distributed_planner/task_estimator.rs +++ b/src/distributed_planner/task_estimator.rs @@ -1,5 +1,4 @@ use crate::DistributedConfig; -use crate::config_extension_ext::set_distributed_option_extension; use crate::execution_plans::DistributedLeafExec; use TaskCountAnnotation::*; use datafusion::catalog::memory::DataSourceExec; @@ -206,23 +205,12 @@ pub(crate) fn set_distributed_task_estimator( cfg: &mut SessionConfig, estimator: impl TaskEstimator + Send + Sync + 'static, ) { - let opts = cfg.options_mut(); - if let Some(distributed_cfg) = opts.extensions.get_mut::() { - distributed_cfg - .__private_task_estimator - .user_provided - .push(Arc::new(estimator)); - } else { - let mut estimators = CombinedTaskEstimator::default(); - estimators.user_provided.push(Arc::new(estimator)); - set_distributed_option_extension( - cfg, - DistributedConfig { - __private_task_estimator: estimators, - ..Default::default() - }, - ) - } + let mut combined = cfg + .get_extension::() + .map(|existing| existing.as_ref().clone()) + .unwrap_or_default(); + combined.user_provided.push(Arc::new(estimator)); + cfg.set_extension(Arc::new(combined)); } /// [TaskEstimator] implementation that acts on [DataSourceExec] nodes that contain @@ -331,6 +319,13 @@ pub(crate) struct CombinedTaskEstimator { pub(crate) user_provided: Vec>, } +impl CombinedTaskEstimator { + pub(crate) fn from_session_config(cfg: &SessionConfig) -> Arc { + cfg.get_extension::() + .unwrap_or_else(|| Arc::new(Self::default())) + } +} + impl TaskEstimator for CombinedTaskEstimator { fn task_estimation( &self, @@ -388,9 +383,7 @@ impl TaskEstimator for CombinedTaskEstimator { #[cfg(test)] mod tests { use super::*; - use crate::test_utils::in_memory_channel_resolver::InMemoryWorkerResolver; use crate::test_utils::parquet::register_parquet_tables; - use crate::worker_resolver::WorkerResolverExtension; use datafusion::error::DataFusionError; use datafusion::prelude::SessionContext; @@ -479,9 +472,6 @@ mod tests { cfg.execution.target_partitions = 1; let d_cfg = DistributedConfig { file_scan_config_bytes_per_partition: 1, - __private_worker_resolver: WorkerResolverExtension(Arc::new( - InMemoryWorkerResolver::new(3), - )), ..Default::default() }; cfg.extensions.insert(f(d_cfg)); diff --git a/src/lib.rs b/src/lib.rs index 777dd5c8..b98856d1 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -18,7 +18,7 @@ mod worker_resolver; #[cfg(feature = "grpc")] pub use arrow_ipc::CompressionType; pub use coordinator::DistributedExec; -pub use distributed_ext::DistributedExt; +pub use distributed_ext::{DistributedExt, DistributedGetterExt}; pub use distributed_planner::{ DistributedConfig, NetworkBoundary, NetworkBoundaryExt, SessionStateBuilderExt, TaskCountAnnotation, TaskEstimation, TaskEstimator, TaskRoutingContext, diff --git a/src/protocol/channel_resolver.rs b/src/protocol/channel_resolver.rs index ded82e05..e781aa90 100644 --- a/src/protocol/channel_resolver.rs +++ b/src/protocol/channel_resolver.rs @@ -1,7 +1,7 @@ -use crate::config_extension_ext::set_distributed_option_extension; +use crate::WorkerChannel; +use crate::distributed_planner::DistributedConfig; #[cfg(feature = "grpc")] use crate::protocol::grpc; -use crate::{DistributedConfig, WorkerChannel}; use async_trait::async_trait; use datafusion::common::DataFusionError; use datafusion::execution::TaskContext; @@ -44,27 +44,18 @@ pub(crate) fn set_distributed_channel_resolver( cfg: &mut SessionConfig, channel_resolver: impl ChannelResolver + Send + Sync + 'static, ) { - let opts = cfg.options_mut(); - let channel_resolver = ChannelResolverExtension(Some(Arc::new(channel_resolver))); - if let Some(distributed_cfg) = opts.extensions.get_mut::() { - distributed_cfg.__private_channel_resolver = channel_resolver; - } else { - set_distributed_option_extension( - cfg, - DistributedConfig { - __private_channel_resolver: channel_resolver, - ..Default::default() - }, - ) - } + cfg.set_extension(Arc::new(ChannelResolverExtension(Some(Arc::new( + channel_resolver, + ))))); + DistributedConfig::ensure_in_config(cfg); } pub fn get_distributed_channel_resolver( task_ctx: &TaskContext, ) -> Arc { - let opts = task_ctx.session_config().options(); - if let Some(distributed_cfg) = opts.extensions.get::() - && let Some(cr) = &distributed_cfg.__private_channel_resolver.0 + let session_cfg = task_ctx.session_config(); + if let Some(channel_resolver_ext) = session_cfg.get_extension::() + && let Some(cr) = &channel_resolver_ext.0 { return Arc::clone(cr); } diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index 5c6bdd1a..6df0e11b 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -4,8 +4,8 @@ pub mod grpc; mod channel_resolver; mod worker_channel; +pub(crate) use channel_resolver::set_distributed_channel_resolver; pub use channel_resolver::{ChannelResolver, get_distributed_channel_resolver}; -pub(crate) use channel_resolver::{ChannelResolverExtension, set_distributed_channel_resolver}; pub use worker_channel::{ CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest, GetWorkerInfoResponse, diff --git a/src/test_utils/routing.rs b/src/test_utils/routing.rs index 26dc6c10..40c080b4 100644 --- a/src/test_utils/routing.rs +++ b/src/test_utils/routing.rs @@ -24,10 +24,9 @@ use url::Url; use crate::execution_plans::DistributedLeafExec; use crate::worker::LocalWorkerContext; -use crate::{ - DistributedConfig, DistributedTaskContext, TaskEstimation, TaskEstimator, WorkerResolver, -}; +use crate::{DistributedTaskContext, TaskEstimation, TaskEstimator, WorkerResolver}; +use crate::distributed_ext::DistributedGetterExt; // Table function that creates a `URLEmitterExec` for testing task routing. #[derive(Debug)] pub struct URLEmitterFunction; @@ -291,8 +290,11 @@ impl TaskEstimator for URLEmitterTaskEstimator { } fn route_tasks(&self, routing_ctx: &crate::TaskRoutingContext<'_>) -> Result>> { - let d_cfg = DistributedConfig::from_task_context(&routing_ctx.task_ctx)?; - let mut routed_urls = d_cfg.worker_resolver().get_urls()?; + let mut routed_urls = routing_ctx + .task_ctx + .session_config() + .get_distributed_worker_resolver()? + .get_urls()?; // Trivial routing policy: Assign tasks to URLs in reverse order. routed_urls.reverse(); diff --git a/src/work_unit_feed/work_unit_feed_registry.rs b/src/work_unit_feed/work_unit_feed_registry.rs index 2a99dffb..bc006841 100644 --- a/src/work_unit_feed/work_unit_feed_registry.rs +++ b/src/work_unit_feed/work_unit_feed_registry.rs @@ -1,5 +1,4 @@ -use crate::config_extension_ext::set_distributed_option_extension; -use crate::{DistributedConfig, WorkUnit, WorkUnitFeed, WorkUnitFeedProvider}; +use crate::{WorkUnit, WorkUnitFeed, WorkUnitFeedProvider}; use datafusion::common::Result; use datafusion::execution::TaskContext; use datafusion::physical_plan::ExecutionPlan; @@ -113,21 +112,10 @@ where T::WorkUnit: 'static, F: Fn(&Arc) -> Option<&WorkUnitFeed> + Send + Sync + 'static, { - let opts = cfg.options_mut(); - if let Some(distributed_cfg) = opts.extensions.get_mut::() { - distributed_cfg - .__private_work_unit_feed_registry - .entries - .push(Arc::new(getter)); - } else { - let mut registry = WorkUnitFeedRegistry::default(); - registry.entries.push(Arc::new(getter)); - set_distributed_option_extension( - cfg, - DistributedConfig { - __private_work_unit_feed_registry: registry, - ..Default::default() - }, - ) - } + let mut registry = cfg + .get_extension::() + .map(|existing| existing.as_ref().clone()) + .unwrap_or_default(); + registry.entries.push(Arc::new(getter)); + cfg.set_extension(Arc::new(registry)); } diff --git a/src/worker_resolver.rs b/src/worker_resolver.rs index 8c661a0e..89ea1e3f 100644 --- a/src/worker_resolver.rs +++ b/src/worker_resolver.rs @@ -1,5 +1,4 @@ use crate::DistributedConfig; -use crate::config_extension_ext::set_distributed_option_extension; use datafusion::common::{DataFusionError, exec_err, not_impl_err}; use datafusion::prelude::SessionConfig; use std::any::Any; @@ -25,29 +24,19 @@ pub(crate) fn set_distributed_worker_resolver( cfg: &mut SessionConfig, worker_resolver: impl WorkerResolver + 'static, ) { - let opts = cfg.options_mut(); - let worker_resolver = WorkerResolverExtension(Arc::new(worker_resolver)); - if let Some(distributed_cfg) = opts.extensions.get_mut::() { - distributed_cfg.__private_worker_resolver = worker_resolver; - } else { - set_distributed_option_extension( - cfg, - DistributedConfig { - __private_worker_resolver: worker_resolver, - ..Default::default() - }, - ) - } + DistributedConfig::ensure_in_config(cfg); + cfg.set_extension(Arc::new(WorkerResolverExtension(Arc::new(worker_resolver)))); } +/// Gets the [WorkerResolver] from the [SessionConfig]'s extensions. Typically called inside +/// [TaskEstimator::route_tasks] to resolve the worker URLs available for distributed tasks. pub fn get_distributed_worker_resolver( cfg: &SessionConfig, ) -> Result, DataFusionError> { - let opts = cfg.options(); - let Some(distributed_cfg) = opts.extensions.get::() else { + let Some(ext) = cfg.get_extension::() else { return exec_err!("WorkerResolver not present in the session config"); }; - Ok(Arc::clone(&distributed_cfg.__private_worker_resolver.0)) + Ok(Arc::clone(&ext.0)) } #[derive(Clone)] @@ -63,6 +52,11 @@ impl WorkerResolverExtension { } Self(Arc::new(NotImplementedWorkerResolver)) } + + pub(crate) fn from_session_config(cfg: &SessionConfig) -> Arc { + cfg.get_extension::() + .unwrap_or_else(|| Arc::new(Self::not_implemented())) + } } impl WorkerResolver for Arc {