add default distrib-config

Author: Rich-T-kid

examples/custom_worker_url_routing.rs

@@ -41,8 +41,8 @@ use datafusion_distributed::test_utils::localhost::{
     LocalHostWorkerResolver, spawn_worker_service,
 };
 use datafusion_distributed::{
-    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;

src/distributed_planner/distributed_config.rs

@@ -90,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].
@@ -107,7 +107,6 @@ impl DistributedConfig {
     pub fn from_task_context(ctx: &Arc<TaskContext>) -> Result<&Self, DataFusionError> {
         Self::from_session_config(ctx.session_config())
     }
-
 }
 
 impl ConfigExtension for DistributedConfig {

src/distributed_planner/inject_network_boundaries.rs

@@ -11,7 +11,6 @@ use async_trait::async_trait;
 use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion};
 use datafusion::common::{HashMap, Result, plan_err};
 use datafusion::config::ConfigOptions;
-use datafusion::prelude::SessionConfig;
 use datafusion::physical_expr::Partitioning;
 use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
 use datafusion::physical_plan::execution_plan::CardinalityEffect;
@@ -20,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;

src/distributed_planner/task_estimator.rs

@@ -376,9 +376,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;
 

src/protocol/channel_resolver.rs

@@ -1,6 +1,6 @@
+use crate::WorkerChannel;
 #[cfg(feature = "grpc")]
 use crate::protocol::grpc;
-use crate::WorkerChannel;
 use async_trait::async_trait;
 use datafusion::common::DataFusionError;
 use datafusion::execution::TaskContext;

src/protocol/mod.rs

@@ -4,8 +4,8 @@ pub mod grpc;
 mod channel_resolver;
 mod worker_channel;
 
-pub use channel_resolver::{ChannelResolver, get_distributed_channel_resolver};
 pub(crate) use channel_resolver::set_distributed_channel_resolver;
+pub use channel_resolver::{ChannelResolver, get_distributed_channel_resolver};
 
 pub use worker_channel::{
     CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest, GetWorkerInfoResponse,

src/test_utils/routing.rs

@@ -293,8 +293,7 @@ impl TaskEstimator for URLEmitterTaskEstimator {
 
     fn route_tasks(&self, routing_ctx: &crate::TaskRoutingContext<'_>) -> Result<Option<Vec<Url>>> {
         let mut routed_urls =
-            get_distributed_worker_resolver(routing_ctx.task_ctx.session_config())?
-                .get_urls()?;
+            get_distributed_worker_resolver(routing_ctx.task_ctx.session_config())?.get_urls()?;
 
         // Trivial routing policy: Assign tasks to URLs in reverse order.
         routed_urls.reverse();