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?

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.

gabotechs · 2026-07-14T09:04:32Z

@@ -116,14 +118,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, session_config_opt)?;
        let plan = push_fetch_into_network_coalesce(plan)?;



Can we leave this unrelated variable name changes out of the PR? Not that there's a problem with the new names, but if for some reason we need to revert this PR in the future, not changing unrelated things will yield less chances of conflicts with other PRs.

gabotechs · 2026-07-14T09:15:47Z


-    /// 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 {


Any reason for removing the docs in this function?

gabotechs · 2026-07-14T09:19:58Z

+    if options.extensions.get::<DistributedConfig>().is_none() {
+        options.extensions.insert(DistributedConfig::default());
    }
 }


I'm realizing that setting it like this is not going to be enough, it's not registered as distributed config, it's registered as a normal config, so it will not make it through worker jumps.

We can do here something like this:

pub(crate) fn set_distributed_worker_resolver( cfg: &mut SessionConfig, worker_resolver: impl WorkerResolver + 'static, ) { DistributedConfig::ensure_in_config(cfg); cfg.set_extension(Arc::new(WorkerResolverExtension(Arc::new(worker_resolver)))); }

And let the ensure_in_config method be like this:

pub(crate) fn ensure_in_config(cfg: &mut SessionConfig) { if cfg .options() .extensions .get::<DistributedConfig>() .is_none() { set_distributed_option_extension(cfg, DistributedConfig::default()) } }

We'll need to call the ensure_in_config method here, in set_distributed_channel_resolver and in SessionStateBuilderExt::with_distributed_planner

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

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!

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Rich-T-kid commented Jul 12, 2026 •

edited

Loading

Rich-T-kid commented Jul 13, 2026 •

edited

Loading