Drop velocity-based estimation in favor of driver path completion estimation

Author: gabotechs

benchmarks/cdk/bin/datafusion-bench.ts

@@ -25,7 +25,7 @@ async function main() {
         .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('--dynamic-bytes-per-partition <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')
@@ -49,7 +49,7 @@ async function main() {
     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 dynamicBytesPerPartition = parseInt(options.dynamicBytesPerPartition)
     const debug = options.debug === true || options.debug === 'true' || options.debug === 1
     const warmup = options.warmup === true || options.warmup === 'true' || options.warmup === 1
 
@@ -64,7 +64,7 @@ async function main() {
         broadcastJoins,
         partialReduce,
         dynamicTaskCount,
-        bytesPerPartitionPerSecond,
+        dynamicBytesPerPartition,
         maxTasksPerStage,
         repartitionFileMinSize,
         targetPartitions
@@ -107,7 +107,7 @@ class DataFusionRunner implements BenchmarkRunner {
         broadcastJoins: boolean;
         partialReduce: boolean;
         dynamicTaskCount: boolean;
-        bytesPerPartitionPerSecond: number;
+        dynamicBytesPerPartition: number;
         maxTasksPerStage: number;
         repartitionFileMinSize: number;
         targetPartitions: number;
@@ -195,7 +195,7 @@ class DataFusionRunner implements BenchmarkRunner {
       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.dynamic_bytes_per_partition=${this.options.dynamicBytesPerPartition};
       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};

src/common/recursion.rs

@@ -3,6 +3,7 @@ use crate::{DistributedTaskContext, NetworkBoundaryExt};
 use datafusion::common::Result;
 use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeIterator, TreeNodeRecursion};
 use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_plan::joins::{CrossJoinExec, HashJoinExec, NestedLoopJoinExec};
 use std::cell::RefCell;
 use std::sync::Arc;
 
@@ -26,6 +27,36 @@ pub(crate) trait TreeNodeExt {
         f: F,
     ) -> Result<TreeNodeRecursion>;
 
+    /// Applies `f` top-down (pre-order) to the nodes on the plan's *driver path*: the operators
+    /// whose ongoing data production feeds, batch for batch, into the root's output. The
+    /// [`TreeNodeRecursion`] returned by `f` steers it as in [`TreeNode::apply`].
+    ///
+    /// It differs from a full traversal only at the pipeline-breaking joins ([`HashJoinExec`],
+    /// [`NestedLoopJoinExec`], [`CrossJoinExec`]): it follows the probe side (`right`) and skips the
+    /// build side (`left`), since a join fully materializes its build side before emitting any
+    /// output row, so build-side rows are setup work rather than output progress. Every other
+    /// operator is transparent.
+    ///
+    /// Used to estimate a running stage's completion: the total rows to pull come from the leaves
+    /// on the driver path (see `estimated_driver_path_leaf_rows`).
+    ///
+    /// ```text
+    ///        ┌────────────┐
+    ///        │ HashJoin   │  visited
+    ///        └─────┬──────┘
+    ///        build │ probe
+    ///      (left)  │  │ (right)
+    ///     ┌────────┘  └────────┐
+    ///     ▼                    ▼
+    /// ┌────────┐          ┌────────┐
+    /// │ Scan B │ SKIPPED  │ Scan P │  visited
+    /// └────────┘          └────────┘
+    /// ```
+    fn apply_driver_path<F: FnMut(&Self) -> Result<TreeNodeRecursion>>(
+        &self,
+        f: F,
+    ) -> Result<TreeNodeRecursion>;
+
     /// Recursively rewrite the tree using `f` in a top-down (pre-order) fashion, propagating
     /// the appropriate [DistributedTaskContext] based on the presence of nodes that can isolate
     /// tasks, like [ChildrenIsolatorUnionExec].
@@ -104,6 +135,32 @@ impl TreeNodeExt for Arc<dyn ExecutionPlan> {
         recurse(self, ctx, &mut f)
     }
 
+    fn apply_driver_path<F: FnMut(&Self) -> Result<TreeNodeRecursion>>(
+        &self,
+        mut f: F,
+    ) -> Result<TreeNodeRecursion> {
+        fn recurse<F: FnMut(&Arc<dyn ExecutionPlan>) -> Result<TreeNodeRecursion>>(
+            plan: &Arc<dyn ExecutionPlan>,
+            f: &mut F,
+        ) -> Result<TreeNodeRecursion> {
+            f(plan)?.visit_children(|| {
+                if let Some(hash_join) = plan.downcast_ref::<HashJoinExec>() {
+                    recurse(hash_join.right(), f)
+                } else if let Some(nested_loop_join) = plan.downcast_ref::<NestedLoopJoinExec>() {
+                    recurse(nested_loop_join.right(), f)
+                } else if let Some(cross_join) = plan.downcast_ref::<CrossJoinExec>() {
+                    recurse(cross_join.right(), f)
+                } else {
+                    plan.children()
+                        .into_iter()
+                        .apply_until_stop(|child| recurse(child, f))
+                }
+            })
+        }
+
+        recurse(self, &mut f)
+    }
+
     fn transform_down_with_dt_ctx<
         F: FnMut(Self, DistributedTaskContext) -> Result<Transformed<Self>>,
     >(
@@ -213,9 +270,13 @@ mod tests {
     use crate::execution_plans::ChildWeight;
     use crate::stage::RemoteStage;
     use crate::{NetworkCoalesceExec, Stage};
-    use datafusion::arrow::datatypes::Schema;
+    use datafusion::arrow::datatypes::{DataType, Field, Schema};
+    use datafusion::common::{JoinType, NullEquality};
+    use datafusion::physical_expr::PhysicalExpr;
+    use datafusion::physical_expr::expressions::Column;
     use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
     use datafusion::physical_plan::empty::EmptyExec;
+    use datafusion::physical_plan::joins::PartitionMode;
     use datafusion::physical_plan::union::UnionExec;
     use insta::assert_snapshot;
     // ── apply_with_dt_ctx ────────────────────────────────────────────────────────
@@ -348,6 +409,81 @@ mod tests {
         ");
     }
 
+    // ── apply_driver_path ────────────────────────────────────────────────────
+
+    #[test]
+    fn driver_path_leaf() {
+        let plan = leaf();
+        assert_snapshot!(trace_driver_path(&plan), @"Leaf");
+    }
+
+    #[test]
+    fn driver_path_top_down_order() {
+        let plan = union(vec![single(leaf()), leaf()]);
+        assert_snapshot!(trace_driver_path(&plan), @r"
+        Union
+        Single
+        Leaf
+        Leaf
+        ");
+    }
+
+    #[test]
+    fn driver_path_jump_skips_subtree() {
+        let child = single(leaf());
+        let plan = single(Arc::clone(&child));
+        assert_snapshot!(
+            trace_driver_path_with(&plan, |p| {
+                if Arc::ptr_eq(p, &child) { TreeNodeRecursion::Jump } else { TreeNodeRecursion::Continue }
+            }),
+            @r"
+        Single
+        Single [->jump]
+        ");
+    }
+
+    #[test]
+    fn driver_path_hash_join_uses_probe_side() {
+        let plan = hash_join(
+            single(leaf_i32("l")),
+            union(vec![leaf_i32("r"), leaf_i32("r")]),
+        );
+        assert_snapshot!(trace_driver_path(&plan), @r"
+        HashJoin
+        Union
+        Leaf
+        Leaf
+        ");
+    }
+
+    #[test]
+    fn driver_path_nested_loop_join_uses_probe_side() {
+        let plan = nested_loop_join(
+            single(leaf_i32("l")),
+            union(vec![leaf_i32("r"), leaf_i32("r")]),
+        );
+        assert_snapshot!(trace_driver_path(&plan), @r"
+        NestedLoopJoin
+        Union
+        Leaf
+        Leaf
+        ");
+    }
+
+    #[test]
+    fn driver_path_cross_join_uses_probe_side() {
+        let plan = cross_join(
+            single(leaf_i32("l")),
+            union(vec![leaf_i32("r"), leaf_i32("r")]),
+        );
+        assert_snapshot!(trace_driver_path(&plan), @r"
+        CrossJoin
+        Union
+        Leaf
+        Leaf
+        ");
+    }
+
     // ── transform_down_with_dt_ctx ────────────────────────────────────────────
 
     #[test]
@@ -566,6 +702,14 @@ mod tests {
         Arc::new(EmptyExec::new(Arc::new(Schema::empty())))
     }
 
+    fn leaf_i32(name: &str) -> Arc<dyn ExecutionPlan> {
+        Arc::new(EmptyExec::new(Arc::new(Schema::new(vec![Field::new(
+            name,
+            DataType::Int32,
+            true,
+        )]))))
+    }
+
     fn single(child: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
         Arc::new(CoalescePartitionsExec::new(child))
     }
@@ -581,6 +725,47 @@ mod tests {
         Arc::new(NetworkCoalesceExec::try_new(input, producer_tasks, 1).unwrap())
     }
 
+    fn hash_join(
+        left: Arc<dyn ExecutionPlan>,
+        right: Arc<dyn ExecutionPlan>,
+    ) -> Arc<dyn ExecutionPlan> {
+        Arc::new(
+            HashJoinExec::try_new(
+                left,
+                right,
+                join_on(),
+                None,
+                &JoinType::Inner,
+                None,
+                PartitionMode::CollectLeft,
+                NullEquality::NullEqualsNothing,
+                false,
+            )
+            .unwrap(),
+        )
+    }
+
+    fn nested_loop_join(
+        left: Arc<dyn ExecutionPlan>,
+        right: Arc<dyn ExecutionPlan>,
+    ) -> Arc<dyn ExecutionPlan> {
+        Arc::new(NestedLoopJoinExec::try_new(left, right, None, &JoinType::Inner, None).unwrap())
+    }
+
+    fn cross_join(
+        left: Arc<dyn ExecutionPlan>,
+        right: Arc<dyn ExecutionPlan>,
+    ) -> Arc<dyn ExecutionPlan> {
+        Arc::new(CrossJoinExec::new(left, right))
+    }
+
+    fn join_on() -> Vec<(Arc<dyn PhysicalExpr>, Arc<dyn PhysicalExpr>)> {
+        vec![(
+            Arc::new(Column::new("l", 0)) as Arc<dyn PhysicalExpr>,
+            Arc::new(Column::new("r", 0)) as Arc<dyn PhysicalExpr>,
+        )]
+    }
+
     fn remote_network_boundary() -> Arc<dyn ExecutionPlan> {
         network_boundary(leaf(), 1)
             .as_network_boundary()
@@ -630,6 +815,12 @@ mod tests {
             "CIU"
         } else if p.is::<NetworkCoalesceExec>() {
             "Network"
+        } else if p.is::<HashJoinExec>() {
+            "HashJoin"
+        } else if p.is::<NestedLoopJoinExec>() {
+            "NestedLoopJoin"
+        } else if p.is::<CrossJoinExec>() {
+            "CrossJoin"
         } else {
             "?"
         }
@@ -664,6 +855,29 @@ mod tests {
         lines.join("\n")
     }
 
+    fn trace_driver_path(root: &Arc<dyn ExecutionPlan>) -> String {
+        trace_driver_path_with(root, |_| TreeNodeRecursion::Continue)
+    }
+
+    fn trace_driver_path_with<F: FnMut(&Arc<dyn ExecutionPlan>) -> TreeNodeRecursion>(
+        root: &Arc<dyn ExecutionPlan>,
+        mut decide: F,
+    ) -> String {
+        let mut lines = vec![];
+        root.apply_driver_path(|p| {
+            let rec = decide(p);
+            let suffix = match rec {
+                TreeNodeRecursion::Continue => "",
+                TreeNodeRecursion::Jump => " [->jump]",
+                TreeNodeRecursion::Stop => " [->stop]",
+            };
+            lines.push(format!("{}{suffix}", plan_label(p)));
+            Ok(rec)
+        })
+        .unwrap();
+        lines.join("\n")
+    }
+
     fn trace_dt_ctx_down(root: Arc<dyn ExecutionPlan>, dt_ctx: DistributedTaskContext) -> String {
         trace_dt_ctx_down_with(root, dt_ctx, |_| TreeNodeRecursion::Continue)
     }