Abstract away gRPC details from the worker protocol

[gabotechs]

Today, this project is deeply tied to gRPC and protobuf, as the worker protocol is autogenerated from the following protobuf definition:

• https://github.com/datafusion-contrib/datafusion-distributed/blob/main/src/worker/worker.proto

Typically, this is not a problem, as most projects are fine with just sticking with the classical gRPC+protobuf stack for inter-worker communication, but there are some cases where people would like to use different transportation mechanism.

Providing custom implementations for the prost generated stubs is not acceptable, as that still forces people into protobuf serde when they might be fine with either just in-memory comms or their own proprietary protocols.

The source of truth for all worker interactions is defined in the protobuf specification here:

service WorkerService {
  // Establishes a bidirectional message stream between a coordinator and a worker, over which messages
  // will be exchanged at any time during a query's lifetime. It's expected to be one coordinator channel
  // per task.
  rpc CoordinatorChannel(stream CoordinatorToWorkerMsg) returns (stream WorkerToCoordinatorMsg);
  // Executes the requested partition range of a subplan previously sent by the coordinator channel.
  rpc ExecuteTask(ExecuteTaskRequest) returns (stream FlightData);
  // Returns metadata about a worker. Currently only used for worker versioning.
  rpc GetWorkerInfo(GetWorkerInfoRequest) returns (GetWorkerInfoResponse);
}

And one possible way would be to build a new trait that abstracts away the WorkerServiceClient currently autogenerated by prost, and exposes an interface similar in shape to the one represented in today's worker.proto, but not coupled to gRPC or protobuf.

A structure that comes to mind can be to declare the protocol definition as a set of well documented traits in:

src/
  worker/
    protocol.rs

And specific implementations place somewhere else, like:

src/
  protocol_impl/
    grpc/
    ...

Ideally, that should allow hiding all gRPC-specific deps behind a feature flag, and even omit compiling them if people do not care about gRPC.

[gabotechs]

cc @stuhood

[mdashti]

@gabotechs Thanks for creating this issue.

Here's the design that we've come up with @stuhood:

Goals:

1. Make the transport a real extension point, so a transport other than Arrow Flight is a first-class implementation.
2. Make tonic / arrow-flight optional behind a Cargo feature, so the crate, its traits, and a test suite build with no gRPC at all.

Arrow Flight is the default transport, not the only one. The interface below suits any transport: an in-process or shared-memory mesh, a different RPC stack, or a custom message bus. The crate ships two: Arrow Flight (default, behind flight) and an in-memory transport (the reference implementation, and the default when flight is off).

What a transport does

datafusion-distributed (DF-D) is unopinionated about networking. A transport has four responsibilities; the abstraction makes them pluggable, with Flight as the default.

1. Dispatch (coordinator to worker): deliver each stage's plan to the workers that run it. Behind WorkerDispatch, built once per query by WorkerTransport::dispatcher(). It also carries the work-unit feed and, for Flight, the metrics back-channel.
2. Consume (consumer side): read a producer's output. Behind WorkerConnection, opened per stage by WorkerTransport::open().
3. Produce (worker side): run a task and emit its output partitions. Not a separate trait: produce is pull-driven by the consume side. Reading a partition runs the producer (a local task in-process, a streaming response over Flight).
4. Lifecycle: address workers (by URL for Flight, by task index for a co-located mesh), cancel, tear down. Cancellation is a per-execution token on the TaskContext.

So the umbrella trait is small: open (consume) and dispatcher (dispatch). Produce rides the read side; routing and cancellation are surfaced elsewhere (below).

The design

The transport traits

/// The umbrella. `open` is the consume/read side; `dispatcher` is the dispatch side.
pub trait WorkerTransport: Send + Sync {
    fn open(
        &self,
        input_stage: &RemoteStage,
        target_partitions: Range<usize>,
        target_task: usize,
        producer_head: ProducerHead,
        ctx: &Arc<TaskContext>,
        metrics: &ExecutionPlanMetricsSet,
    ) -> Result<Box<dyn WorkerConnection>>;

    /// Factory, called once per query so the dispatcher can hold per-query state (Flight shares one
    /// set of plan-send metrics and one query start timestamp across every stage).
    fn dispatcher(&self) -> Box<dyn WorkerDispatch>;
}

/// Consume: read a producer's output. The connection is opened per stage (closing over the stage)
/// and covers the whole `target_partitions` range; `execute` keys on the output partition. Produce
/// is pull-driven: reading a partition runs the producer.
pub trait WorkerConnection: Send + Sync {
    fn execute(&self, partition: usize) -> Result<BoxStream<'static, Result<RecordBatch>>>;
}

/// Dispatch: deliver a stage's plans to its workers, and carry the work-unit feed. Per-stage, not
/// per-task: the request carries the whole stage plus the query `join_set`, so delivery failures
/// propagate to the head. Flight resolves each worker URL and ships a `SetPlanRequest`; a co-located
/// transport ignores the URLs and routes by task index.
pub trait WorkerDispatch: Send + Sync {
    fn dispatch(&self, request: WorkerDispatchRequest<'_>) -> Result<()>;
}

#[non_exhaustive]
pub struct WorkerDispatchRequest<'a> {
    pub stage: &'a LocalStage,
    pub routed_urls: &'a [Url],
    pub task_ctx: &'a Arc<TaskContext>,
    pub metrics: &'a ExecutionPlanMetricsSet,
    pub metrics_store: Option<&'a Arc<MetricsStore>>,
    pub join_set: &'a mut JoinSet<Result<()>>,
}

WorkerConnection is the data-plane extension point. In a co-located transport, open returns a connection whose execute(partition) pulls batches straight from the local task registry (buffered through SingleWriteMultiRead), going through the same encode/decode, session building, work-unit feed, and metrics path as a remote read. Flight's execute opens a gRPC stream and the worker produces reactively. The asymmetry is lifecycle, not data flow: a Flight worker is a reactive server that runs a stage on request; a co-located worker has no request to react to and is driven eagerly by the read.

ProducerHead (passed to open) describes the upstream stage's output shape, None / BroadcastExec { output_partitions } / RepartitionExec { partitioning }, so the consume side knows how the producer partitions its output.

Surface the routing

The receive side of a shuffle reads partition slices with a fixed formula, global = P_c * consumer_task + local, internal to the network-boundary operators. Any transport that places a produced partition would otherwise re-derive it. Expose it on the boundary:

pub struct PartitionRoute { pub consumer_task: usize, pub consumer_partition: usize }

pub trait NetworkBoundary: ExecutionPlan {
    /// P_c, partitions per consumer task. Defaults to the boundary's output partition count.
    fn partitions_per_consumer_task(&self) -> usize { /* partition_count() */ }
    /// Producer output partition -> (consumer task, local partition): q -> (q / P_c, q % P_c).
    /// Fallible: a boundary whose consumers read whole per-producer-task groups
    /// (`NetworkCoalesceExec`) overrides it with an error instead of misrouting to task 0.
    fn route_partition(&self, output_partition: usize) -> Result<PartitionRoute>;
}

The operators' receive loops read partitions_per_consumer_task(), so there is one source of truth. Task-to-worker assignment stays the planner's (route_tasks, round-robin default, writing the URLs onto RemoteStage). The crate owns route_partition (fixed by the plan's partitioning); assignment is the pluggable part.

Cancellation

One CancellationToken per execution, on the TaskContext, not on the transport (Flight is shared process-wide) and not per connection (too granular). The crate creates it when the plan starts; dropping the head stream fires it; producers and consumers watch it through get_distributed_cancellation_token(ctx). The trait has no cancellation() method.

Work-unit feed

The work-unit feed streams runtime items (a file to read, a query to run) to each worker's leaf node as the query progresses, instead of fixing them at plan time. It is core, not Flight-specific: the dispatcher collects the per-task feeds and pumps units into the worker channels, so it rides dispatch in the same coordinator-to-worker direction as plan delivery. The in-memory transport feeds the same channels with flight off.

Decisions

1. Produce is pull-based, driven by the read side. There is no producer sink trait. The consumer's execute(partition) runs the producer: in-process it pulls from the local task registry (buffered through SingleWriteMultiRead); over Flight the worker streams reactively. One data-plane extension point (WorkerConnection) instead of two ends to keep in sync.
2. Dispatch is a per-query factory. dispatcher() mints a WorkerDispatch before the first stage; every stage reuses it. Flight needs this to share one set of plan-send metrics and one query start timestamp across stages, which a borrowed singleton could not hold.
3. Cancellation: one token per execution. On the TaskContext, not on the transport (shared process-wide) and not per connection (too granular). The crate creates it when the plan starts; dropping the head stream fires it; producers and consumers watch it.
4. Routing: the crate owns the route; assignment stays pluggable. route_partition is fixed by the plan's partitioning, so the crate always owns it, fallibly: a boundary that doesn't read the sliced layout errors instead of misrouting. Task-to-worker assignment is a real choice, so the crate ships a default (route_tasks plus round-robin) and lets a transport override it.
5. Work-unit feed is core and rides dispatch. The abstraction and its provider live in core; delivery goes through the transport's dispatch, so each transport carries the feed its own way (Flight over its bidi stream, the in-memory transport straight into the worker channels).

Transports shipped

• In-memory (InMemoryWorkerTransport, always built, the default when flight is off): every worker is the current process; plans delivered by a function call, partitions read from the local task registry, through the same encode/decode, session building, work-unit feed, and metrics path as a remote transport. It is the reference implementation for the extension points, and the integration suite runs against it.
• Arrow Flight (FlightWorkerTransport, default, behind flight): the reference remote transport.

An out-of-crate transport implements WorkerTransport (open + dispatcher) and registers it through DistributedExt::with_distributed_worker_transport. The crate re-exports the wire frames (crate::proto) and the helpers a transport needs (plan encode/decode, local-task execution, metrics collection).

And the implementation is in #508

We'd appreciate if you could review it. Thanks!

[mdashti]

#508 is meant to introduce the traits and their usages, with minimum code churn, and then it can be followed up with a refactor-only PR that moves the relevant code for each transport (e.g., flight-only) to its own sub-module.

[gabotechs]

Really cool! I think I'll be able to review this by the end of this week / beginning of the next one, thanks for all this work