refactor: SessionConfig.extensions instead ConfigOptions.extensions

src/config_extension_ext.rs

@@ -1,3 +1,4 @@
+use crate::DistributedConfig;
 use datafusion::common::{DataFusionError, HashSet, internal_datafusion_err};
 use datafusion::config::ConfigExtension;
 use datafusion::prelude::SessionConfig;
@@ -128,6 +129,95 @@ pub(crate) fn get_config_extension_propagation_headers(
     Ok(headers)
 }
 
+/// Serializes the scalar fields of [DistributedConfig] from `SessionConfig.extensions`
+/// into gRPC metadata headers so workers can reconstruct the config.
+pub(crate) fn get_distributed_config_propagation_headers(
+    cfg: &SessionConfig,
+) -> Result<HeaderMap, DataFusionError> {
+    fn insert(
+        headers: &mut HeaderMap,
+        field: &str,
+        value: &str,
+    ) -> Result<(), DataFusionError> {
+        let name = format!("{FLIGHT_METADATA_CONFIG_PREFIX}distributed.{field}");
+        headers.insert(
+            HeaderName::from_str(&name)
+                .map_err(|e| internal_datafusion_err!("invalid header name: {e}"))?,
+            value
+                .parse()
+                .map_err(|e| internal_datafusion_err!("invalid header value: {e}"))?,
+        );
+        Ok(())
+    }
+
+    let Some(d) = cfg.get_extension::<DistributedConfig>() else {
+        return Ok(HeaderMap::new());
+    };
+    let mut headers = HeaderMap::new();
+    let h = &mut headers;
+    insert(h, "file_scan_config_bytes_per_partition", &d.file_scan_config_bytes_per_partition.to_string())?;
+    insert(h, "cardinality_task_count_factor", &d.cardinality_task_count_factor.to_string())?;
+    insert(h, "children_isolator_unions", &d.children_isolator_unions.to_string())?;
+    insert(h, "collect_metrics", &d.collect_metrics.to_string())?;
+    insert(h, "broadcast_joins", &d.broadcast_joins.to_string())?;
+    insert(h, "compression", &d.compression)?;
+    insert(h, "shuffle_batch_size", &d.shuffle_batch_size.to_string())?;
+    insert(h, "max_tasks_per_stage", &d.max_tasks_per_stage.to_string())?;
+    insert(h, "partial_reduce", &d.partial_reduce.to_string())?;
+    insert(h, "worker_connection_buffer_budget_bytes", &d.worker_connection_buffer_budget_bytes.to_string())?;
+    insert(h, "dynamic_task_count", &d.dynamic_task_count.to_string())?;
+    insert(h, "bytes_per_partition_per_second", &d.bytes_per_partition_per_second.to_string())?;
+    Ok(headers)
+}
+
+/// Deserializes [DistributedConfig] scalar fields from gRPC metadata headers into
+/// `SessionConfig.extensions`.  Called on the worker side when a task arrives.
+pub(crate) fn set_distributed_config_from_headers(
+    cfg: &mut SessionConfig,
+    headers: &HeaderMap,
+) -> Result<(), DataFusionError> {
+    let prefix = format!("{FLIGHT_METADATA_CONFIG_PREFIX}distributed.");
+    let mut d = cfg
+        .get_extension::<DistributedConfig>()
+        .map(|a| a.as_ref().clone())
+        .unwrap_or_default();
+
+    for (k, v) in headers.iter() {
+        let key = k.as_str();
+        if !key.starts_with(&prefix) {
+            continue;
+        }
+        let field = &key[prefix.len()..];
+        let value = v
+            .to_str()
+            .map_err(|e| internal_datafusion_err!("invalid header value: {e}"))?;
+        macro_rules! parse {
+            ($target:expr) => {{
+                $target = value
+                    .parse()
+                    .map_err(|e| internal_datafusion_err!("bad header {field}: {e}"))?;
+            }};
+        }
+        match field {
+            "file_scan_config_bytes_per_partition" => parse!(d.file_scan_config_bytes_per_partition),
+            "cardinality_task_count_factor" => parse!(d.cardinality_task_count_factor),
+            "children_isolator_unions" => parse!(d.children_isolator_unions),
+            "collect_metrics" => parse!(d.collect_metrics),
+            "broadcast_joins" => parse!(d.broadcast_joins),
+            "compression" => d.compression = value.to_string(),
+            "shuffle_batch_size" => parse!(d.shuffle_batch_size),
+            "max_tasks_per_stage" => parse!(d.max_tasks_per_stage),
+            "partial_reduce" => parse!(d.partial_reduce),
+            "worker_connection_buffer_budget_bytes" => parse!(d.worker_connection_buffer_budget_bytes),
+            "dynamic_task_count" => parse!(d.dynamic_task_count),
+            "bytes_per_partition_per_second" => parse!(d.bytes_per_partition_per_second),
+            _ => {}
+        }
+    }
+    cfg.set_extension(Arc::new(d));
+    Ok(())
+}
+
 #[cfg(test)]
 mod tests {
     use crate::config_extension_ext::{

src/coordinator/distributed.rs

@@ -209,7 +209,10 @@ impl ExecutionPlan for DistributedExec {
             //    channel is closed.
             let guard = query_coordinator.end_query_guard();
 
-            let d_cfg = DistributedConfig::from_config_options(context.session_config().options())?;
+            let d_cfg = context
+                .session_config()
+                .get_extension::<DistributedConfig>()
+                .expect("DistributedConfig should be set");
             let result = match d_cfg.dynamic_task_count {
                 true => prepare_dynamic_plan(&query_coordinator, &base_plan).await?,
                 false => prepare_static_plan(&query_coordinator, &base_plan)?,

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

@@ -1,5 +1,7 @@
 use crate::common::{TreeNodeExt, now_ns, task_ctx_with_extension};
-use crate::config_extension_ext::get_config_extension_propagation_headers;
+use crate::config_extension_ext::{
+    get_config_extension_propagation_headers, get_distributed_config_propagation_headers,
+};
 use crate::coordinator::MetricsStore;
 use crate::coordinator::latency_metric::LatencyMetric;
 use crate::execution_plans::{ChildrenIsolatorUnionExec, DistributedLeafExec};
@@ -171,6 +173,7 @@ impl<'a> StageCoordinator<'a> {
         let channel_resolver = get_distributed_channel_resolver(self.task_ctx.as_ref());
 
         let mut headers = get_config_extension_propagation_headers(session_config)?;
+        headers.extend(get_distributed_config_propagation_headers(session_config)?);
         headers.extend(get_passthrough_headers(session_config));
 
         let coordinator_to_worker_stream = futures::stream::once(async { msg })
@@ -261,7 +264,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 d_cfg = session_config
+            .get_extension::<DistributedConfig>()
+            .expect("DistributedConfig should be set");
         let wuf_registry = &d_cfg.__private_work_unit_feed_registry;
 
         let d_ctx = DistributedTaskContext {
@@ -333,7 +338,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 d_cfg = session_config
+            .get_extension::<DistributedConfig>()
+            .expect("DistributedConfig should be set");
         let wuf_registry = &d_cfg.__private_work_unit_feed_registry;
 
         let mut work_unit_feed_declarations = vec![];
@@ -373,7 +380,9 @@ 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 d_cfg = session_config
+            .get_extension::<DistributedConfig>()
+            .expect("DistributedConfig should be set");
         let worker_resolver = get_distributed_worker_resolver(session_config)?;
         let task_estimator = &d_cfg.__private_task_estimator;
 

src/distributed_ext.rs

@@ -645,38 +645,43 @@ impl DistributedExt for SessionConfig {
         &mut self,
         bytes_per_partition: usize,
     ) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        let mut d_cfg = self.get_extension::<DistributedConfig>().map(|a| a.as_ref().clone()).unwrap_or_default();
         d_cfg.file_scan_config_bytes_per_partition = bytes_per_partition;
+        self.set_extension(Arc::new(d_cfg));
         Ok(())
     }
 
     fn set_distributed_cardinality_effect_task_scale_factor(
         &mut self,
         factor: f64,
     ) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        let mut d_cfg = self.get_extension::<DistributedConfig>().map(|a| a.as_ref().clone()).unwrap_or_default();
         d_cfg.cardinality_task_count_factor = factor;
+        self.set_extension(Arc::new(d_cfg));
         Ok(())
     }
 
     fn set_distributed_metrics_collection(&mut self, enabled: bool) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        let mut d_cfg = self.get_extension::<DistributedConfig>().map(|a| a.as_ref().clone()).unwrap_or_default();
         d_cfg.collect_metrics = enabled;
+        self.set_extension(Arc::new(d_cfg));
         Ok(())
     }
 
     fn set_distributed_children_isolator_unions(
         &mut self,
         enabled: bool,
     ) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        let mut d_cfg = self.get_extension::<DistributedConfig>().map(|a| a.as_ref().clone()).unwrap_or_default();
         d_cfg.children_isolator_unions = enabled;
+        self.set_extension(Arc::new(d_cfg));
         Ok(())
     }
 
     fn set_distributed_broadcast_joins(&mut self, enabled: bool) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        let mut d_cfg = self.get_extension::<DistributedConfig>().map(|a| a.as_ref().clone()).unwrap_or_default();
         d_cfg.broadcast_joins = enabled;
+        self.set_extension(Arc::new(d_cfg));
         Ok(())
     }
 
@@ -685,21 +690,23 @@ impl DistributedExt for SessionConfig {
         &mut self,
         compression: Option<arrow_ipc::CompressionType>,
     ) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        let mut d_cfg = self.get_extension::<DistributedConfig>().map(|a| a.as_ref().clone()).unwrap_or_default();
         d_cfg.compression = match compression {
             Some(arrow_ipc::CompressionType::ZSTD) => "zstd".to_string(),
             Some(arrow_ipc::CompressionType::LZ4_FRAME) => "lz4".to_string(),
             _ => "none".to_string(),
         };
+        self.set_extension(Arc::new(d_cfg));
         Ok(())
     }
 
     fn set_distributed_shuffle_batch_size(
         &mut self,
         batch_size: usize,
     ) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        let mut d_cfg = self.get_extension::<DistributedConfig>().map(|a| a.as_ref().clone()).unwrap_or_default();
         d_cfg.shuffle_batch_size = batch_size;
+        self.set_extension(Arc::new(d_cfg));
         Ok(())
     }
 
@@ -714,23 +721,26 @@ impl DistributedExt for SessionConfig {
         &mut self,
         max_tasks_per_stage: usize,
     ) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        let mut d_cfg = self.get_extension::<DistributedConfig>().map(|a| a.as_ref().clone()).unwrap_or_default();
         d_cfg.max_tasks_per_stage = max_tasks_per_stage;
+        self.set_extension(Arc::new(d_cfg));
         Ok(())
     }
 
     fn set_distributed_partial_reduce(&mut self, enabled: bool) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        let mut d_cfg = self.get_extension::<DistributedConfig>().map(|a| a.as_ref().clone()).unwrap_or_default();
         d_cfg.partial_reduce = enabled;
+        self.set_extension(Arc::new(d_cfg));
         Ok(())
     }
 
     fn set_distributed_worker_connection_buffer_budget_bytes(
         &mut self,
         budget_bytes: usize,
     ) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        let mut d_cfg = self.get_extension::<DistributedConfig>().map(|a| a.as_ref().clone()).unwrap_or_default();
         d_cfg.worker_connection_buffer_budget_bytes = budget_bytes;
+        self.set_extension(Arc::new(d_cfg));
         Ok(())
     }
 
@@ -747,17 +757,19 @@ impl DistributedExt for SessionConfig {
     }
 
     fn set_distributed_dynamic_task_count(&mut self, enabled: bool) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        let mut d_cfg = self.get_extension::<DistributedConfig>().map(|a| a.as_ref().clone()).unwrap_or_default();
         d_cfg.dynamic_task_count = enabled;
+        self.set_extension(Arc::new(d_cfg));
         Ok(())
     }
 
     fn set_distributed_bytes_per_partition_per_second(
         &mut self,
         bytes_per_partition_per_second: usize,
     ) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        let mut d_cfg = self.get_extension::<DistributedConfig>().map(|a| a.as_ref().clone()).unwrap_or_default();
         d_cfg.bytes_per_partition_per_second = bytes_per_partition_per_second;
+        self.set_extension(Arc::new(d_cfg));
         Ok(())
     }