fix: let the shm transport absorb the coordinator's plan_proto encode

.github/workflows/ci.yml

@@ -41,27 +41,54 @@ jobs:
       - uses: ./.github/actions/setup
       - run: cargo test --features integration
 
-  tpch-test:
+  # Builds the lib with `grpc` off and runs the in-memory and shared-memory transport tests over it,
+  # the proof that the abstraction runs without gRPC. Scoped to those tests on purpose: the rest of
+  # the lib's unit tests load parquet testdata (git-LFS), which is the gRPC `unit-test` job's domain.
+  # The full lib still compiles with `grpc` off here, since the test binary builds every module.
+  unit-test-no-grpc:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
       - uses: ./.github/actions/setup
-      - run: cargo test --features tpch --test 'tpch_*'
+      - run: cargo test --no-default-features --lib -- protocol::in_process shm::in_process
+
+  tpch-correctness-test:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        planning_mode: [ "adaptive", "static" ]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: ./.github/actions/setup
+      - run: cargo test -p datafusion-distributed-benchmarks --features tpch --test tpch_correctness_test
+        env:
+          ADAPTIVE: ${{ matrix.planning_mode == 'adaptive' }}
+
+  tpch-plans-test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: ./.github/actions/setup
+      - run: cargo test -p datafusion-distributed-benchmarks --features tpch --test tpch_plans_test
 
   tpcds-correctness-test:
     runs-on: ubuntu-latest
     strategy:
       fail-fast: false
       matrix:
-        shard: ["01", "02", "03", "04", "05", "06", "07", "08", "09", "10"]
+        shard: [ "01", "02", "03", "04", "05", "06", "07", "08", "09", "10" ]
+        planning_mode: [ "adaptive", "static" ]
     steps:
       - uses: actions/checkout@v4
       - 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' }}
 
   tpcds-plans-test:
     runs-on: ubuntu-latest
@@ -70,20 +97,37 @@ 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
+    strategy:
+      fail-fast: false
+      matrix:
+        planning_mode: [ "adaptive", "static" ]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: ./.github/actions/setup
+      - uses: actions/cache@v4
+        with:
+          path: benchmarks/testdata/clickbench/
+          key: "data"
+      - run: cargo test -p datafusion-distributed-benchmarks --features clickbench --test clickbench_correctness_test
+        env:
+          ADAPTIVE: ${{ matrix.planning_mode == 'adaptive' }}
 
-  clickbench-test:
+  clickbench-plans-test:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
       - 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_plans_test
 
   format-check:
     runs-on: ubuntu-latest

.gitignore

@@ -8,7 +8,7 @@ testdata/tpcds/*
 !testdata/tpcds/README.md
 testdata/clickbench/*
 !testdata/clickbench/queries
-src/observability/gen/target/*
-src/observability/gen/Cargo.lock
-src/worker/gen/target/*
-src/worker/gen/Cargo.lock
+src/protocol/grpc/observability/gen/target/*
+src/protocol/grpc/observability/gen/Cargo.lock
+src/protocol/grpc/gen/target/*
+src/protocol/grpc/gen/Cargo.lock

Cargo.toml

@@ -23,13 +23,10 @@ datafusion = { workspace = true, features = [
   "datetime_expressions",
 ] }
 datafusion-proto = { workspace = true }
-arrow-flight = "58"
-arrow-select = "58"
-arrow-ipc = { version = "58", features = ["zstd"] }
 async-trait = "0.1.89"
+async-stream = "0.3"
+log = "0.4"
 tokio = { version = "1.48", features = ["full"] }
-tonic = { version = "0.14.1", features = ["transport"] }
-tower = "0.5.2"
 http = "1.3.1"
 itertools = "0.14.0"
 futures = "0.3.31"
@@ -48,29 +45,41 @@ moka = { version = "0.12", features = ["sync", "future"] }
 crossbeam-queue = "0.3"
 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"
 
+# grpc-specific features
+arrow-flight = { version = "58", optional = true }
+arrow-select = { version = "58", optional = true }
+arrow-ipc = { version = "58", features = ["zstd"], optional = true }
+tonic = { version = "0.14.1", features = ["transport"], optional = true }
+tower = { version = "0.5.2", optional = true }
+
 # integration_tests deps
 insta = { version = "1.46.0", features = ["filters"], optional = true }
 parquet = { version = "58", optional = true }
 arrow = { version = "58", optional = true, features = ["test_utils"] }
 hyper-util = { version = "0.1.16", optional = true }
 
 [features]
+default = ["grpc"]
 avro = ["datafusion/avro"]
-integration = ["insta", "parquet", "arrow", "hyper-util"]
+grpc = [
+    "arrow-flight",
+    "arrow-select",
+    "arrow-ipc",
+    "tonic",
+    "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"
@@ -80,6 +89,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"] }

benchmarks/cdk/bin/datafusion-bench.ts

@@ -24,6 +24,8 @@ async function main() {
         .option('--max-tasks-per-stage <number>', 'Max tasks per stage', '0')
         .option('--repartition-file-min-size <number>', 'repartition_file_min_size DF option', '10485760' /* upstream default */)
         .option('--target-partitions <number>', 'target_partitions DF option', '8')
+        .option('--dynamic <boolean>', 'Use the dynamic task count assigner', 'false')
+        .option('--bytes-per-partition-per-second <number>', 'Target throughput in bytes per partition per second for the dynamic task count allocator', `${16 * 1024 * 1024}`)
         .option('--queries <string>', 'Specific queries to run', undefined)
         .option('--debug <boolean>', 'Print the generated plans to stdout')
         .option('--warmup <boolean>', 'Perform a warmup query before the benchmarks', 'true')
@@ -46,6 +48,8 @@ async function main() {
     const childrenIsolatorUnions = options.childrenIsolatorUnions === 'true' || options.childrenIsolatorUnions === 1
     const broadcastJoins = options.broadcastJoins === 'true' || options.broadcastJoins === 1
     const partialReduce = options.partialReduce === 'true' || options.partialReduce === 1
+    const dynamicTaskCount = options.dynamic === 'true' || options.dynamic === 1
+    const bytesPerPartitionPerSecond = parseInt(options.bytesPerPartitionPerSecond)
     const debug = options.debug === true || options.debug === 'true' || options.debug === 1
     const warmup = options.warmup === true || options.warmup === 'true' || options.warmup === 1
 
@@ -59,6 +63,8 @@ async function main() {
         compression,
         broadcastJoins,
         partialReduce,
+        dynamicTaskCount,
+        bytesPerPartitionPerSecond,
         maxTasksPerStage,
         repartitionFileMinSize,
         targetPartitions
@@ -98,6 +104,8 @@ class DataFusionRunner implements BenchmarkRunner {
         childrenIsolatorUnions: boolean;
         broadcastJoins: boolean;
         partialReduce: boolean;
+        dynamicTaskCount: boolean;
+        bytesPerPartitionPerSecond: number;
         maxTasksPerStage: number;
         repartitionFileMinSize: number;
         targetPartitions: number;
@@ -177,6 +185,8 @@ class DataFusionRunner implements BenchmarkRunner {
       SET distributed.children_isolator_unions=${this.options.childrenIsolatorUnions};
       SET distributed.broadcast_joins=${this.options.broadcastJoins};
       SET distributed.partial_reduce=${this.options.partialReduce};
+      SET distributed.dynamic_task_count=${this.options.dynamicTaskCount};
+      SET distributed.bytes_per_partition_per_second=${this.options.bytesPerPartitionPerSecond};
       SET distributed.max_tasks_per_stage=${this.options.maxTasksPerStage};
       SET datafusion.optimizer.repartition_file_min_size=${this.options.repartitionFileMinSize};
       SET datafusion.execution.target_partitions=${this.options.targetPartitions};

benchmarks/src/run.rs

@@ -106,6 +106,10 @@ pub struct RunOpt {
     #[structopt(long, default_value = "0")]
     max_tasks_per_stage: usize,
 
+    /// Activate dynamic task count
+    #[structopt(long)]
+    dynamic: bool,
+
     /// Number of iterations of each test run
     #[structopt(short = "i", long = "iterations", default_value = "5")]
     iterations: usize,
@@ -203,6 +207,7 @@ impl RunOpt {
             .with_distributed_cardinality_effect_task_scale_factor(
                 self.cardinality_task_sf.unwrap_or(1.0),
             )?
+            .with_distributed_dynamic_task_count(self.dynamic)?
             .with_distributed_compression(match self.compression.as_str() {
                 "zstd" => Some(CompressionType::ZSTD),
                 "lz4" => Some(CompressionType::LZ4_FRAME),