feat: add an in-process worker transport (no gRPC)

.github/workflows/ci.yml

@@ -41,6 +41,17 @@ jobs:
       - uses: ./.github/actions/setup
       - run: cargo test --features integration
 
+  # Runs the whole lib unit suite with `grpc` off, the proof that the abstraction builds and runs
+  # without gRPC. `lfs: true` because some planner unit tests load parquet testdata.
+  unit-test-no-grpc:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          lfs: true
+      - uses: ./.github/actions/setup
+      - run: cargo test --no-default-features --lib
+
   tpch-correctness-test:
     runs-on: ubuntu-latest
     strategy:
@@ -50,7 +61,7 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - uses: ./.github/actions/setup
-      - run: cargo test --features tpch --test tpch_correctness_test
+      - run: cargo test -p datafusion-distributed-benchmarks --features tpch --test tpch_correctness_test
         env:
           ADAPTIVE: ${{ matrix.planning_mode == 'adaptive' }}
 
@@ -59,7 +70,7 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - uses: ./.github/actions/setup
-      - run: cargo test --features tpch --test tpch_plans_test
+      - run: cargo test -p datafusion-distributed-benchmarks --features tpch --test tpch_plans_test
 
   tpcds-correctness-test:
     runs-on: ubuntu-latest
@@ -73,9 +84,9 @@ jobs:
       - uses: ./.github/actions/setup
       - uses: actions/cache@v4
         with:
-          path: testdata/tpcds/main.zip
+          path: benchmarks/testdata/tpcds/main.zip
           key: "main.zip"
-      - run: cargo test --features tpcds --test tpcds_correctness_test shard${{ matrix.shard }}
+      - run: cargo test -p datafusion-distributed-benchmarks --features tpcds --test tpcds_correctness_test shard${{ matrix.shard }}
         env:
           ADAPTIVE: ${{ matrix.planning_mode == 'adaptive' }}
 
@@ -86,9 +97,9 @@ jobs:
       - uses: ./.github/actions/setup
       - uses: actions/cache@v4
         with:
-          path: testdata/tpcds/main.zip
+          path: benchmarks/testdata/tpcds/main.zip
           key: "main.zip"
-      - run: cargo test --features tpcds --test tpcds_plans_test
+      - run: cargo test -p datafusion-distributed-benchmarks --features tpcds --test tpcds_plans_test
 
   clickbench-correctness-test:
     runs-on: ubuntu-latest
@@ -101,9 +112,9 @@ jobs:
       - uses: ./.github/actions/setup
       - uses: actions/cache@v4
         with:
-          path: testdata/clickbench/
+          path: benchmarks/testdata/clickbench/
           key: "data"
-      - run: cargo test --features clickbench --test clickbench_correctness_test
+      - run: cargo test -p datafusion-distributed-benchmarks --features clickbench --test clickbench_correctness_test
         env:
           ADAPTIVE: ${{ matrix.planning_mode == 'adaptive' }}
 
@@ -114,9 +125,9 @@ jobs:
       - uses: ./.github/actions/setup
       - uses: actions/cache@v4
         with:
-          path: testdata/clickbench/
+          path: benchmarks/testdata/clickbench/
           key: "data"
-      - run: cargo test --features clickbench --test clickbench_plans_test
+      - run: cargo test -p datafusion-distributed-benchmarks --features clickbench --test clickbench_plans_test
 
   format-check:
     runs-on: ubuntu-latest

Cargo.toml

@@ -45,7 +45,7 @@ sysinfo = { version = "0.30", optional = true }
 sketches-ddsketch = { version = "0.3", features = ["use_serde"] }
 num-traits = "0.2"
 bincode = "1"
-tonic-prost = "0.14.2"
+tonic-prost = { version = "0.14.2", optional = true }
 
 # grpc-specific features
 arrow-flight = { version = "58", optional = true }
@@ -68,21 +68,17 @@ grpc = [
     "arrow-select",
     "arrow-ipc",
     "tonic",
+    "tonic-prost",
     "tower"
 ]
 
 integration = ["insta", "parquet", "arrow", "hyper-util", "grpc"]
 
 system-metrics = ["sysinfo"]
 
-tpch = ["integration"]
-tpcds = ["integration"]
-clickbench = ["integration"]
-slow-tests = []
 sysinfo = ["dep:sysinfo"]
 
 [dev-dependencies]
-datafusion-distributed-benchmarks = { path = "benchmarks" }
 structopt = "0.3"
 insta = { version = "1.46.0", features = ["filters"] }
 parquet = "58"
@@ -92,6 +88,43 @@ hyper-util = "0.1.16"
 pretty_assertions = "1.4"
 test-case = "3.3.1"
 
+# Every example drives the gRPC transport, so a no-grpc build skips them.
+[[example]]
+name = "custom_distributed_partial_reduction_tree"
+required-features = ["integration"]
+
+[[example]]
+name = "custom_execution_plan"
+required-features = ["integration"]
+
+[[example]]
+name = "custom_worker_url_routing"
+required-features = ["integration"]
+
+[[example]]
+name = "work_unit_feed"
+required-features = ["integration"]
+
+[[example]]
+name = "in_memory_cluster"
+required-features = ["grpc"]
+
+[[example]]
+name = "localhost_run"
+required-features = ["grpc"]
+
+[[example]]
+name = "localhost_versioned_run"
+required-features = ["grpc"]
+
+[[example]]
+name = "localhost_worker"
+required-features = ["grpc"]
+
+[[example]]
+name = "localhost_versioned_worker"
+required-features = ["grpc"]
+
 [workspace.lints.clippy]
 disallowed_types = "deny"
 disallowed-methods = "deny"

benchmarks/Cargo.toml

@@ -33,9 +33,20 @@ aws-sdk-ec2 = "1"
 openssl = { version = "0.10", features = ["vendored"] } # Keep this. Necessary for the remote benchmarks worker.
 mimalloc = "0.1"
 
+[features]
+# Gates for the dataset test suites under `tests/`. They live here rather than in the library
+# because, as a dev-dependency of the library, this crate re-enables `grpc` on every test build
+# through feature unification, which made the no-grpc config untestable.
+tpch = []
+tpcds = []
+clickbench = []
+slow-tests = []
+
 [dev-dependencies]
 criterion = "0.5"
 sysinfo = "0.30"
+pretty_assertions = "1.4"
+test-case = "3.3.1"
 
 [build-dependencies]
 built = { version = "0.8", features = ["git2", "chrono"] }

tests/clickbench_correctness_test.rs → benchmarks/tests/clickbench_correctness_test.rs

@@ -1,4 +1,4 @@
-#[cfg(all(feature = "integration", feature = "clickbench", test))]
+#[cfg(all(feature = "clickbench", test))]
 mod tests {
     use datafusion::arrow::array::RecordBatch;
     use datafusion::common::plan_err;

tests/clickbench_plans_test.rs → benchmarks/tests/clickbench_plans_test.rs

@@ -1,4 +1,4 @@
-#[cfg(all(feature = "integration", feature = "clickbench", test))]
+#[cfg(all(feature = "clickbench", test))]
 mod tests {
     use datafusion::error::Result;
     use datafusion_distributed::test_utils::in_memory_channel_resolver::start_in_memory_context;

tests/stateful_data_cleanup.rs → benchmarks/tests/stateful_data_cleanup.rs

@@ -1,4 +1,4 @@
-#[cfg(all(feature = "integration", feature = "tpch", test))]
+#[cfg(all(feature = "tpch", test))]
 mod tests {
     use datafusion::common::instant::Instant;
     use datafusion::error::Result;

tests/tpcds_correctness_test.rs → benchmarks/tests/tpcds_correctness_test.rs

@@ -1,4 +1,4 @@
-#[cfg(all(feature = "integration", feature = "tpcds", test))]
+#[cfg(all(feature = "tpcds", test))]
 mod tests {
     use datafusion::arrow::array::RecordBatch;
     use datafusion::common::plan_err;

tests/tpcds_plans_test.rs → benchmarks/tests/tpcds_plans_test.rs

@@ -1,4 +1,4 @@
-#[cfg(all(feature = "integration", feature = "tpcds", test))]
+#[cfg(all(feature = "tpcds", test))]
 mod tests {
     use datafusion::error::Result;
     use datafusion_distributed::test_utils::in_memory_channel_resolver::start_in_memory_context;

tests/tpch_correctness_test.rs → benchmarks/tests/tpch_correctness_test.rs

@@ -1,4 +1,4 @@
-#[cfg(all(feature = "integration", feature = "tpch", test))]
+#[cfg(all(feature = "tpch", test))]
 mod tests {
     use datafusion::physical_plan::execute_stream;
     use datafusion::prelude::SessionContext;

tests/tpch_plans_test.rs → benchmarks/tests/tpch_plans_test.rs

@@ -1,4 +1,4 @@
-#[cfg(all(feature = "integration", feature = "tpch", test))]
+#[cfg(all(feature = "tpch", test))]
 mod tests {
     use datafusion_distributed::test_utils::in_memory_channel_resolver::start_in_memory_context;
     use datafusion_distributed::{

src/coordinator/metrics_store.rs

@@ -27,7 +27,7 @@ impl MetricsStore {
         self.rx.borrow().get(key).cloned()
     }
 
-    #[cfg(test)]
+    #[cfg(all(test, feature = "grpc"))]
     pub(crate) fn from_entries(entries: impl IntoIterator<Item = (TaskKey, TaskMetrics)>) -> Self {
         let map: HashMap<_, _> = entries.into_iter().collect();
         let (tx, rx) = watch::channel(map);

src/execution_plans/metrics.rs

@@ -21,7 +21,7 @@ impl MetricsWrapperExec {
         Self { inner, metrics }
     }
 
-    #[cfg(test)]
+    #[cfg(all(test, feature = "grpc"))]
     pub(crate) fn inner(&self) -> &Arc<dyn ExecutionPlan> {
         &self.inner
     }

src/execution_plans/mod.rs

@@ -8,7 +8,9 @@ mod network_coalesce;
 mod network_shuffle;
 mod sampler;
 
-#[cfg(any(test, feature = "integration"))]
+// The benchmark fixtures spin up real gRPC workers (`tonic` servers, Flight channels), so they
+// only exist when that transport is compiled in.
+#[cfg(all(feature = "grpc", any(test, feature = "integration")))]
 pub mod benchmarks;
 
 pub use broadcast::BroadcastExec;

src/lib.rs

@@ -39,14 +39,19 @@ pub mod test_utils;
 #[cfg(feature = "grpc")]
 pub use protocol::grpc;
 
+/// The worker-protocol prost message types, independent of any transport. A non-gRPC transport
+/// reaches for these to speak the same wire shape the gRPC path serializes. Unstable: this
+/// mirrors `worker.proto`, which is regenerated freely.
+pub use protocol::generated::worker as proto;
+
 pub use codec::DistributedCodec;
 pub use worker_resolver::{WorkerResolver, get_distributed_worker_resolver};
 
 pub use protocol::{
     ChannelResolver, CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest,
-    GetWorkerInfoResponse, LoadInfo, ProducerHeadSpec, SetPlanRequest, TaskKey, TaskMetrics,
-    WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, WorkerToCoordinatorMsg,
-    get_distributed_channel_resolver,
+    GetWorkerInfoResponse, InProcessChannelResolver, LoadInfo, ProducerHeadSpec, SetPlanRequest,
+    TaskKey, TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel,
+    WorkerToCoordinatorMsg, get_distributed_channel_resolver,
 };
 pub use stage::{
     DistributedTaskContext, Stage, display_plan_ascii, display_plan_graphviz, explain_analyze,
@@ -59,7 +64,7 @@ pub use worker::{
     Worker, WorkerQueryContext, WorkerSessionBuilder,
 };
 
-#[cfg(any(feature = "integration", test))]
+#[cfg(all(feature = "grpc", any(feature = "integration", test)))]
 pub use execution_plans::benchmarks::{
     LocalRepartitionBench, LocalRepartitionFixture, LocalRepartitionMode, ShuffleBench,
     ShuffleFixture, TransportBench, TransportBenchMode, TransportFixture,

src/metrics/task_metrics_collector.rs

@@ -19,7 +19,8 @@ pub fn collect_plan_metrics(plan: &Arc<dyn ExecutionPlan>) -> Result<Vec<Metrics
     Ok(metrics)
 }
 
-#[cfg(test)]
+// These tests execute over the in-memory gRPC transport, so they need that transport compiled in.
+#[cfg(all(test, feature = "grpc"))]
 mod tests {
 
     use super::*;

src/metrics/task_metrics_rewriter.rs

@@ -280,7 +280,8 @@ pub fn stage_metrics_rewriter(
         .map(|v| v.data)
 }
 
-#[cfg(test)]
+// These tests execute over the in-memory gRPC transport, so they need that transport compiled in.
+#[cfg(all(test, feature = "grpc"))]
 mod tests {
     use crate::DistributedExt;
     use crate::coordinator::MetricsStore;

src/protocol/channel_resolver.rs

@@ -79,6 +79,9 @@ pub fn get_distributed_channel_resolver(
 
     #[cfg(not(feature = "grpc"))]
     {
+        // With `grpc` off there is no built-in default. A co-located deployment can register
+        // [`crate::InProcessChannelResolver`] (or another transport) via
+        // `with_distributed_channel_resolver`.
         panic!(
             "gRPC feature is not enabled, and no channel resolver was provided, so no default ChannelResolver can be provided"
         );

src/protocol/generated/mod.rs

@@ -0,0 +1 @@
+pub mod worker;

src/protocol/grpc/generated/worker.rs → src/protocol/generated/worker.rs

@@ -468,6 +468,7 @@ pub struct MaxGauge {
     pub value: u64,
 }
 /// Generated client implementations.
+#[cfg(feature = "grpc")]
 pub mod worker_service_client {
     #![allow(
         unused_variables,
@@ -617,6 +618,7 @@ pub mod worker_service_client {
     }
 }
 /// Generated server implementations.
+#[cfg(feature = "grpc")]
 pub mod worker_service_server {
     #![allow(
         unused_variables,

src/protocol/grpc/gen/src/main.rs

@@ -6,7 +6,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     let proto_dir = repo_root.join("src/protocol/grpc");
     let proto_file = proto_dir.join("worker.proto");
-    let out_dir = repo_root.join("src/protocol/grpc/generated");
+    let out_dir = repo_root.join("src/protocol/generated");
 
     fs::create_dir_all(&out_dir)?;
 
@@ -18,6 +18,10 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     tonic_prost_build::configure()
         .build_server(true)
         .build_client(true)
+        // The generated messages build with `grpc` off; only the tonic client and server carry
+        // the feature gate. Emitted here so a regeneration cannot drop the gates.
+        .client_mod_attribute(".", "#[cfg(feature = \"grpc\")]")
+        .server_mod_attribute(".", "#[cfg(feature = \"grpc\")]")
         .out_dir(&out_dir)
         .extern_path(".worker.FlightData", "::arrow_flight::FlightData")
         .extern_path(

src/protocol/grpc/metrics_proto.rs

@@ -1,4 +1,4 @@
-use super::generated::worker as pb;
+use crate::protocol::generated::worker as pb;
 use chrono::DateTime;
 use datafusion::common::internal_err;
 use datafusion::error::DataFusionError;

src/protocol/grpc/mod.rs

@@ -1,6 +1,5 @@
 mod channel_resolver;
 mod errors;
-mod generated;
 mod metrics_proto;
 mod observability;
 mod on_drop_stream;

src/protocol/grpc/observability/gen/src/main.rs

@@ -21,7 +21,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         .out_dir(&out_dir)
         .extern_path(
             ".observability.TaskKey",
-            "crate::protocol::grpc::generated::worker::TaskKey",
+            "crate::protocol::generated::worker::TaskKey",
         )
         .compile_protos(&[proto_file], &[proto_dir])?;