revert from_config_options_mut

Author: Rich-T-kid

examples/custom_worker_url_routing.rs

@@ -37,6 +37,7 @@ 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,
 };
@@ -220,8 +221,7 @@ impl TaskEstimator for CachedFileScanConfigTaskEstimator {
 
     fn route_tasks(&self, ctx: &TaskRoutingContext<'_>) -> Result<Option<Vec<Url>>> {
         let available_urls =
-            datafusion_distributed::get_distributed_worker_resolver(ctx.task_ctx.session_config())?
-                .get_urls()?;
+            get_distributed_worker_resolver(ctx.task_ctx.session_config())?.get_urls()?;
 
         let mut routed = None;
         ctx.plan.apply(|node| {

src/config_extension_ext.rs

@@ -21,19 +21,6 @@ pub(crate) fn set_distributed_option_extension<T: ConfigExtension + Default>(
     cfg.set_extension(Arc::new(propagation_ctx));
 }
 
-/// Registers `prefix` in [`ConfigExtensionPropagationContext`] so that headers for that
-/// extension are included in outgoing gRPC requests. Idempotent: safe to call multiple times.
-pub(crate) fn register_config_extension_prefix(cfg: &mut SessionConfig, prefix: &'static str) {
-    let mut ctx = cfg
-        .get_extension::<ConfigExtensionPropagationContext>()
-        .map(|existing| existing.as_ref().clone())
-        .unwrap_or_default();
-    if !ctx.prefixes.contains(&prefix) {
-        ctx.prefixes.push(prefix);
-        cfg.set_extension(Arc::new(ctx));
-    }
-}
-
 pub(crate) fn set_distributed_option_extension_from_headers<'a, T: ConfigExtension + Default>(
     cfg: &'a mut SessionConfig,
     headers: &HeaderMap,

src/coordinator/query_coordinator.rs

@@ -378,10 +378,7 @@ impl<'a> StageCoordinator<'a> {
     pub(super) fn routed_urls(&self) -> Result<Vec<Url>> {
         let session_config = self.task_ctx.session_config();
         let worker_resolver = get_distributed_worker_resolver(session_config)?;
-        let task_estimator = session_config
-            .get_extension::<CombinedTaskEstimator>()
-            .map(|a| a.as_ref().clone())
-            .unwrap_or_default();
+        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_ext.rs

@@ -645,7 +645,7 @@ impl DistributedExt for SessionConfig {
         &mut self,
         bytes_per_partition: usize,
     ) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_session_config_mut(self)?;
+        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
         d_cfg.file_scan_config_bytes_per_partition = bytes_per_partition;
         Ok(())
     }
@@ -654,13 +654,13 @@ impl DistributedExt for SessionConfig {
         &mut self,
         factor: f64,
     ) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_session_config_mut(self)?;
+        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
         d_cfg.cardinality_task_count_factor = factor;
         Ok(())
     }
 
     fn set_distributed_metrics_collection(&mut self, enabled: bool) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_session_config_mut(self)?;
+        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
         d_cfg.collect_metrics = enabled;
         Ok(())
     }
@@ -669,13 +669,13 @@ impl DistributedExt for SessionConfig {
         &mut self,
         enabled: bool,
     ) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_session_config_mut(self)?;
+        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
         d_cfg.children_isolator_unions = enabled;
         Ok(())
     }
 
     fn set_distributed_broadcast_joins(&mut self, enabled: bool) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_session_config_mut(self)?;
+        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
         d_cfg.broadcast_joins = enabled;
         Ok(())
     }
@@ -685,7 +685,7 @@ impl DistributedExt for SessionConfig {
         &mut self,
         compression: Option<arrow_ipc::CompressionType>,
     ) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_session_config_mut(self)?;
+        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
         d_cfg.compression = match compression {
             Some(arrow_ipc::CompressionType::ZSTD) => "zstd".to_string(),
             Some(arrow_ipc::CompressionType::LZ4_FRAME) => "lz4".to_string(),
@@ -698,7 +698,7 @@ impl DistributedExt for SessionConfig {
         &mut self,
         batch_size: usize,
     ) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_session_config_mut(self)?;
+        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
         d_cfg.shuffle_batch_size = batch_size;
         Ok(())
     }
@@ -714,13 +714,13 @@ impl DistributedExt for SessionConfig {
         &mut self,
         max_tasks_per_stage: usize,
     ) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_session_config_mut(self)?;
+        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
         d_cfg.max_tasks_per_stage = max_tasks_per_stage;
         Ok(())
     }
 
     fn set_distributed_partial_reduce(&mut self, enabled: bool) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_session_config_mut(self)?;
+        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
         d_cfg.partial_reduce = enabled;
         Ok(())
     }
@@ -729,7 +729,7 @@ impl DistributedExt for SessionConfig {
         &mut self,
         budget_bytes: usize,
     ) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_session_config_mut(self)?;
+        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
         d_cfg.worker_connection_buffer_budget_bytes = budget_bytes;
         Ok(())
     }
@@ -747,7 +747,7 @@ impl DistributedExt for SessionConfig {
     }
 
     fn set_distributed_dynamic_task_count(&mut self, enabled: bool) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_session_config_mut(self)?;
+        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
         d_cfg.dynamic_task_count = enabled;
         Ok(())
     }
@@ -756,7 +756,7 @@ impl DistributedExt for SessionConfig {
         &mut self,
         bytes_per_partition_per_second: usize,
     ) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_session_config_mut(self)?;
+        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
         d_cfg.bytes_per_partition_per_second = bytes_per_partition_per_second;
         Ok(())
     }

src/distributed_planner/distributed_config.rs

@@ -1,5 +1,3 @@
-use crate::config_extension_ext::register_config_extension_prefix;
-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;
@@ -92,25 +90,11 @@ impl DistributedConfig {
         Ok(distributed_cfg)
     }
 
-    /// Gets the [DistributedConfig] from the [SessionConfig], inserting a default if not present.
-    /// Always registers the `"distributed"` propagation prefix so coordinator-side settings are
-    /// included in gRPC headers sent to workers, regardless of call ordering.
-    pub fn from_session_config_mut(cfg: &mut SessionConfig) -> Result<&mut Self, DataFusionError> {
-        if cfg
-            .options()
-            .extensions
-            .get::<DistributedConfig>()
-            .is_none()
-        {
-            set_distributed_option_extension(cfg, DistributedConfig::default());
-        } else {
-            register_config_extension_prefix(cfg, DistributedConfig::PREFIX);
-        }
-        Ok(cfg
-            .options_mut()
-            .extensions
-            .get_mut::<DistributedConfig>()
-            .unwrap())
+    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)
     }
 
     /// Gets the [DistributedConfig] from the [ConfigOptions]'s in the provided [SessionConfig].

src/distributed_planner/distributed_query_planner.rs

@@ -79,9 +79,7 @@ impl QueryPlanner for DistributedQueryPlanner {
         // 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 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));

src/execution_plans/benchmarks/shuffle_bench.rs

@@ -5,7 +5,7 @@ use crate::common::task_ctx_with_extension;
 use crate::stage::RemoteStage;
 use crate::worker::WorkerConnectionPool;
 use crate::worker::test_utils::worker_handles::MemoryWorkerHandle;
-use crate::{DistributedExt, DistributedTaskContext, NetworkShuffleExec, Stage};
+use crate::{DistributedConfig, DistributedExt, DistributedTaskContext, NetworkShuffleExec, Stage};
 use arrow::datatypes::Schema;
 use arrow_ipc::CompressionType;
 use datafusion::common::Result;
@@ -161,6 +161,7 @@ impl ShuffleBench {
 
         let channel_resolver = InMemoryChannelsResolver { channels };
         let task_ctx = SessionStateBuilder::new()
+            .with_distributed_option_extension(DistributedConfig::default())
             .with_distributed_channel_resolver(channel_resolver)
             .with_distributed_compression(bench.compression)?
             .build()

src/execution_plans/benchmarks/transport_bench.rs

@@ -4,7 +4,9 @@ use super::fixture::{
 use crate::common::task_ctx_with_extension;
 use crate::stage::RemoteStage;
 use crate::worker::test_utils::worker_handles::{MemoryWorkerHandle, TcpWorkerHandle};
-use crate::{DistributedExt, DistributedTaskContext, NetworkShuffleExec, Stage, grpc};
+use crate::{
+    DistributedConfig, DistributedExt, DistributedTaskContext, NetworkShuffleExec, Stage, grpc,
+};
 use arrow::datatypes::Schema;
 use arrow_ipc::CompressionType;
 use datafusion::common::Result;
@@ -195,6 +197,7 @@ impl TransportBench {
             bench: self.clone(),
             schema,
             task_ctx: SessionStateBuilder::new()
+                .with_distributed_option_extension(DistributedConfig::default())
                 .with_distributed_channel_resolver(InMemoryChannelsResolver { channels })
                 .with_distributed_compression(self.compression)?
                 .build()
@@ -231,6 +234,7 @@ impl TransportBench {
             bench: self.clone(),
             schema,
             task_ctx: SessionStateBuilder::new()
+                .with_distributed_option_extension(DistributedConfig::default())
                 .with_distributed_channel_resolver(grpc::DefaultChannelResolver::default())
                 .with_distributed_compression(self.compression)?
                 .build()

src/worker_resolver.rs

@@ -25,7 +25,10 @@ pub(crate) fn set_distributed_worker_resolver(
     worker_resolver: impl WorkerResolver + 'static,
 ) {
     cfg.set_extension(Arc::new(WorkerResolverExtension(Arc::new(worker_resolver))));
-    let _ = DistributedConfig::from_session_config_mut(cfg);
+    let options = cfg.options_mut();
+    if options.extensions.get::<DistributedConfig>().is_none() {
+        options.extensions.insert(DistributedConfig::default());
+    }
 }
 
 pub fn get_distributed_worker_resolver(
@@ -62,49 +65,3 @@ impl WorkerResolver for Arc<dyn WorkerResolver> {
         self.as_ref().get_urls()
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use crate::config_extension_ext::get_config_extension_propagation_headers;
-    use crate::distributed_ext::DistributedExt;
-
-    struct NoopWorkerResolver;
-    impl WorkerResolver for NoopWorkerResolver {
-        fn get_urls(&self) -> Result<Vec<Url>, DataFusionError> {
-            Ok(vec![])
-        }
-    }
-
-    // Regression test: scalars set before the worker resolver must still appear in outgoing gRPC headers.
-    #[test]
-    fn distributed_config_header_propagation_scalar_before_resolver() {
-        let mut cfg = SessionConfig::new();
-        cfg.set_distributed_metrics_collection(false).unwrap();
-        set_distributed_worker_resolver(&mut cfg, NoopWorkerResolver);
-
-        let headers = get_config_extension_propagation_headers(&cfg).unwrap();
-        let value = headers
-            .get("x-datafusion-distributed-config-distributed.collect_metrics")
-            .expect("collect_metrics header must be present")
-            .to_str()
-            .unwrap();
-        assert_eq!(value, "false");
-    }
-
-    // Regression test: scalars set after the resolver must also appear in gRPC headers.
-    #[test]
-    fn distributed_config_header_propagation_scalar_after_resolver() {
-        let mut cfg = SessionConfig::new();
-        set_distributed_worker_resolver(&mut cfg, NoopWorkerResolver);
-        cfg.set_distributed_metrics_collection(false).unwrap();
-
-        let headers = get_config_extension_propagation_headers(&cfg).unwrap();
-        let value = headers
-            .get("x-datafusion-distributed-config-distributed.collect_metrics")
-            .expect("collect_metrics header must be present")
-            .to_str()
-            .unwrap();
-        assert_eq!(value, "false");
-    }
-}