Added an in-process InMemoryWorkerTransport.

This PR adds a `WorkerTransport` that hosts its workers in the current process: plans are
delivered with a direct `Worker::set_task_plan` call and partitions are read straight from
the local task registry, with no gRPC underneath. It is the reference implementation for
the transport extension points and the basis for running distributed plans without the
Flight stack.

To keep the dispatch paths from drifting, plan delivery is factored into transport-neutral
pieces both transports share: `encode_task_plan` (task specialization + codec), the
`Worker::set_task_plan` core that the Flight coordinator stream now wraps, and
`collect_task_work_unit_feeds`. The in-memory read side runs one `execute_local_task` over
the whole partition range and pumps each partition into a buffer, so a consumer that
interleaves partition polls of a partitioned join can't leave partitions empty.

Author: mdashti

src/coordinator/mod.rs

@@ -1,10 +1,12 @@
 mod distributed;
 mod latency_metric;
 mod metrics_store;
+mod plan_encoding;
 mod prepare_static_plan;
 mod query_coordinator;
 
 pub use distributed::DistributedExec;
 pub use latency_metric::LatencyMetric;
 pub use metrics_store::MetricsStore;
-pub(crate) use query_coordinator::FlightWorkerDispatch;
+pub use plan_encoding::{EncodedTaskPlan, encode_task_plan};
+pub(crate) use query_coordinator::{CoordinatorToWorkerMetrics, FlightWorkerDispatch};

src/coordinator/plan_encoding.rs

@@ -0,0 +1,78 @@
+use crate::common::{TreeNodeExt, serialize_uuid};
+use crate::execution_plans::{ChildrenIsolatorUnionExec, DistributedLeafExec};
+use crate::worker::generated::worker::set_plan_request::WorkUnitFeedDeclaration;
+use crate::{DistributedCodec, DistributedConfig, DistributedTaskContext};
+use datafusion::common::Result;
+use datafusion::common::tree_node::Transformed;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::prelude::SessionConfig;
+use datafusion_proto::physical_plan::AsExecutionPlan;
+use datafusion_proto::protobuf::PhysicalPlanNode;
+use prost::Message;
+use std::sync::Arc;
+
+/// A stage plan specialized for one task and encoded for delivery, together with the work-unit
+/// feeds it declares.
+pub struct EncodedTaskPlan {
+    pub plan_proto: Vec<u8>,
+    pub feed_declarations: Vec<WorkUnitFeedDeclaration>,
+    /// Output partitions of the specialized plan. Task-isolated nodes (a union running one child
+    /// per task) make this differ from the unspecialized stage plan, so a push-based transport
+    /// sizes its per-task sinks from this. The pull transports ignore it; a push transport (such
+    /// as the shared-memory one) reads it.
+    pub partitions: usize,
+}
+
+/// Specializes `plan` for `task_index` and encodes it with the session's combined codec.
+///
+/// Shared by every transport's dispatch path so task specialization and codec selection cannot
+/// drift between transports: each task must see its own slice of task-isolated nodes (a union
+/// executing one child per task, an unexecuted [DistributedLeafExec] variant), and the worker
+/// decodes with the same codec stack the coordinator encoded with.
+pub fn encode_task_plan(
+    plan: &Arc<dyn ExecutionPlan>,
+    task_index: usize,
+    task_count: usize,
+    cfg: &SessionConfig,
+) -> Result<EncodedTaskPlan> {
+    let d_cfg = DistributedConfig::from_config_options(cfg.options())?;
+    let wuf_registry = &d_cfg.__private_work_unit_feed_registry;
+
+    let mut feed_declarations = vec![];
+    let d_ctx = DistributedTaskContext {
+        task_index,
+        task_count,
+    };
+
+    let specialized = Arc::clone(plan).transform_down_with_dt_ctx(d_ctx, |plan, d_ctx| {
+        if let Some(wuf) = wuf_registry.get_work_unit_feed(&plan) {
+            feed_declarations.push(WorkUnitFeedDeclaration {
+                id: serialize_uuid(&wuf.id()),
+                partitions: plan.properties().partitioning.partition_count() as u64,
+            });
+        };
+
+        if let Some(ciu) = plan.downcast_ref::<ChildrenIsolatorUnionExec>() {
+            let ciu = ciu.to_task_specialized(d_ctx.task_index);
+            return Ok(Transformed::yes(Arc::new(ciu)));
+        };
+
+        if let Some(dle) = plan.downcast_ref::<DistributedLeafExec>() {
+            let specialized = dle.to_task_specialized(d_ctx.task_index);
+            return Ok(Transformed::yes(specialized));
+        }
+
+        Ok(Transformed::no(plan))
+    })?;
+
+    let codec = DistributedCodec::new_combined_with_user(cfg);
+    let partitions = specialized.data.properties().partitioning.partition_count();
+    let plan_proto =
+        PhysicalPlanNode::try_from_physical_plan(specialized.data, &codec)?.encode_to_vec();
+
+    Ok(EncodedTaskPlan {
+        plan_proto,
+        feed_declarations,
+        partitions,
+    })
+}

src/coordinator/query_coordinator.rs

@@ -2,32 +2,28 @@ use crate::common::{TreeNodeExt, now_ns, serialize_uuid, task_ctx_with_extension
 use crate::config_extension_ext::get_config_extension_propagation_headers;
 use crate::coordinator::MetricsStore;
 use crate::coordinator::latency_metric::LatencyMetric;
-use crate::execution_plans::{ChildrenIsolatorUnionExec, DistributedLeafExec};
+use crate::coordinator::plan_encoding::encode_task_plan;
 use crate::passthrough_headers::get_passthrough_headers;
 use crate::protobuf::tonic_status_to_datafusion_error;
 use crate::work_unit_feed::{build_work_unit_batch_msg, set_work_unit_send_time};
 use crate::worker::generated::worker as pb;
 use crate::worker::generated::worker::coordinator_to_worker_msg::Inner;
-use crate::worker::generated::worker::set_plan_request::WorkUnitFeedDeclaration;
 use crate::worker::{WorkerDispatch, WorkerDispatchRequest};
 use crate::{
-    BytesCounterMetric, BytesMetricExt, DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedCodec,
-    DistributedConfig, DistributedTaskContext, DistributedWorkUnitFeedContext, TaskKey,
+    BytesCounterMetric, BytesMetricExt, DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedConfig,
+    DistributedTaskContext, DistributedWorkUnitFeedContext, TaskKey,
     get_distributed_channel_resolver,
 };
 use datafusion::common::Result;
 use datafusion::common::instant::Instant;
 use datafusion::common::runtime::JoinSet;
-use datafusion::common::tree_node::{Transformed, TreeNodeRecursion};
+use datafusion::common::tree_node::TreeNodeRecursion;
 use datafusion::common::{DataFusionError, exec_datafusion_err};
 use datafusion::execution::TaskContext;
 use datafusion::physical_expr_common::metrics::{ExecutionPlanMetricsSet, Label, MetricBuilder};
 use datafusion::physical_plan::ExecutionPlan;
-use datafusion_proto::physical_plan::AsExecutionPlan;
-use datafusion_proto::protobuf::PhysicalPlanNode;
 use futures::{Stream, StreamExt};
 use http::Extensions;
-use prost::Message;
 use std::sync::{Arc, OnceLock};
 use tokio::sync::Notify;
 use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
@@ -128,13 +124,9 @@ impl<'a> StageCoordinator<'a> {
         UnboundedReceiver<pb::WorkerToCoordinatorMsg>,
     )> {
         let session_config = self.task_ctx.session_config();
-        let codec = DistributedCodec::new_combined_with_user(session_config);
 
-        let (specialized, work_unit_feed_declarations) = self.task_specialized_plan(task_i)?;
-
-        let plan_proto =
-            PhysicalPlanNode::try_from_physical_plan(specialized, &codec)?.encode_to_vec();
-        let plan_size = plan_proto.len();
+        let encoded = encode_task_plan(self.plan, task_i, self.task_count, session_config)?;
+        let plan_size = encoded.plan_proto.len();
 
         let task_key = TaskKey {
             query_id: serialize_uuid(&self.query_id),
@@ -143,10 +135,10 @@ impl<'a> StageCoordinator<'a> {
         };
         let msg = pb::CoordinatorToWorkerMsg {
             inner: Some(Inner::SetPlanRequest(pb::SetPlanRequest {
-                plan_proto,
+                plan_proto: encoded.plan_proto,
                 task_count: self.task_count as u64,
                 task_key: Some(task_key.clone()),
-                work_unit_feed_declarations,
+                work_unit_feed_declarations: encoded.feed_declarations,
                 target_worker_url: url.to_string(),
                 query_start_time_ns: self.metrics.instantiation_time,
             })),
@@ -309,48 +301,6 @@ impl<'a> StageCoordinator<'a> {
         });
         Ok(())
     }
-
-    /// Specializes the [Arc<dyn ExecutionPlan>] for this stage to provided task index. This implies
-    /// trimming down any unnecessary information that the specific `task_i` task is not going to
-    /// need, like unexecuted branches in [ChildrenIsolatorUnionExec], or unexecuted variants of
-    /// [DistributedLeafExec].
-    fn task_specialized_plan(
-        &self,
-        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 wuf_registry = &d_cfg.__private_work_unit_feed_registry;
-
-        let mut work_unit_feed_declarations = vec![];
-        let d_ctx = DistributedTaskContext {
-            task_index: task_i,
-            task_count: self.task_count,
-        };
-
-        let plan = Arc::clone(self.plan);
-        let transformed = plan.transform_down_with_dt_ctx(d_ctx, |plan, d_ctx| {
-            if let Some(wuf) = wuf_registry.get_work_unit_feed(&plan) {
-                work_unit_feed_declarations.push(WorkUnitFeedDeclaration {
-                    id: serialize_uuid(&wuf.id()),
-                    partitions: plan.properties().partitioning.partition_count() as u64,
-                });
-            };
-
-            if let Some(ciu) = plan.downcast_ref::<ChildrenIsolatorUnionExec>() {
-                let ciu = ciu.to_task_specialized(d_ctx.task_index);
-                return Ok(Transformed::yes(Arc::new(ciu)));
-            };
-
-            if let Some(dle) = plan.downcast_ref::<DistributedLeafExec>() {
-                let specialized = dle.to_task_specialized(d_ctx.task_index);
-                return Ok(Transformed::yes(specialized));
-            }
-
-            Ok(Transformed::no(plan))
-        })?;
-        Ok((transformed.data, work_unit_feed_declarations))
-    }
 }
 
 fn keep_stream_alive<T: 'static>(notify: Arc<Notify>) -> impl Stream<Item = T> + 'static {
@@ -359,14 +309,14 @@ fn keep_stream_alive<T: 'static>(notify: Arc<Notify>) -> impl Stream<Item = T> +
 
 /// Metrics that measure network details about communications between [DistributedExec] and a worker.
 #[derive(Clone)]
-pub(super) struct CoordinatorToWorkerMetrics {
-    pub(super) plan_bytes_sent: BytesCounterMetric,
-    pub(super) plan_send_latency: Arc<LatencyMetric>,
-    pub(super) instantiation_time: u64,
+pub(crate) struct CoordinatorToWorkerMetrics {
+    pub(crate) plan_bytes_sent: BytesCounterMetric,
+    pub(crate) plan_send_latency: Arc<LatencyMetric>,
+    pub(crate) instantiation_time: u64,
 }
 
 impl CoordinatorToWorkerMetrics {
-    pub(super) fn new(metrics: &ExecutionPlanMetricsSet) -> Self {
+    pub(crate) fn new(metrics: &ExecutionPlanMetricsSet) -> Self {
         Self {
             // Metric that measures to total sum of bytes worth of subplans sent.
             plan_bytes_sent: MetricBuilder::new(metrics)

src/lib.rs

@@ -24,7 +24,9 @@ pub use common::{
     TreeNodeExt, deserialize_uuid, get_distributed_cancellation_token, serialize_uuid,
 };
 pub use config_extension_ext::get_config_extension_propagation_headers;
-pub use coordinator::{DistributedExec, LatencyMetric, MetricsStore};
+pub use coordinator::{
+    DistributedExec, EncodedTaskPlan, LatencyMetric, MetricsStore, encode_task_plan,
+};
 pub use distributed_ext::DistributedExt;
 pub use distributed_planner::{
     DistributedConfig, NetworkBoundary, NetworkBoundaryExt, PartitionRoute, ProducerHead,
@@ -51,15 +53,18 @@ pub use stage::{
     explain_analyze,
 };
 pub use work_unit_feed::{
-    DistributedWorkUnitFeedContext, WorkUnit, WorkUnitFeed, WorkUnitFeedProto, WorkUnitFeedProvider,
+    DistributedWorkUnitFeedContext, RemoteWorkUnitFeedRegistry, RemoteWorkUnitFeedRxs,
+    RemoteWorkUnitFeedTxs, WorkUnit, WorkUnitFeed, WorkUnitFeedProto, WorkUnitFeedProvider,
+    WorkUnitRx, WorkUnitTx, collect_task_work_unit_feeds, set_received_time, set_sent_time,
 };
 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, FlightWorkerTransport, MappedWorkerSessionBuilder,
-    MappedWorkerSessionBuilderExt, TaskData, Worker, WorkerConnection, WorkerDispatch,
-    WorkerDispatchRequest, WorkerQueryContext, WorkerSessionBuilder, WorkerTransport,
+    DefaultSessionBuilder, FlightWorkerTransport, InMemoryWorkerTransport,
+    MappedWorkerSessionBuilder, MappedWorkerSessionBuilderExt, TaskData, Worker, WorkerConnection,
+    WorkerDispatch, WorkerDispatchRequest, WorkerQueryContext, WorkerSessionBuilder,
+    WorkerTransport,
 };
 
 pub use observability::{

src/work_unit_feed/coordinator.rs

@@ -0,0 +1,61 @@
+use crate::common::{TreeNodeExt, task_ctx_with_extension};
+use crate::work_unit_feed::remote_work_unit_feed::build_work_unit;
+use crate::worker::generated::worker as pb;
+use crate::{DistributedConfig, DistributedTaskContext, DistributedWorkUnitFeedContext};
+use datafusion::common::Result;
+use datafusion::common::tree_node::TreeNodeRecursion;
+use datafusion::execution::TaskContext;
+use datafusion::physical_plan::ExecutionPlan;
+use futures::StreamExt;
+use futures::stream::BoxStream;
+use std::sync::Arc;
+
+/// Drives every registered [crate::WorkUnitFeed] in `plan` for one task and returns its
+/// per-partition work-unit streams, already encoded.
+///
+/// Transport-neutral: the Flight dispatch wraps each [pb::WorkUnit] in its envelope and pushes it
+/// over the coordinator-to-worker gRPC stream; an in-process transport hands them straight to the
+/// worker-side channels. Each unit carries its `(feed id, task-local partition)`; two tasks emit
+/// identical pairs, so routing to the right task stays the transport's job, as with Flight's
+/// per-task stream. A task owns a non-overlapping window of partition feeds, offset by its task
+/// index, and each partition feed can be consumed once.
+pub fn collect_task_work_unit_feeds(
+    plan: &Arc<dyn ExecutionPlan>,
+    ctx: &Arc<TaskContext>,
+    task_index: usize,
+    task_count: usize,
+) -> Result<Vec<BoxStream<'static, Result<pb::WorkUnit>>>> {
+    let d_cfg = DistributedConfig::from_config_options(ctx.session_config().options())?;
+    let registry = &d_cfg.__private_work_unit_feed_registry;
+
+    let d_ctx = DistributedTaskContext {
+        task_index,
+        task_count,
+    };
+    let mut streams = vec![];
+    plan.apply_with_dt_ctx(d_ctx, |plan, d_ctx| {
+        let Some(wuf) = registry.get_work_unit_feed(plan) else {
+            return Ok(TreeNodeRecursion::Continue);
+        };
+
+        let partitions = plan.properties().partitioning.partition_count();
+        let start_partition = partitions * d_ctx.task_index;
+        let end_partition = start_partition + partitions;
+
+        let dist_feed_ctx = DistributedWorkUnitFeedContext {
+            fan_out_tasks: d_ctx.task_count,
+        };
+        let t_ctx = Arc::new(task_ctx_with_extension(ctx, dist_feed_ctx));
+        let id = wuf.id();
+
+        for (partition, feed_idx) in (start_partition..end_partition).enumerate() {
+            let stream = wuf
+                .feed(feed_idx, Arc::clone(&t_ctx))?
+                .map(move |res| res.map(|wu| build_work_unit(&id, partition, wu)))
+                .boxed();
+            streams.push(stream);
+        }
+        Ok(TreeNodeRecursion::Continue)
+    })?;
+    Ok(streams)
+}

src/work_unit_feed/mod.rs

@@ -1,13 +1,18 @@
+mod coordinator;
 mod remote_work_unit_feed;
 mod work_unit;
 #[allow(clippy::module_inception)]
 mod work_unit_feed;
 mod work_unit_feed_provider;
 mod work_unit_feed_registry;
 
+pub use coordinator::collect_task_work_unit_feeds;
+pub use remote_work_unit_feed::{
+    RemoteWorkUnitFeedRegistry, RemoteWorkUnitFeedRxs, RemoteWorkUnitFeedTxs, WorkUnitRx,
+    WorkUnitTx, set_received_time, set_sent_time,
+};
 pub(crate) use remote_work_unit_feed::{
-    RemoteWorkUnitFeedRegistry, build_work_unit_batch_msg, set_work_unit_received_time,
-    set_work_unit_send_time,
+    build_work_unit_batch_msg, set_work_unit_received_time, set_work_unit_send_time,
 };
 pub(crate) use work_unit_feed_registry::{WorkUnitFeedRegistry, set_distributed_work_unit_feed};