refactor: SessionConfig.extensions instead ConfigOptions.extensions by Rich-T-kid · Pull Request #550 · datafusion-contrib/datafusion-distributed

Rich-T-kid · 2026-07-12T20:44:34Z

closes #538

What Changed

Move 4 non-serializable DistributedConfig fields to SessionConfig.extensions

__private_task_estimator, __private_channel_resolver, __private_worker_resolver, and __private_work_unit_feed_registry no longer live in ConfigOptions.extensions. They are stored directly in SessionConfig.extensions via set_extension/get_extension, which removes all the hollow ConfigField + Debug stub impls these types required.

The main structural change is inject_network_boundaries taking &SessionConfig instead of &ConfigOptions so it can reach the AnyMap. All other changes are mechanical callsite updates.

Rich-T-kid

left some comments for reviewers on some of my rational in case it helps 🚀

gabotechs · 2026-07-13T07:20:53Z

🤔 I think you might have misunderstood the issue description in #538.

The scope is not to move DistributedConfig to SessionConfig.extensions, the idea is to just move those fields that were transported as __private_* fields in DistributedConfig as normal SessionConfig.extensions.

Citing the issue description

it should be very easy to transport these structs as normal SessionConfig.extensions objects, instead of hacking them into ConfigOptions.

Here, structs mean CombinedTaskEstimator, ChannelResolverExtension, WorkerResolverExtension and WorkUnitFeedRegistry

gabotechs

Nice blood red PR!

Rich-T-kid · 2026-07-13T16:21:45Z

Any chance to leave this one untouched? or is this change needed for this PR?

@gabotechs

from_config_options_mut() took &mut ConfigOptions, which has no access to the SessionConfig-level propagation prefix registry. So any DistributedConfig field set before the worker resolver, things like compression, collect_metrics, broadcast_joins, was never propagated to workers because the prefix was only registered inside set_distributed_worker_resolver(). Changing the signature to &mut SessionConfig lets the method register the prefix and insert DistributedConfig in one place, making both orderings safe.
I added two regression test to catch this behavior in the future.
distributed_config_header_propagation_scalar_before_resolver() and distributed_config_header_propagation_scalar_after_resolver() in worker_resolver.rs.

gabotechs · 2026-07-13T16:23:50Z

-    /// 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");
-        };
-        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())
    }


This change also seems unrelated to this PR, ca we keep the from_config_options_mut method?

Also, that way of injecting DistributedConfig in the config.extensions... looks a bit too LLMy, it's overall pretty convoluted and unnecessary.

For ensuring that the DistributedConfig is always present in the SessionConfig, rather than doing on each from_session_config_mut, I'd instead just put in in SessionStateBuilderExt, because that's the only mandatory place people should go enable distributed-datafusion:

impl SessionStateBuilderExt for SessionStateBuilder { fn with_distributed_planner(mut self) -> Self { let config = self.config().get_or_insert_default().options_mut(); config .optimizer .enable_physical_uncorrelated_scalar_subquery = false; if config.extensions.get::<DistributedConfig>().is_none() { config.extensions.insert(DistributedConfig::default()); } let prev = std::mem::take(self.query_planner()); self.with_query_planner(Arc::new(DistributedQueryPlanner { prev })) } }

But not sure if this is necessary.

Reverted this change. There were some test breaking due to this but I resolved them by adding
.with_distributed_option_extension(DistributedConfig::default()) to the sessionState builder

I'm afraid we cannot do that, because that means that in the same way this is blowing up in our tests here, it will blow up in other people's codebase, se we need to find an alternative.

gabotechs · 2026-07-14T08:59:44Z

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


🤔 Damn... this is actually how people are expected to access their worker resolver, and the API is taking a hit on usability now...

What comes to mind is, in the same way that we have the DistributedExt::with_distributed_worker_resolver and DistributedExt::set_distributed_worker_resolver, we can have the DistributedExt::get_distributed_worker_resolver for retrieving the worker resolver.

How does that sound?

This look's good but it requires adding this stub for impl DistributedExt for SessionStateBuilder

fn get_distributed_worker_resolver(&self) -> Result<Arc<dyn WorkerResolver>, DataFusionError> { Err(DataFusionError::Execution( "get_distributed_worker_resolver is not available on SessionStateBuilder; use SessionState or SessionConfig".to_string(), )) }

gabotechs · 2026-07-14T09:01:00Z

-    /// 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");
-        };
-        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())
    }


I'm afraid we cannot do that, because that means that in the same way this is blowing up in our tests here, it will blow up in other people's codebase, se we need to find an alternative.

Rich-T-kid · 2026-07-15T15:36:32Z

 use crate::{
    ChannelResolver, DistributedConfig, TaskEstimator, WorkUnitFeed, WorkUnitFeedProvider,
-    WorkerResolver,
+    WorkerResolver, get_distributed_worker_resolver as get_worker_resolver,


renamed this because otherwise
impl DistributedExt for SessionConfig :: get_distributed_worker_resolver() would called get_distributed_worker_resolver() which may be confusing.

Why would it be confusing? this is in the same situation as set_distributed_task_estimator, set_distributed_work_unit_feed, set_distributed_worker_resolver, set_distributed_channel_resolver, which have the same names than their counterpart methods in DistributedExt.

For consistency with the other functions, I'd just import this one without aliasing.

gabotechs · 2026-07-15T16:48:24Z

    /// ```
    fn with_distributed_worker_resolver<T: WorkerResolver + 'static>(self, resolver: T) -> Self;

+    fn get_distributed_worker_resolver(&self) -> Result<Arc<dyn WorkerResolver>, DataFusionError>;


We are missing a doc comment here. I'd explain what it does and one example of where is typically used.

For example, it can be typically used inside a TaskEstimator::route_tasks implementation for retrieving the available URLs at any given moment.

gabotechs · 2026-07-15T16:50:56Z

 use crate::{
    ChannelResolver, DistributedConfig, TaskEstimator, WorkUnitFeed, WorkUnitFeedProvider,
-    WorkerResolver,
+    WorkerResolver, get_distributed_worker_resolver as get_worker_resolver,


Why would it be confusing? this is in the same situation as set_distributed_task_estimator, set_distributed_work_unit_feed, set_distributed_worker_resolver, set_distributed_channel_resolver, which have the same names than their counterpart methods in DistributedExt.

For consistency with the other functions, I'd just import this one without aliasing.

gabotechs · 2026-07-15T16:54:50Z

+    fn get_distributed_worker_resolver(&self) -> Result<Arc<dyn WorkerResolver>, DataFusionError> {
+        Err(DataFusionError::Execution(
+            "get_distributed_worker_resolver is not available on SessionStateBuilder; use SessionState or SessionConfig".to_string(),
+        ))
+    }
+


🤔 damn, it's true, there's no way of accessing the worker resolver here...

Then, I think it's probably worth it to create another extension trait for getters. I bet that we are going to be in this same situation more times in the future, so it might be better to tackle this now.

Here's an idea:

pub trait DistributedGetterExt: Sized { /// TODO: nice description fn get_distributed_worker_resolver(&self) -> Result<Arc<dyn WorkerResolver>, DataFusionError>; } impl DistributedGetterExt for SessionConfig { fn get_distributed_worker_resolver(&self) -> Result<Arc<dyn WorkerResolver>, DataFusionError> { get_distributed_worker_resolver(self) } } impl DistributedGetterExt for SessionState { delegate! { to self.config() { fn get_distributed_worker_resolver(&self) -> Result<Arc<dyn WorkerResolver>, DataFusionError>; } } } impl DistributedGetterExt for SessionContext { delegate! { to self.state_ref().read().config() { fn get_distributed_worker_resolver(&self) -> Result<Arc<dyn WorkerResolver>, DataFusionError>; } } }

And expose it in lib.rs like this, along with the other DistributedExt:

pub use distributed_ext::{DistributedExt, DistributedGetterExt};

And omit the DistributedGetterExt impl for SessionStateBuilder

gabotechs · 2026-07-15T17:05:43Z

cc @timsaucer, this PR removes that hack we were doing for adding extensions as ConfigOptions. Does this help with the datafusion-python integration?

Rich-T-kid · 2026-07-15T17:42:41Z

I've introduced the changes you mentioned in beda7de

gabotechs

💯 Nice work @Rich-T-kid!

Rich-T-kid added 3 commits July 11, 2026 15:43

first working draft

fa3fb46

clean up PR

068b0da

fix lint ci

ff757be

Rich-T-kid commented Jul 12, 2026

View reviewed changes

Comment thread src/distributed_planner/distributed_query_planner.rs Outdated

Comment thread src/distributed_planner/task_estimator.rs Outdated

Comment thread src/test_utils/work_unit_file_scan.rs Outdated

Comment thread src/distributed_ext.rs Outdated

fix breaking CI test

302551b

Rich-T-kid force-pushed the rich-T-kid/use-SessionCOnfig branch from d393923 to 302551b Compare July 12, 2026 21:46

revert changes

aa5231b

Rich-T-kid force-pushed the rich-T-kid/use-SessionCOnfig branch from e066089 to aa5231b Compare July 13, 2026 13:11

Rich-T-kid added 2 commits July 13, 2026 09:34

move structs out of configOptions.extentions

86635ce

add default distrib-config

93990e6

Rich-T-kid force-pushed the rich-T-kid/use-SessionCOnfig branch from 56b6672 to 93990e6 Compare July 13, 2026 14:12

Rich-T-kid commented Jul 13, 2026

View reviewed changes

Comment thread src/distributed_planner/task_estimator.rs

Rich-T-kid commented Jul 13, 2026

View reviewed changes

Comment thread src/distributed_planner/inject_network_boundaries.rs Outdated

Rich-T-kid commented Jul 13, 2026

View reviewed changes

Comment thread src/distributed_planner/inject_network_boundaries.rs Outdated

Rich-T-kid commented Jul 13, 2026

View reviewed changes

Comment thread src/distributed_planner/distributed_query_planner.rs Outdated

provide sane defaults

b4df221

gabotechs reviewed Jul 13, 2026

View reviewed changes

Comment thread src/distributed_planner/distributed_config.rs

Comment thread src/distributed_planner/distributed_config.rs

Comment thread src/distributed_planner/inject_network_boundaries.rs Outdated

Rich-T-kid added 2 commits July 13, 2026 12:02

PR review comments

6eab687

typo

d495385

gabotechs reviewed Jul 13, 2026

View reviewed changes

revert from_config_options_mut

efb8121

Rich-T-kid force-pushed the rich-T-kid/use-SessionCOnfig branch from 15a9f98 to efb8121 Compare July 13, 2026 17:46

gabotechs reviewed Jul 14, 2026

View reviewed changes

Comment thread src/distributed_planner/distributed_config.rs Outdated

Comment thread src/worker_resolver.rs Outdated

shrink api surface & fix test

65cf1f0

Rich-T-kid commented Jul 15, 2026

View reviewed changes

Rich-T-kid added 2 commits July 15, 2026 11:38

clean up PR before review

1ba1f5b

fix panics

ffb4f23

gabotechs reviewed Jul 15, 2026

View reviewed changes

Rich-T-kid added 2 commits July 15, 2026 13:26

add comments to public methods

48a6376

introduce helper trait

beda7de

Rich-T-kid requested a review from gabotechs July 15, 2026 17:42

gabotechs approved these changes Jul 15, 2026

View reviewed changes

Comment thread src/distributed_planner/distributed_config.rs

add white space for clarity

5c1163f

gabotechs merged commit 5851c72 into datafusion-contrib:main Jul 15, 2026
31 checks passed

gabotechs mentioned this pull request Jul 24, 2026

refactor: move planning APIs to session crate apache/datafusion#23842

Open

Uh oh!

Conversation

Rich-T-kid commented Jul 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What Changed

Uh oh!

Rich-T-kid left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

gabotechs commented Jul 13, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

gabotechs left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Rich-T-kid commented Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gabotechs commented Jul 15, 2026

Uh oh!

Rich-T-kid commented Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gabotechs left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Rich-T-kid commented Jul 12, 2026 •

edited

Loading

Rich-T-kid commented Jul 13, 2026 •

edited

Loading

Rich-T-kid commented Jul 15, 2026 •

edited

Loading