refactor: SessionConfig.extensions instead ConfigOptions.extensions

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, DistributedLeafExec, SessionStateBuilderExt, TaskEstimation, TaskEstimator,
+    TaskRoutingContext, WorkerQueryContext, display_plan_ascii,
 };
 use datafusion_proto::physical_plan::PhysicalExtensionCodec;
 use datafusion_proto::protobuf;
@@ -219,8 +219,9 @@ 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 =
+            datafusion_distributed::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,11 @@ 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 = session_config
+            .get_extension::<CombinedTaskEstimator>()
+            .map(|a| a.as_ref().clone())
+            .unwrap_or_default();
 
         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,12 +90,12 @@ impl DistributedConfig {
         Ok(distributed_cfg)
     }
 
-    /// Gets the [DistributedConfig] from the [ConfigOptions]'s extensions.
+    /// Gets the [DistributedConfig] from the [ConfigOptions]'s extensions, inserting a default if not present.
     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");
-        };
-        Ok(distributed_cfg)
+        if cfg.extensions.get::<DistributedConfig>().is_none() {
+            cfg.extensions.insert(DistributedConfig::default());
+        }
+        Ok(cfg.extensions.get_mut::<DistributedConfig>().unwrap())
     }
 
     /// Gets the [DistributedConfig] from the [ConfigOptions]'s in the provided [SessionConfig].
@@ -138,85 +107,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,21 @@ 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 = session_cfg
+                .get_extension::<CombinedTaskEstimator>()
+                .unwrap_or_else(|| Arc::new(CombinedTaskEstimator::default()));
             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 +108,9 @@ impl QueryPlanner for DistributedQueryPlanner {
             plan = Arc::new(CoalescePartitionsExec::new(plan));
         }
 
-        let cfg = session_state.config_options();
+        let arc_cfg = session_state.config_options();
 
-        plan = insert_broadcast_execs(plan, cfg)?;
+        plan = insert_broadcast_execs(plan, arc_cfg)?;
 
         if d_cfg.dynamic_task_count {
             // The task count will be decided dynamically at execution time.
@@ -116,14 +120,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, arc_cfg)?;
         let plan = push_fetch_into_network_coalesce(plan)?;
 
         Ok(Arc::new(