feat(distributed): in_process_mode skips coordinator-to-worker gRPC plumbing

Adds `DistributedConfig::in_process_mode` (default false) and the
`with_distributed_in_process_mode` / `set_distributed_in_process_mode`
setters on `DistributedExt`.

When set, `DistributedExec::prepare_plan` still walks every network
boundary and converts `Stage::Local` to `Stage::Remote` with task→URL
routing (so transports that key off `target_task` continue to work
unchanged), but skips the per-worker gRPC plumbing:

  - `send_plan_task` (ships the worker subplan over the coordinator
    channel),
  - `metrics_collection_task` (background coordinator-channel reader for
    metrics),
  - `work_unit_feed_task` (work-unit-feed RPC).

Intended for embedders that ship the worker plan over a side channel
and provide their own [`WorkerTransport`] (e.g. a shared-memory mesh
inside a single process). The URL recorded on `RemoteStage::workers`
is meaningful to such transports only insofar as it indexes
`target_task`; the transport is free to ignore the URL string.

Default behaviour for users that don't set the flag is unchanged.

Author: mdashti

src/distributed_ext.rs

@@ -505,6 +505,15 @@ pub trait DistributedExt: Sized {
     /// Same as [DistributedExt::with_distributed_broadcast_joins_enabled] but with an in-place mutation.
     fn set_distributed_broadcast_joins(&mut self, enabled: bool) -> Result<(), DataFusionError>;
 
+    /// Skip the gRPC plan-send / metrics-collection / work-unit-feed tasks in
+    /// `DistributedExec::prepare_plan`. Intended for embedders that ship the worker plan over a
+    /// side channel and provide their own [crate::WorkerTransport]; input stages are still
+    /// converted from `Local` to `Remote` so transport-keyed routing works the same way.
+    fn with_distributed_in_process_mode(self, enabled: bool) -> Result<Self, DataFusionError>;
+
+    /// Same as [DistributedExt::with_distributed_in_process_mode] but with an in-place mutation.
+    fn set_distributed_in_process_mode(&mut self, enabled: bool) -> Result<(), DataFusionError>;
+
     /// The compression type to use for sending data over the wire.
     ///
     /// The default is [CompressionType::LZ4_FRAME].
@@ -721,6 +730,12 @@ impl DistributedExt for SessionConfig {
         Ok(())
     }
 
+    fn set_distributed_in_process_mode(&mut self, enabled: bool) -> Result<(), DataFusionError> {
+        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        d_cfg.in_process_mode = enabled;
+        Ok(())
+    }
+
     fn set_distributed_compression(
         &mut self,
         compression: Option<CompressionType>,
@@ -840,6 +855,10 @@ impl DistributedExt for SessionConfig {
             #[expr($?;Ok(self))]
             fn with_distributed_broadcast_joins(mut self, enabled: bool) -> Result<Self, DataFusionError>;
 
+            #[call(set_distributed_in_process_mode)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_in_process_mode(mut self, enabled: bool) -> Result<Self, DataFusionError>;
+
             #[call(set_distributed_compression)]
             #[expr($?;Ok(self))]
             fn with_distributed_compression(mut self, compression: Option<CompressionType>) -> Result<Self, DataFusionError>;
@@ -944,6 +963,11 @@ impl DistributedExt for SessionStateBuilder {
             #[expr($?;Ok(self))]
             fn with_distributed_broadcast_joins(mut self, enabled: bool) -> Result<Self, DataFusionError>;
 
+            fn set_distributed_in_process_mode(&mut self, enabled: bool) -> Result<(), DataFusionError>;
+            #[call(set_distributed_in_process_mode)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_in_process_mode(mut self, enabled: bool) -> Result<Self, DataFusionError>;
+
             fn set_distributed_compression(&mut self, compression: Option<CompressionType>) -> Result<(), DataFusionError>;
             #[call(set_distributed_compression)]
             #[expr($?;Ok(self))]
@@ -1060,6 +1084,11 @@ impl DistributedExt for SessionState {
             #[expr($?;Ok(self))]
             fn with_distributed_broadcast_joins(mut self, enabled: bool) -> Result<Self, DataFusionError>;
 
+            fn set_distributed_in_process_mode(&mut self, enabled: bool) -> Result<(), DataFusionError>;
+            #[call(set_distributed_in_process_mode)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_in_process_mode(mut self, enabled: bool) -> Result<Self, DataFusionError>;
+
             fn set_distributed_compression(&mut self, compression: Option<CompressionType>) -> Result<(), DataFusionError>;
             #[call(set_distributed_compression)]
             #[expr($?;Ok(self))]
@@ -1176,6 +1205,11 @@ impl DistributedExt for SessionContext {
             #[expr($?;Ok(self))]
             fn with_distributed_broadcast_joins(self, enabled: bool) -> Result<Self, DataFusionError>;
 
+            fn set_distributed_in_process_mode(&mut self, enabled: bool) -> Result<(), DataFusionError>;
+            #[call(set_distributed_in_process_mode)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_in_process_mode(self, enabled: bool) -> Result<Self, DataFusionError>;
+
             fn set_distributed_compression(&mut self, compression: Option<CompressionType>) -> Result<(), DataFusionError>;
             #[call(set_distributed_compression)]
             #[expr($?;Ok(self))]

src/distributed_planner/distributed_config.rs

@@ -41,6 +41,15 @@ extensions_options! {
         /// use broadcasting like checking build side size.
         /// For now, broadcasting all CollectLeft joins is not always beneficial.
         pub broadcast_joins: bool, default = false
+        /// When set, `DistributedExec::prepare_plan` skips the gRPC plan-send / metrics-collection
+        /// / work-unit-feed tasks. Tasks are still assigned to URLs (so input stages are still
+        /// converted from `Local` to `Remote`), but no out-of-process worker is contacted.
+        ///
+        /// Intended for embedders that ship the worker plan over a side channel and provide their
+        /// own [crate::WorkerTransport] (e.g. a shared-memory mesh inside a single process). The
+        /// transport receives the assigned URLs but is free to ignore them and route by
+        /// `target_task` directly.
+        pub in_process_mode: bool, default = false
         /// The compression used for sending data over the network between workers.
         /// It can be set to either `zstd`, `lz4` or `none`.
         pub compression: String, default = "lz4".to_string()

src/execution_plans/distributed.rs

@@ -170,6 +170,9 @@ impl DistributedExec {
     fn prepare_plan(&self, ctx: &Arc<TaskContext>) -> Result<PreparedPlan> {
         let worker_resolver = get_distributed_worker_resolver(ctx.session_config())?;
         let codec = DistributedCodec::new_combined_with_user(ctx.session_config());
+        let in_process = DistributedConfig::from_config_options(ctx.session_config().options())
+            .map(|c| c.in_process_mode)
+            .unwrap_or(false);
 
         let available_urls = worker_resolver.get_urls()?;
 
@@ -236,6 +239,13 @@ impl DistributedExec {
             let mut workers = Vec::with_capacity(stage.tasks);
             for (i, routed_url) in routed_urls.into_iter().enumerate() {
                 workers.push(routed_url.clone());
+                if in_process {
+                    // Skip the coordinator → worker gRPC plumbing: the embedder ships the
+                    // worker plan over a side channel and exposes its own `WorkerTransport`.
+                    // The URL is still recorded on the `RemoteStage` so the transport can
+                    // key off `target_task` (the index into `RemoteStage::workers`).
+                    continue;
+                }
                 // Spawn a task that sends the subplan to the chosen URL.
                 // There will be as many spawned tasks as workers.
                 let (tx, worker_rx) = spawner.send_plan_task(Arc::clone(ctx), i, routed_url)?;