feat: added a pluggable WorkerTransport, with Arrow Flight optional.

.github/workflows/ci.yml

@@ -32,6 +32,8 @@ jobs:
           components: clippy
       - run: cargo clippy --all-targets --all-features -- -D warnings
 
+  # The generic suite runs over the in-memory transport (what `start_localhost_context` builds by
+  # default), so it covers the distributed paths without a gRPC server.
   unit-test:
     runs-on: ubuntu-latest
     steps:
@@ -41,12 +43,50 @@ jobs:
       - uses: ./.github/actions/setup
       - run: cargo test --features integration
 
+  # The same suite over the Arrow-Flight gRPC transport. `start_localhost_context` switches to it
+  # when `DATAFUSION_DISTRIBUTED_TEST_TRANSPORT=flight`, so both transports get full coverage.
+  unit-test-flight-transport:
+    runs-on: ubuntu-latest
+    env:
+      DATAFUSION_DISTRIBUTED_TEST_TRANSPORT: flight
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          lfs: true
+      - uses: ./.github/actions/setup
+      - run: cargo test --features integration
+
+  # Builds the lib and the suite with `flight` off: no `tonic` / `arrow-flight`, in-memory transport
+  # as the default. The dataset suites moved to the benchmarks crate, so the suite no longer drags
+  # the benchmarks dev-dependency in (which re-enabled `flight` through feature unification).
+  # Examples are left out: they demo the Flight cluster, so `--all-features` clippy covers them.
+  no-flight:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: ./.github/actions/setup
+        with:
+          components: clippy
+      - run: cargo build --no-default-features --features integration
+      - run: cargo clippy --no-default-features --features integration --lib --tests -- -D warnings
+
+  # The integration suite over the in-memory transport with `flight` off, so the no-flight runtime
+  # path is run, not just built. Flight-specific tests gate themselves out.
+  unit-test-no-flight:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          lfs: true
+      - uses: ./.github/actions/setup
+      - run: cargo test --no-default-features --features integration --lib --tests
+
   tpch-test:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
       - uses: ./.github/actions/setup
-      - run: cargo test --features tpch --test 'tpch_*'
+      - run: cargo test -p datafusion-distributed-benchmarks --features tpch --test 'tpch_*'
 
   tpcds-correctness-test:
     runs-on: ubuntu-latest
@@ -59,9 +99,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 }}
 
   tpcds-plans-test:
     runs-on: ubuntu-latest
@@ -70,9 +110,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-test:
     runs-on: ubuntu-latest
@@ -81,9 +121,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_*'
+      - run: cargo test -p datafusion-distributed-benchmarks --features clickbench --test 'clickbench_*'
 
   format-check:
     runs-on: ubuntu-latest

Cargo.toml

@@ -23,13 +23,13 @@ datafusion = { workspace = true, features = [
   "datetime_expressions",
 ] }
 datafusion-proto = { workspace = true }
-arrow-flight = "58"
+arrow-flight = { version = "58", optional = true }
 arrow-select = "58"
 arrow-ipc = { version = "58", features = ["zstd"] }
 async-trait = "0.1.89"
 tokio = { version = "1.48", features = ["full"] }
-tonic = { version = "0.14.1", features = ["transport"] }
-tower = "0.5.2"
+tonic = { version = "0.14.1", features = ["transport"], optional = true }
+tower = { version = "0.5.2", optional = true }
 http = "1.3.1"
 itertools = "0.14.0"
 futures = "0.3.31"
@@ -49,7 +49,7 @@ crossbeam-queue = "0.3"
 sysinfo = { version = "0.30", optional = true }
 sketches-ddsketch = { version = "0.3", features = ["use_serde"] }
 bincode = "1"
-tonic-prost = "0.14.2"
+tonic-prost = { version = "0.14.2", optional = true }
 
 # integration_tests deps
 insta = { version = "1.46.0", features = ["filters"], optional = true }
@@ -58,19 +58,22 @@ arrow = { version = "58", optional = true, features = ["test_utils"] }
 hyper-util = { version = "0.1.16", optional = true }
 
 [features]
+default = ["flight"]
+# Arrow Flight gRPC transport. When off, the crate builds with no `tonic` / `arrow-flight`
+# and the in-memory transport becomes the default: distributed plans run inside the current
+# process. Multi-process execution then needs an embedder-registered transport.
+flight = ["dep:arrow-flight", "dep:tonic", "dep:tonic-prost", "dep:tower"]
 avro = ["datafusion/avro"]
+# Independent from `flight` so the integration suite runs against both the gRPC transport
+# (`--features integration`) and the in-memory one (`--no-default-features --features
+# integration`). The Flight-specific tests gate themselves on `flight`.
 integration = ["insta", "parquet", "arrow", "hyper-util"]
 
 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"

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 instead of the library because
+# they depend on it: as a dev-dependency of the library they re-enable `flight` on every test build
+# through feature unification, which makes the no-flight 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/common/cancellation.rs

@@ -0,0 +1,29 @@
+use datafusion::execution::TaskContext;
+use tokio_util::sync::CancellationToken;
+
+/// Per-execution cancellation token, attached to the [TaskContext] by `DistributedExec` when a
+/// plan starts. One token per execution: not on the transport (a shared, process-wide instance
+/// cannot own per-execution state) and not per connection (too granular).
+#[derive(Clone)]
+pub(crate) struct DistributedCancellationToken(pub(crate) CancellationToken);
+
+/// Returns the per-execution [CancellationToken] attached to `ctx`, or a fresh never-cancelled one
+/// if none is set (a context that did not come through `DistributedExec`). A transport's producers
+/// and consumers watch this instead of the transport carrying a `cancellation()` method.
+///
+/// The token does not replace transport-level stream close. Close propagates hop by hop and
+/// surfaces only when the next blocking operation fails; the token reaches every in-process
+/// participant the moment the head stream drops, including ones whose channels were never opened
+/// or never read.
+///
+/// Two caveats for watchers:
+/// - The token fires on any drop of the head stream, including after normal exhaustion. Treat it
+///   as teardown, not failure.
+/// - It lives in a session-config extension, so it is process-local: it does not cross plan
+///   serialization to remote workers. Out-of-process producers need their own teardown signal.
+pub fn get_distributed_cancellation_token(ctx: &TaskContext) -> CancellationToken {
+    ctx.session_config()
+        .get_extension::<DistributedCancellationToken>()
+        .map(|t| t.0.clone())
+        .unwrap_or_default()
+}

src/common/mod.rs

@@ -1,3 +1,4 @@
+mod cancellation;
 mod children_helpers;
 mod on_drop_stream;
 mod once_lock;
@@ -6,10 +7,12 @@ mod task_context_helpers;
 mod time;
 mod uuid;
 
+pub(crate) use cancellation::DistributedCancellationToken;
+pub use cancellation::get_distributed_cancellation_token;
 pub(crate) use children_helpers::require_one_child;
 pub(crate) use on_drop_stream::on_drop_stream;
 pub(crate) use once_lock::OnceLockResult;
-pub(crate) use recursion::TreeNodeExt;
+pub use recursion::TreeNodeExt;
 pub(crate) use task_context_helpers::task_ctx_with_extension;
 pub(crate) use time::now_ns;
-pub(crate) use uuid::{deserialize_uuid, serialize_uuid};
+pub use uuid::{deserialize_uuid, serialize_uuid};

src/common/recursion.rs

@@ -6,7 +6,7 @@ use datafusion::physical_plan::ExecutionPlan;
 use std::cell::RefCell;
 use std::sync::Arc;
 
-pub(crate) trait TreeNodeExt {
+pub trait TreeNodeExt {
     /// Applies `f` to the node then each of its children, recursively (a
     /// top-down, pre-order traversal), propagating the [DistributedTaskContext] correctly
     /// across nodes that mutate this context, and ignoring nodes that do not belong to

src/common/uuid.rs

@@ -2,10 +2,10 @@ use datafusion::common::Result;
 use datafusion_proto::protobuf::proto_error;
 use uuid::Uuid;
 
-pub(crate) fn deserialize_uuid(id: &[u8]) -> Result<Uuid> {
+pub fn deserialize_uuid(id: &[u8]) -> Result<Uuid> {
     Uuid::from_slice(id).map_err(|err| proto_error(format!("Invalid Uuid bytes: {err}")))
 }
 
-pub(crate) fn serialize_uuid(id: &Uuid) -> Vec<u8> {
+pub fn serialize_uuid(id: &Uuid) -> Vec<u8> {
     id.as_bytes().to_vec()
 }

src/config_extension_ext.rs

@@ -80,7 +80,7 @@ struct ConfigExtensionPropagationContext {
     prefixes: Vec<&'static str>,
 }
 
-pub(crate) fn get_config_extension_propagation_headers(
+pub fn get_config_extension_propagation_headers(
     cfg: &SessionConfig,
 ) -> Result<HeaderMap, DataFusionError> {
     fn parse_err(err: impl Error) -> DataFusionError {

src/coordinator/dispatch_metrics.rs

@@ -0,0 +1,33 @@
+use crate::common::now_ns;
+use crate::coordinator::latency_metric::LatencyMetric;
+use crate::{BytesCounterMetric, BytesMetricExt, DISTRIBUTED_DATAFUSION_TASK_ID_LABEL};
+use datafusion::physical_expr_common::metrics::{ExecutionPlanMetricsSet, Label, MetricBuilder};
+use std::sync::Arc;
+
+/// Metrics that measure network details about communications between [crate::DistributedExec] and a
+/// worker. Shared by every transport's dispatch path, so the plan-send byte/latency counters read
+/// the same regardless of how plans reach the workers.
+#[derive(Clone)]
+pub struct CoordinatorToWorkerMetrics {
+    pub plan_bytes_sent: BytesCounterMetric,
+    pub plan_send_latency: Arc<LatencyMetric>,
+    pub instantiation_time: u64,
+}
+
+impl CoordinatorToWorkerMetrics {
+    pub fn new(metrics: &ExecutionPlanMetricsSet) -> Self {
+        Self {
+            // Metric that measures to total sum of bytes worth of subplans sent.
+            plan_bytes_sent: MetricBuilder::new(metrics)
+                .with_label(Label::new(DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, "0"))
+                .bytes_counter("plan_bytes_sent"),
+            // Latency statistics about the network calls issued to the workers for feeding subplans.
+            plan_send_latency: Arc::new(LatencyMetric::new(
+                "plan_send_latency",
+                |b| b.with_label(Label::new(DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, "0")),
+                metrics,
+            )),
+            instantiation_time: now_ns(),
+        }
+    }
+}