refactor: SessionConfig.extensions instead ConfigOptions.extensions

examples/custom_worker_url_routing.rs

@@ -37,12 +37,13 @@ use datafusion::physical_plan::stream::{
 };
 use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
 use datafusion::prelude::{ParquetReadOptions, SessionContext};
+use datafusion_distributed::get_distributed_worker_resolver;
 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, DistributedLeafExec, SessionStateBuilderExt, TaskEstimation, TaskEstimator,
+    TaskRoutingContext, WorkerQueryContext, display_plan_ascii,
 };
 use datafusion_proto::physical_plan::PhysicalExtensionCodec;
 use datafusion_proto::protobuf;
@@ -219,8 +220,8 @@ impl TaskEstimator for CachedFileScanConfigTaskEstimator {
     }
 
     fn route_tasks(&self, ctx: &TaskRoutingContext<'_>) -> Result<Option<Vec<Url>>> {
-        let d_cfg = DistributedConfig::from_task_context(&ctx.task_ctx)?;
-        let available_urls = d_cfg.worker_resolver().get_urls()?;
+        let available_urls =
+            get_distributed_worker_resolver(ctx.task_ctx.session_config())?.get_urls()?;
 
         let mut routed = None;
         ctx.plan.apply(|node| {

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?;
 

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<CoordinatorToWorkerMsg>,
     ) -> 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::<WorkUnitFeedRegistry>()
+            .unwrap_or_default();
 
         let d_ctx = DistributedTaskContext {
             task_index: task_i,
@@ -333,8 +336,9 @@ impl<'a> StageCoordinator<'a> {
         task_i: usize,
     ) -> Result<(Arc<dyn ExecutionPlan>, Vec<WorkUnitFeedDeclaration>)> {
         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::<WorkUnitFeedRegistry>()
+            .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<Vec<Url>> {
         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),

src/distributed_planner/distributed_config.rs

@@ -1,13 +1,7 @@
-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 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 +70,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 +82,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::<DistributedConfig>() else {
@@ -121,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::<DistributedConfig>() else {
             return plan_err!("DistributedConfig is not in ConfigOptions.extensions");
@@ -138,85 +106,8 @@ impl DistributedConfig {
     pub fn from_task_context(ctx: &Arc<TaskContext>) -> Result<&Self, DataFusionError> {
         Self::from_session_config(ctx.session_config())
     }
-
-    /// Returns the [WorkerResolver] currently in scope for this [DistributedConfig].
-    pub fn worker_resolver(&self) -> &Arc<dyn WorkerResolver> {
-        &self.__private_worker_resolver.0
-    }
 }
 
 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<V: 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<V: 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<V: 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<V: 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")
-    }
-}

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,18 +71,19 @@ 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
         // unique identifiers for each stage, and move forward with it.
         if original_plan.exists(|plan| Ok(plan.is_network_boundary()))? {
             // Ensure the leafs are appropriately scaled up.
+            let task_estimator = CombinedTaskEstimator::from_session_config(session_cfg);
             let scaled = original_plan.transform_down_with_task_count(1, |plan, task_count| {
                 if !plan.children().is_empty() {
                     return Ok(Transformed::no(plan));
                 }
-                let task_estimator = &d_cfg.__private_task_estimator;
                 match task_estimator.scale_up_leaf_node(&plan, task_count, cfg)? {
                     None => Ok(Transformed::no(plan)),
                     Some(scaled) => Ok(Transformed::yes(scaled)),
@@ -104,9 +106,9 @@ impl QueryPlanner for DistributedQueryPlanner {
             plan = Arc::new(CoalescePartitionsExec::new(plan));
         }
 
-        let cfg = session_state.config_options();
+        let session_config_opt = session_state.config_options();
 
-        plan = insert_broadcast_execs(plan, cfg)?;
+        plan = insert_broadcast_execs(plan, session_config_opt)?;
 
         if d_cfg.dynamic_task_count {
             // The task count will be decided dynamically at execution time.
@@ -116,14 +118,15 @@ 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()))? {
             return Ok(original_plan);
         }
 
-        let plan = partial_reduce_below_network_shuffles(plan, cfg)?;
+        let plan = partial_reduce_below_network_shuffles(plan, session_config_opt)?;
         let plan = push_fetch_into_network_coalesce(plan)?;
 
         Ok(Arc::new(

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<dyn ExecutionPlan>,
     nb_builder: impl NetworkBoundaryBuilder + Send + Sync,
-    cfg: &ConfigOptions,
+    session_cfg: &SessionConfig,
 ) -> Result<Arc<dyn ExecutionPlan>> {
+    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<WorkerResolverExtension>,
+    task_estimator: Arc<CombinedTaskEstimator>,
     nb_builder: &'a (dyn NetworkBoundaryBuilder + Send + Sync),
     task_counts: &'a Mutex<HashMap<usize, TaskCountAnnotation>>,
     query_id: Uuid,
@@ -164,13 +172,7 @@ pub(crate) struct InjectNetworkBoundaryContext<'a> {
 impl<'a> InjectNetworkBoundaryContext<'a> {
     pub(crate) fn max_tasks(&self) -> Result<usize> {
         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<Arc<dyn ExecutionPlan>> {
     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<Arc<dyn ExecutionPlan>> {
         // 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),

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};