Adaptive task count assignation

.github/workflows/ci.yml

@@ -41,19 +41,33 @@ jobs:
       - uses: ./.github/actions/setup
       - run: cargo test --features integration
 
-  tpch-test:
+  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 --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 --features tpch --test 'tpch_*'
+      - run: cargo test --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
@@ -62,6 +76,8 @@ jobs:
           path: testdata/tpcds/main.zip
           key: "main.zip"
       - run: cargo test --features tpcds --test tpcds_correctness_test shard${{ matrix.shard }}
+        env:
+          ADAPTIVE: ${{ matrix.planning_mode == 'adaptive' }}
 
   tpcds-plans-test:
     runs-on: ubuntu-latest
@@ -74,7 +90,24 @@ jobs:
           key: "main.zip"
       - run: cargo test --features tpcds --test tpcds_plans_test
 
-  clickbench-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: testdata/clickbench/
+          key: "data"
+      - run: cargo test --features clickbench --test clickbench_correctness_test
+        env:
+          ADAPTIVE: ${{ matrix.planning_mode == 'adaptive' }}
+
+  clickbench-plans-test:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
@@ -83,7 +116,7 @@ jobs:
         with:
           path: testdata/clickbench/
           key: "data"
-      - run: cargo test --features clickbench --test 'clickbench_*'
+      - run: cargo test --features clickbench --test clickbench_plans_test
 
   format-check:
     runs-on: ubuntu-latest

Cargo.toml

@@ -48,6 +48,7 @@ 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"
 

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),

src/common/mod.rs

@@ -5,6 +5,7 @@ mod recursion;
 mod task_context_helpers;
 mod time;
 mod uuid;
+mod vec;
 
 pub(crate) use children_helpers::require_one_child;
 pub(crate) use on_drop_stream::on_drop_stream;
@@ -13,3 +14,4 @@ pub(crate) 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(crate) use vec::{element_wise_sum, vec_avg_reduce, vec_cast, vec_div, vec_mul};

src/common/recursion.rs

@@ -589,6 +589,7 @@ mod tests {
                 query_id: uuid::Uuid::nil(),
                 num: 0,
                 workers: vec![],
+                runtime_stats: None,
             }))
             .unwrap()
     }

src/common/vec.rs

@@ -0,0 +1,80 @@
+use datafusion::common::internal_err;
+use datafusion::error::Result;
+use num_traits::AsPrimitive;
+use std::ops::{AddAssign, DivAssign, MulAssign};
+
+/// Converts a slice of type `I` to a `Vec<O>` using `as`-style primitive casting.
+pub(crate) fn vec_cast<I, O>(input: &[I]) -> Vec<O>
+where
+    I: AsPrimitive<O>,
+    O: Copy + 'static,
+{
+    input.iter().map(|v| v.as_()).collect()
+}
+
+/// Adds each element of `other` into the corresponding element of `one`, converting types via `AsPrimitive`.
+pub(crate) fn element_wise_sum<I, O>(mut one: Vec<I>, other: &[O]) -> Result<Vec<I>>
+where
+    I: AddAssign + Copy + 'static,
+    O: AsPrimitive<I> + 'static,
+{
+    if one.len() != other.len() {
+        return internal_err!("Cannot do an element wise sum of two vectors of different lengths");
+    }
+    for i in 0..one.len() {
+        one[i] += other[i].as_();
+    }
+    Ok(one)
+}
+
+/// Multiplies every element of `one` by the scalar `other`, converting types via `AsPrimitive`.
+pub(crate) fn vec_mul<I, O>(mut one: Vec<I>, other: O) -> Vec<I>
+where
+    I: MulAssign + Copy + 'static,
+    O: AsPrimitive<I> + 'static,
+{
+    for el in one.iter_mut() {
+        *el *= other.as_();
+    }
+    one
+}
+
+/// Divides every element of `one` by the scalar `other`, converting types via `AsPrimitive`.
+pub(crate) fn vec_div<I, O>(mut one: Vec<I>, other: O) -> Vec<I>
+where
+    I: DivAssign + Copy + 'static,
+    O: AsPrimitive<I> + 'static,
+{
+    for el in one.iter_mut() {
+        *el /= other.as_();
+    }
+    one
+}
+
+/// Reduces a collection of same-length `f32` vectors into a single vector by averaging element-wise.
+/// Empty inner vecs are skipped; returns an empty vec if all inputs are empty.
+pub(crate) fn vec_avg_reduce(vecs: Vec<Vec<f32>>) -> Result<Vec<f32>> {
+    let sample_count = vecs.len();
+    let mut iter = vecs.into_iter();
+    let mut acc = loop {
+        let Some(v) = iter.next() else {
+            return Ok(vec![]);
+        };
+        if !v.is_empty() {
+            break v;
+        }
+    };
+    for v in iter {
+        if v.is_empty() {
+            continue;
+        } else if acc.len() != v.len() {
+            return internal_err!(
+                "vec_avg_reduce: length mismatch — first vec has {} elements, got {}",
+                acc.len(),
+                v.len()
+            );
+        }
+        acc = element_wise_sum(acc, &v)?;
+    }
+    Ok(vec_div(acc, sample_count as f32))
+}

src/coordinator/distributed.rs

@@ -1,5 +1,7 @@
+use crate::DistributedConfig;
 use crate::common::{require_one_child, serialize_uuid};
 use crate::coordinator::metrics_store::MetricsStore;
+use crate::coordinator::prepare_dynamic_plan::prepare_dynamic_plan;
 use crate::coordinator::prepare_static_plan::prepare_static_plan;
 use crate::coordinator::query_coordinator::QueryCoordinator;
 use crate::distributed_planner::NetworkBoundaryExt;
@@ -198,7 +200,11 @@ impl ExecutionPlan for DistributedExec {
         builder.spawn(async move {
             let _guard = query_coordinator.end_query_guard();
 
-            let result = prepare_static_plan(&query_coordinator, &base_plan)?;
+            let d_cfg = DistributedConfig::from_config_options(context.session_config().options())?;
+            let result = match d_cfg.dynamic_task_count {
+                true => prepare_dynamic_plan(&query_coordinator, &base_plan).await?,
+                false => prepare_static_plan(&query_coordinator, &base_plan)?,
+            };
 
             plan_for_viz
                 .lock()

src/coordinator/mod.rs

@@ -1,6 +1,7 @@
 mod distributed;
 mod latency_metric;
 mod metrics_store;
+mod prepare_dynamic_plan;
 mod prepare_static_plan;
 mod query_coordinator;