Introduce WorkerTransport + WorkerConnection traits

Promote the previously concrete `WorkerConnection` struct to a pair of
traits so embedded executors can plug their own non-Flight transport
into `WorkerConnectionPool` without forking the per-operator code paths.

- `WorkerConnection`: per-task connection that demultiplexes the
  underlying transport into one record-batch stream per partition.
- `WorkerTransport`: factory that opens a `WorkerConnection` for a
  given input stage / partition range / target task.

The default implementation (`FlightWorkerTransport` /
`FlightWorkerConnection`) is the existing Arrow-Flight gRPC code,
moved verbatim behind the trait. The pool's slot type changes to
`Box<dyn WorkerConnection + Send + Sync>` and `get_or_init_worker_connection`
now consults the transport registered on the `TaskContext` via a new
`get_distributed_worker_transport` accessor; callers who haven't
registered one fall back to a process-wide `FlightWorkerTransport`.

Custom transports plug in through
`DistributedExt::with_distributed_worker_transport`, mirroring the
existing channel- and worker-resolver hooks.

Drive-bys: drop the unused `on_metadata` callback from
`stream_partition` (every existing caller passed `|_meta| {}`); rename
the inner struct to `FlightWorkerConnection` and `WorkerConnection::init`
to `FlightWorkerConnection::open`; tighten the connection-slot type
into a `ConnectionSlot` alias for clippy.

All 190 lib tests + 39 integration tests pass; clippy + cargo fmt
clean.

Author: mdashti

src/distributed_ext.rs

@@ -2,13 +2,16 @@ use crate::config_extension_ext::{
     set_distributed_option_extension, set_distributed_option_extension_from_headers,
 };
 use crate::distributed_planner::set_distributed_task_estimator;
-use crate::networking::{set_distributed_channel_resolver, set_distributed_worker_resolver};
+use crate::networking::{
+    set_distributed_channel_resolver, set_distributed_worker_resolver,
+    set_distributed_worker_transport,
+};
 use crate::passthrough_headers::set_passthrough_headers;
 use crate::protobuf::{set_distributed_user_codec, set_distributed_user_codec_arc};
 use crate::work_unit_feed::set_distributed_work_unit_feed;
 use crate::{
     ChannelResolver, DistributedConfig, TaskEstimator, WorkUnitFeed, WorkUnitFeedProvider,
-    WorkerResolver,
+    WorkerResolver, WorkerTransport,
 };
 use arrow_ipc::CompressionType;
 use datafusion::common::DataFusionError;
@@ -284,6 +287,52 @@ pub trait DistributedExt: Sized {
         resolver: T,
     );
 
+    /// Overrides the [WorkerTransport] used by [crate::worker::WorkerConnectionPool] when opening
+    /// connections to remote workers. The default transport is the Arrow-Flight gRPC implementation
+    /// shipped in this crate; embedded executors can plug in their own (e.g. shared-memory queues)
+    /// without forking the per-operator code paths.
+    ///
+    /// Example:
+    ///
+    /// ```ignore
+    /// # use std::ops::Range;
+    /// # use std::sync::Arc;
+    /// # use datafusion::common::Result;
+    /// # use datafusion::execution::{SessionStateBuilder, TaskContext};
+    /// # use datafusion::physical_expr_common::metrics::ExecutionPlanMetricsSet;
+    /// # use datafusion_distributed::{
+    /// #     DistributedExt, Stage, WorkerConnection, WorkerPartitionStream, WorkerTransport,
+    /// # };
+    ///
+    /// struct MyTransport;
+    /// impl WorkerTransport for MyTransport {
+    ///     fn open(
+    ///         &self,
+    ///         _input_stage: &Stage,
+    ///         _target_partitions: Range<usize>,
+    ///         _target_task: usize,
+    ///         _ctx: &Arc<TaskContext>,
+    ///         _metrics: &ExecutionPlanMetricsSet,
+    ///     ) -> Result<Box<dyn WorkerConnection + Send + Sync>> {
+    ///         todo!()
+    ///     }
+    /// }
+    ///
+    /// let state = SessionStateBuilder::new()
+    ///     .with_distributed_worker_transport(MyTransport)
+    ///     .build();
+    /// ```
+    fn with_distributed_worker_transport<T: WorkerTransport + Send + Sync + 'static>(
+        self,
+        transport: T,
+    ) -> Self;
+
+    /// Same as [DistributedExt::with_distributed_worker_transport] but with an in-place mutation.
+    fn set_distributed_worker_transport<T: WorkerTransport + Send + Sync + 'static>(
+        &mut self,
+        transport: T,
+    );
+
     /// Adds a distributed task count estimator. [TaskEstimator]s are executed on each node
     /// sequentially until one returns an estimation on the number of tasks that should be
     /// used for the stage containing that node.
@@ -619,6 +668,13 @@ impl DistributedExt for SessionConfig {
         set_distributed_channel_resolver(self, resolver);
     }
 
+    fn set_distributed_worker_transport<T: WorkerTransport + Send + Sync + 'static>(
+        &mut self,
+        transport: T,
+    ) {
+        set_distributed_worker_transport(self, transport);
+    }
+
     fn set_distributed_task_estimator<T: TaskEstimator + Send + Sync + 'static>(
         &mut self,
         estimator: T,
@@ -756,6 +812,10 @@ impl DistributedExt for SessionConfig {
             #[expr($;self)]
             fn with_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(mut self, resolver: T) -> Self;
 
+            #[call(set_distributed_worker_transport)]
+            #[expr($;self)]
+            fn with_distributed_worker_transport<T: WorkerTransport + Send + Sync + 'static>(mut self, transport: T) -> Self;
+
             #[call(set_distributed_task_estimator)]
             #[expr($;self)]
             fn with_distributed_task_estimator<T: TaskEstimator + Send + Sync + 'static>(mut self, estimator: T) -> Self;
@@ -849,6 +909,11 @@ impl DistributedExt for SessionStateBuilder {
             #[expr($;self)]
             fn with_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(mut self, resolver: T) -> Self;
 
+            fn set_distributed_worker_transport<T: WorkerTransport + Send + Sync + 'static>(&mut self, transport: T);
+            #[call(set_distributed_worker_transport)]
+            #[expr($;self)]
+            fn with_distributed_worker_transport<T: WorkerTransport + Send + Sync + 'static>(mut self, transport: T) -> Self;
+
             fn set_distributed_task_estimator<T: TaskEstimator + Send + Sync + 'static>(&mut self, estimator: T);
             #[call(set_distributed_task_estimator)]
             #[expr($;self)]
@@ -960,6 +1025,11 @@ impl DistributedExt for SessionState {
             #[expr($;self)]
             fn with_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(mut self, resolver: T) -> Self;
 
+            fn set_distributed_worker_transport<T: WorkerTransport + Send + Sync + 'static>(&mut self, transport: T);
+            #[call(set_distributed_worker_transport)]
+            #[expr($;self)]
+            fn with_distributed_worker_transport<T: WorkerTransport + Send + Sync + 'static>(mut self, transport: T) -> Self;
+
             fn set_distributed_task_estimator<T: TaskEstimator + Send + Sync + 'static>(&mut self, estimator: T);
             #[call(set_distributed_task_estimator)]
             #[expr($;self)]
@@ -1071,6 +1141,11 @@ impl DistributedExt for SessionContext {
             #[expr($;self)]
             fn with_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(self, resolver: T) -> Self;
 
+            fn set_distributed_worker_transport<T: WorkerTransport + Send + Sync + 'static>(&mut self, transport: T);
+            #[call(set_distributed_worker_transport)]
+            #[expr($;self)]
+            fn with_distributed_worker_transport<T: WorkerTransport + Send + Sync + 'static>(self, transport: T) -> Self;
+
             fn set_distributed_task_estimator<T: TaskEstimator + Send + Sync + 'static>(&mut self, estimator: T);
             #[call(set_distributed_task_estimator)]
             #[expr($;self)]

src/distributed_planner/distributed_config.rs

@@ -1,6 +1,8 @@
 use crate::TaskEstimator;
 use crate::distributed_planner::task_estimator::CombinedTaskEstimator;
-use crate::networking::{ChannelResolverExtension, WorkerResolverExtension};
+use crate::networking::{
+    ChannelResolverExtension, WorkerResolverExtension, WorkerTransportExtension,
+};
 use crate::work_unit_feed::WorkUnitFeedRegistry;
 use datafusion::common::utils::get_available_parallelism;
 use datafusion::common::{DataFusionError, extensions_options, not_impl_err, plan_err};
@@ -74,6 +76,10 @@ extensions_options! {
         /// [WorkerResolver] implementation that tells the distributed planner information about
         /// the available workers ready to execute distributed tasks.
         pub(crate) __private_worker_resolver: WorkerResolverExtension, default = WorkerResolverExtension::not_implemented()
+        /// Optional [crate::WorkerTransport] override used by [crate::worker::WorkerConnectionPool]
+        /// when opening connections to remote workers. When unset, callers fall back to a process-
+        /// wide [crate::FlightWorkerTransport].
+        pub(crate) __private_worker_transport: WorkerTransportExtension, default = WorkerTransportExtension::default()
         /// [WorkUnitFeedRegistry] that contains a set of getters that, applied to each node in a
         /// plan, will return the [crate::WorkUnitFeed]s present in all nodes.
         pub(crate) __private_work_unit_feed_registry: WorkUnitFeedRegistry, default = WorkUnitFeedRegistry::default()
@@ -171,6 +177,22 @@ impl Debug for WorkerResolverExtension {
     }
 }
 
+impl ConfigField for WorkerTransportExtension {
+    fn visit<V: Visit>(&self, _: &mut V, _: &str, _: &'static str) {
+        // nothing to do.
+    }
+
+    fn set(&mut self, _: &str, _: &str) -> datafusion::common::Result<()> {
+        not_impl_err!("Not implemented")
+    }
+}
+
+impl Debug for WorkerTransportExtension {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "WorkerTransportExtension")
+    }
+}
+
 impl ConfigField for CombinedTaskEstimator {
     fn visit<V: Visit>(&self, _: &mut V, _: &str, _: &'static str) {
         //nothing to do.

src/execution_plans/network_broadcast.rs

@@ -251,7 +251,7 @@ impl ExecutionPlan for NetworkBroadcastExec {
                 &context,
             )?;
 
-            let stream = worker_connection.stream_partition(off + partition, |_meta| {})?;
+            let stream = worker_connection.stream_partition(off + partition)?;
             streams.push(stream);
         }
 

src/execution_plans/network_coalesce.rs

@@ -243,7 +243,7 @@ impl ExecutionPlan for NetworkCoalesceExec {
             &context,
         )?;
 
-        let stream = worker_connection.stream_partition(target_partition, |_meta| {})?;
+        let stream = worker_connection.stream_partition(target_partition)?;
 
         Ok(Box::pin(RecordBatchStreamAdapter::new(
             self.schema(),

src/execution_plans/network_shuffle.rs

@@ -237,7 +237,7 @@ impl ExecutionPlan for NetworkShuffleExec {
                 &context,
             )?;
 
-            let stream = worker_connection.stream_partition(off + partition, |_meta| {})?;
+            let stream = worker_connection.stream_partition(off + partition)?;
             streams.push(stream);
         }
 

src/lib.rs

@@ -37,6 +37,7 @@ pub use metrics::{
 pub use networking::{
     BoxCloneSyncChannel, ChannelResolver, DefaultChannelResolver, WorkerResolver,
     create_worker_client, get_distributed_channel_resolver, get_distributed_worker_resolver,
+    get_distributed_worker_transport,
 };
 pub use stage::{
     DistributedTaskContext, Stage, display_plan_ascii, display_plan_graphviz, explain_analyze,
@@ -48,8 +49,9 @@ pub use worker::generated::worker::worker_service_client::WorkerServiceClient;
 pub use worker::generated::worker::worker_service_server::WorkerServiceServer;
 pub use worker::generated::worker::{GetWorkerInfoRequest, GetWorkerInfoResponse, TaskKey};
 pub use worker::{
-    DefaultSessionBuilder, MappedWorkerSessionBuilder, MappedWorkerSessionBuilderExt, TaskData,
-    Worker, WorkerQueryContext, WorkerSessionBuilder,
+    DefaultSessionBuilder, FlightWorkerTransport, MappedWorkerSessionBuilder,
+    MappedWorkerSessionBuilderExt, TaskData, Worker, WorkerConnection, WorkerPartitionStream,
+    WorkerQueryContext, WorkerSessionBuilder, WorkerTransport,
 };
 
 pub use observability::{

src/networking/mod.rs

@@ -1,5 +1,6 @@
 mod channel_resolver;
 mod worker_resolver;
+mod worker_transport;
 
 pub use channel_resolver::{
     BoxCloneSyncChannel, ChannelResolver, DefaultChannelResolver, create_worker_client,
@@ -9,3 +10,6 @@ pub(crate) use channel_resolver::{ChannelResolverExtension, set_distributed_chan
 
 pub use worker_resolver::{WorkerResolver, get_distributed_worker_resolver};
 pub(crate) use worker_resolver::{WorkerResolverExtension, set_distributed_worker_resolver};
+
+pub use worker_transport::get_distributed_worker_transport;
+pub(crate) use worker_transport::{WorkerTransportExtension, set_distributed_worker_transport};

src/networking/worker_transport.rs

@@ -0,0 +1,52 @@
+use crate::DistributedConfig;
+use crate::config_extension_ext::set_distributed_option_extension;
+use crate::worker::{FlightWorkerTransport, WorkerTransport};
+use datafusion::execution::TaskContext;
+use datafusion::prelude::SessionConfig;
+use std::sync::{Arc, LazyLock};
+
+pub(crate) fn set_distributed_worker_transport(
+    cfg: &mut SessionConfig,
+    transport: impl WorkerTransport + Send + Sync + 'static,
+) {
+    let opts = cfg.options_mut();
+    let extension = WorkerTransportExtension(Some(Arc::new(transport)));
+    if let Some(distributed_cfg) = opts.extensions.get_mut::<DistributedConfig>() {
+        distributed_cfg.__private_worker_transport = extension;
+    } else {
+        set_distributed_option_extension(
+            cfg,
+            DistributedConfig {
+                __private_worker_transport: extension,
+                ..Default::default()
+            },
+        )
+    }
+}
+
+// The default Flight transport carries no per-runtime state (it consults the channel resolver each
+// time), so a single process-wide instance is sufficient for callers that have not registered
+// their own.
+static DEFAULT_WORKER_TRANSPORT: LazyLock<Arc<dyn WorkerTransport + Send + Sync>> =
+    LazyLock::new(|| Arc::new(FlightWorkerTransport));
+
+/// Returns the [WorkerTransport] registered on the session config attached to `task_ctx`, or a
+/// process-wide [FlightWorkerTransport] if none has been set. This is what
+/// [crate::worker::WorkerConnectionPool] consults at execute time when opening connections to
+/// remote workers.
+pub fn get_distributed_worker_transport(
+    task_ctx: &TaskContext,
+) -> Arc<dyn WorkerTransport + Send + Sync> {
+    let opts = task_ctx.session_config().options();
+    if let Some(distributed_cfg) = opts.extensions.get::<DistributedConfig>()
+        && let Some(t) = &distributed_cfg.__private_worker_transport.0
+    {
+        return Arc::clone(t);
+    }
+    Arc::clone(&DEFAULT_WORKER_TRANSPORT)
+}
+
+#[derive(Clone, Default)]
+pub(crate) struct WorkerTransportExtension(
+    pub(crate) Option<Arc<dyn WorkerTransport + Send + Sync>>,
+);

src/worker/mod.rs

@@ -6,10 +6,13 @@ mod single_write_multi_read;
 mod spawn_select_all;
 #[cfg(any(test, feature = "integration"))]
 pub(crate) mod test_utils;
+pub(crate) mod transport;
 mod worker_connection_pool;
 mod worker_service;
 
 pub(crate) use single_write_multi_read::SingleWriteMultiRead;
+pub use transport::{WorkerConnection, WorkerPartitionStream, WorkerTransport};
+pub use worker_connection_pool::FlightWorkerTransport;
 pub(crate) use worker_connection_pool::{LocalWorkerContext, WorkerConnectionPool};
 
 pub use session_builder::{

src/worker/transport.rs

@@ -0,0 +1,63 @@
+use crate::stage::RemoteStage;
+use datafusion::arrow::array::RecordBatch;
+use datafusion::common::Result;
+use datafusion::execution::TaskContext;
+use datafusion::physical_expr_common::metrics::ExecutionPlanMetricsSet;
+use futures::Stream;
+use std::ops::Range;
+use std::pin::Pin;
+use std::sync::Arc;
+
+/// A schema-less stream of record batches produced by a single partition of a [WorkerConnection].
+///
+/// Operators wrap this with [datafusion::physical_plan::stream::RecordBatchStreamAdapter] when they
+/// need a schema-bearing [datafusion::execution::SendableRecordBatchStream]. Keeping the transport
+/// free of schema responsibilities makes alternative implementations (e.g. shared-memory queues
+/// backing embedded executors) easier to write.
+pub type WorkerPartitionStream = Pin<Box<dyn Stream<Item = Result<RecordBatch>> + Send + 'static>>;
+
+/// A live connection to a single worker that demultiplexes the underlying transport into one
+/// [WorkerPartitionStream] per partition.
+///
+/// One connection handles every partition in the `target_partition_range` requested at open time,
+/// so the implementation can reuse a single underlying network/IPC stream and fan messages out to
+/// per-partition queues. Each partition can be streamed exactly once.
+pub trait WorkerConnection {
+    /// Returns the stream of record batches for `partition`. Calling this twice for the same
+    /// partition is an error.
+    fn stream_partition(&self, partition: usize) -> Result<WorkerPartitionStream>;
+}
+
+/// Factory that opens a [WorkerConnection] to a single worker task.
+///
+/// The default implementation is the Arrow-Flight gRPC transport baked into this crate. Custom
+/// transports (e.g. shared-memory queues for an embedded execution context) plug in via
+/// [crate::DistributedExt::with_distributed_worker_transport].
+pub trait WorkerTransport {
+    /// Opens a connection to the worker hosting `target_task` of `input_stage` covering the
+    /// partitions in `target_partitions`. The returned [WorkerConnection] takes ownership of any
+    /// background resources (gRPC streams, demux tasks, cancellation tokens, ...) and is expected
+    /// to clean them up on drop.
+    fn open(
+        &self,
+        input_stage: &RemoteStage,
+        target_partitions: Range<usize>,
+        target_task: usize,
+        ctx: &Arc<TaskContext>,
+        metrics: &ExecutionPlanMetricsSet,
+    ) -> Result<Box<dyn WorkerConnection + Send + Sync>>;
+}
+
+impl WorkerTransport for Arc<dyn WorkerTransport + Send + Sync> {
+    fn open(
+        &self,
+        input_stage: &RemoteStage,
+        target_partitions: Range<usize>,
+        target_task: usize,
+        ctx: &Arc<TaskContext>,
+        metrics: &ExecutionPlanMetricsSet,
+    ) -> Result<Box<dyn WorkerConnection + Send + Sync>> {
+        self.as_ref()
+            .open(input_stage, target_partitions, target_task, ctx, metrics)
+    }
+}