Skip to content

Commit 9867814

Browse files
gabotechsgabotechs
authored
Drop velocity-based estimation in favor of driver path completion estimation (#552)
## Summary Replaces the **velocity-based** runtime-statistics estimate used by the dynamic task-count planner with a **driver-path completion** estimate. The sampler on top of each stage feeds statistics (final row/byte counts) back to the coordinator, which uses them to size the stage above. The old estimate extrapolated those final sizes from a sampled throughput multiplied by a fixed, assumed query duration, which systematically overestimated — often by orders of magnitude. The new estimate instead scales what a stage has already produced by how far it is through its *known* total work. ## How it worked before The `SamplerExec` reported, per partition: - `*_ready` — rows / bytes already materialized at sampling time, and - `*_per_second` — a sampled production **velocity** (rows/bytes per second). `gather_runtime_statistics` extrapolated the final size as: ```text total_rows = rows_ready + rows_per_second * ESTIMATED_QUERY_TIME_S // = 10 total_bytes = bytes_ready + bytes_per_second * ESTIMATED_QUERY_TIME_S ``` i.e. *"what's ready now, plus the sampled rate projected over 10 more seconds of execution."* While doing some benchmarks, I've seen that this approach heavily over-estimates row count. ## How the new estimate avoids it Instead of projecting a velocity over a made-up duration, it computes an actual **completion fraction** and divides by it: ```text pct_completed = rows_pulled_from_leaf / estimated_driver_path_leaf_rows total_rows = rows_ready / pct_completed total_bytes = bytes_ready / pct_completed ``` - `rows_pulled_from_leaf` — how many rows the stage has actually pulled from its leaf scans so far (now reported by the sampler in place of the velocity fields). - `estimated_driver_path_leaf_rows` — the total number of rows the stage must pull to finish, obtained by summing the leaf `num_rows` statistics along the **driver path**: the probe-side chain that actually streams to the output. Join build sides are excluded because they are one-off setup work, not output progress (see the new `apply_driver_path` helper). These counts come from real source statistics (e.g. parquet row counts). ## Caveats If a stage's leaves expose no row-count statistics, it falls back to a fixed `FALLBACK_ESTIMATED_PCT_SAMPLED = 0.2` (assume ~20 % sampled). ## Benchmarks TPCH-SF1: 1.02 faster TPCH-SF10: 1.04 faster TPCH-SF100: 1.01 faster TPC-DS-SF1: 1.01 slower ClickBench: 1.01 slower Performance is not impacted because both in `main` and this PR max out on tasks for a cluster of 12 workers. Regarding q-error on stats estimation, here are the results: TPCH-SF1: QERR P50: prev=1.80x, new=1.80x QERR P95: prev=6.75x, new=2.29x TPCH-SF10: QERR P50: prev=1.80x, new=1.45x QERR P95: prev=8.94x, new=2.12x TPCH-SF100: QERR P50: prev=2.12x, new=1.73x QERR P95: prev=9.68x, new=2.16x TPC-DS-SF1: QERR P50: prev=1.54x, new=1.45x QERR P95: prev=4.62x, new=4.01x ClickBench: QERR P50: prev=1.42x, new=1.45x QERR P95: prev=1.42x, new=1.45x
1 parent fb7b8c5 commit 9867814

13 files changed

Lines changed: 416 additions & 193 deletions

File tree

‎benchmarks/cdk/bin/datafusion-bench.ts‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ async function main() {
2525
.option('--repartition-file-min-size <number>', 'repartition_file_min_size DF option', '10485760' /* upstream default */)
2626
.option('--target-partitions <number>', 'target_partitions DF option', '8')
2727
.option('--dynamic <boolean>', 'Use the dynamic task count assigner', 'false')
28-
.option('--bytes-per-partition-per-second <number>', 'Target throughput in bytes per partition per second for the dynamic task count allocator', `${16 * 1024 * 1024}`)
28+
.option('--dynamic-bytes-per-partition <number>', 'Target throughput in bytes per partition per second for the dynamic task count allocator', `${16 * 1024 * 1024}`)
2929
.option('--queries <string>', 'Specific queries to run', undefined)
3030
.option('--debug <boolean>', 'Print the generated plans to stdout')
3131
.option('--warmup <boolean>', 'Perform a warmup query before the benchmarks', 'true')
@@ -49,7 +49,7 @@ async function main() {
4949
const broadcastJoins = options.broadcastJoins === 'true' || options.broadcastJoins === 1
5050
const partialReduce = options.partialReduce === 'true' || options.partialReduce === 1
5151
const dynamicTaskCount = options.dynamic === 'true' || options.dynamic === 1
52-
const bytesPerPartitionPerSecond = parseInt(options.bytesPerPartitionPerSecond)
52+
const dynamicBytesPerPartition = parseInt(options.dynamicBytesPerPartition)
5353
const debug = options.debug === true || options.debug === 'true' || options.debug === 1
5454
const warmup = options.warmup === true || options.warmup === 'true' || options.warmup === 1
5555

@@ -64,7 +64,7 @@ async function main() {
6464
broadcastJoins,
6565
partialReduce,
6666
dynamicTaskCount,
67-
bytesPerPartitionPerSecond,
67+
dynamicBytesPerPartition,
6868
maxTasksPerStage,
6969
repartitionFileMinSize,
7070
targetPartitions
@@ -107,7 +107,7 @@ class DataFusionRunner implements BenchmarkRunner {
107107
broadcastJoins: boolean;
108108
partialReduce: boolean;
109109
dynamicTaskCount: boolean;
110-
bytesPerPartitionPerSecond: number;
110+
dynamicBytesPerPartition: number;
111111
maxTasksPerStage: number;
112112
repartitionFileMinSize: number;
113113
targetPartitions: number;
@@ -195,7 +195,7 @@ class DataFusionRunner implements BenchmarkRunner {
195195
SET distributed.broadcast_joins=${this.options.broadcastJoins};
196196
SET distributed.partial_reduce=${this.options.partialReduce};
197197
SET distributed.dynamic_task_count=${this.options.dynamicTaskCount};
198-
SET distributed.bytes_per_partition_per_second=${this.options.bytesPerPartitionPerSecond};
198+
SET distributed.dynamic_bytes_per_partition=${this.options.dynamicBytesPerPartition};
199199
SET distributed.max_tasks_per_stage=${this.options.maxTasksPerStage};
200200
SET datafusion.optimizer.repartition_file_min_size=${this.options.repartitionFileMinSize};
201201
SET datafusion.execution.target_partitions=${this.options.targetPartitions};

‎src/common/recursion.rs‎

Lines changed: 215 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ use crate::{DistributedTaskContext, NetworkBoundaryExt};
33
use datafusion::common::Result;
44
use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeIterator, TreeNodeRecursion};
55
use datafusion::physical_plan::ExecutionPlan;
6+
use datafusion::physical_plan::joins::{CrossJoinExec, HashJoinExec, NestedLoopJoinExec};
67
use std::cell::RefCell;
78
use std::sync::Arc;
89

@@ -26,12 +27,41 @@ pub(crate) trait TreeNodeExt {
2627
f: F,
2728
) -> Result<TreeNodeRecursion>;
2829

30+
/// Applies `f` top-down (pre-order) to the nodes on the plan's *driver path*: the operators
31+
/// whose ongoing data production feeds, batch for batch, into the root's output. The
32+
/// [`TreeNodeRecursion`] returned by `f` steers it as in [`TreeNode::apply`].
33+
///
34+
/// It differs from a full traversal only at the pipeline-breaking joins ([`HashJoinExec`],
35+
/// [`NestedLoopJoinExec`], [`CrossJoinExec`]): it follows the probe side (`right`) and skips the
36+
/// build side (`left`), since a join fully materializes its build side before emitting any
37+
/// output row, so build-side rows are setup work rather than output progress. Every other
38+
/// operator is transparent.
39+
///
40+
/// Used to estimate a running stage's completion: the total rows to pull come from the leaves
41+
/// on the driver path (see `estimated_driver_path_leaf_rows`).
42+
///
43+
/// ```text
44+
/// ┌────────────┐
45+
/// │ HashJoin │ visited
46+
/// └─────┬──────┘
47+
/// build │ probe
48+
/// (left) │ │ (right)
49+
/// ┌────────┘ └────────┐
50+
/// ▼ ▼
51+
/// ┌────────┐ ┌────────┐
52+
/// │ Scan B │ SKIPPED │ Scan P │ visited
53+
/// └────────┘ └────────┘
54+
/// ```
55+
fn apply_driver_path<F: FnMut(&Self) -> Result<TreeNodeRecursion>>(
56+
&self,
57+
f: F,
58+
) -> Result<TreeNodeRecursion>;
59+
2960
/// Recursively rewrite the tree using `f` in a top-down (pre-order) fashion, propagating
3061
/// the appropriate [DistributedTaskContext] based on the presence of nodes that can isolate
3162
/// tasks, like [ChildrenIsolatorUnionExec].
3263
///
3364
/// `f` is applied to the node first, and then its children.
34-
#[allow(dead_code)] // Used in follow up work.
3565
fn transform_down_with_dt_ctx<
3666
F: FnMut(Self, DistributedTaskContext) -> Result<Transformed<Self>>,
3767
>(
@@ -62,7 +92,6 @@ pub(crate) trait TreeNodeExt {
6292
/// count.
6393
///
6494
/// `f` is applied to the node first, and then its children.
65-
#[allow(dead_code)] // Used in follow up work.
6695
fn transform_down_with_task_count<F: FnMut(Self, usize) -> Result<Transformed<Self>>>(
6796
self,
6897
task_count: usize,
@@ -104,6 +133,32 @@ impl TreeNodeExt for Arc<dyn ExecutionPlan> {
104133
recurse(self, ctx, &mut f)
105134
}
106135

136+
fn apply_driver_path<F: FnMut(&Self) -> Result<TreeNodeRecursion>>(
137+
&self,
138+
mut f: F,
139+
) -> Result<TreeNodeRecursion> {
140+
fn recurse<F: FnMut(&Arc<dyn ExecutionPlan>) -> Result<TreeNodeRecursion>>(
141+
plan: &Arc<dyn ExecutionPlan>,
142+
f: &mut F,
143+
) -> Result<TreeNodeRecursion> {
144+
f(plan)?.visit_children(|| {
145+
if let Some(hash_join) = plan.downcast_ref::<HashJoinExec>() {
146+
recurse(hash_join.right(), f)
147+
} else if let Some(nested_loop_join) = plan.downcast_ref::<NestedLoopJoinExec>() {
148+
recurse(nested_loop_join.right(), f)
149+
} else if let Some(cross_join) = plan.downcast_ref::<CrossJoinExec>() {
150+
recurse(cross_join.right(), f)
151+
} else {
152+
plan.children()
153+
.into_iter()
154+
.apply_until_stop(|child| recurse(child, f))
155+
}
156+
})
157+
}
158+
159+
recurse(self, &mut f)
160+
}
161+
107162
fn transform_down_with_dt_ctx<
108163
F: FnMut(Self, DistributedTaskContext) -> Result<Transformed<Self>>,
109164
>(
@@ -213,9 +268,13 @@ mod tests {
213268
use crate::execution_plans::ChildWeight;
214269
use crate::stage::RemoteStage;
215270
use crate::{NetworkCoalesceExec, Stage};
216-
use datafusion::arrow::datatypes::Schema;
271+
use datafusion::arrow::datatypes::{DataType, Field, Schema};
272+
use datafusion::common::{JoinType, NullEquality};
273+
use datafusion::physical_expr::PhysicalExpr;
274+
use datafusion::physical_expr::expressions::Column;
217275
use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
218276
use datafusion::physical_plan::empty::EmptyExec;
277+
use datafusion::physical_plan::joins::PartitionMode;
219278
use datafusion::physical_plan::union::UnionExec;
220279
use insta::assert_snapshot;
221280
// ── apply_with_dt_ctx ────────────────────────────────────────────────────────
@@ -348,6 +407,81 @@ mod tests {
348407
");
349408
}
350409

410+
// ── apply_driver_path ────────────────────────────────────────────────────
411+
412+
#[test]
413+
fn driver_path_leaf() {
414+
let plan = leaf();
415+
assert_snapshot!(trace_driver_path(&plan), @"Leaf");
416+
}
417+
418+
#[test]
419+
fn driver_path_top_down_order() {
420+
let plan = union(vec![single(leaf()), leaf()]);
421+
assert_snapshot!(trace_driver_path(&plan), @r"
422+
Union
423+
Single
424+
Leaf
425+
Leaf
426+
");
427+
}
428+
429+
#[test]
430+
fn driver_path_jump_skips_subtree() {
431+
let child = single(leaf());
432+
let plan = single(Arc::clone(&child));
433+
assert_snapshot!(
434+
trace_driver_path_with(&plan, |p| {
435+
if Arc::ptr_eq(p, &child) { TreeNodeRecursion::Jump } else { TreeNodeRecursion::Continue }
436+
}),
437+
@r"
438+
Single
439+
Single [->jump]
440+
");
441+
}
442+
443+
#[test]
444+
fn driver_path_hash_join_uses_probe_side() {
445+
let plan = hash_join(
446+
single(leaf_i32("l")),
447+
union(vec![leaf_i32("r"), leaf_i32("r")]),
448+
);
449+
assert_snapshot!(trace_driver_path(&plan), @r"
450+
HashJoin
451+
Union
452+
Leaf
453+
Leaf
454+
");
455+
}
456+
457+
#[test]
458+
fn driver_path_nested_loop_join_uses_probe_side() {
459+
let plan = nested_loop_join(
460+
single(leaf_i32("l")),
461+
union(vec![leaf_i32("r"), leaf_i32("r")]),
462+
);
463+
assert_snapshot!(trace_driver_path(&plan), @r"
464+
NestedLoopJoin
465+
Union
466+
Leaf
467+
Leaf
468+
");
469+
}
470+
471+
#[test]
472+
fn driver_path_cross_join_uses_probe_side() {
473+
let plan = cross_join(
474+
single(leaf_i32("l")),
475+
union(vec![leaf_i32("r"), leaf_i32("r")]),
476+
);
477+
assert_snapshot!(trace_driver_path(&plan), @r"
478+
CrossJoin
479+
Union
480+
Leaf
481+
Leaf
482+
");
483+
}
484+
351485
// ── transform_down_with_dt_ctx ────────────────────────────────────────────
352486

353487
#[test]
@@ -566,6 +700,14 @@ mod tests {
566700
Arc::new(EmptyExec::new(Arc::new(Schema::empty())))
567701
}
568702

703+
fn leaf_i32(name: &str) -> Arc<dyn ExecutionPlan> {
704+
Arc::new(EmptyExec::new(Arc::new(Schema::new(vec![Field::new(
705+
name,
706+
DataType::Int32,
707+
true,
708+
)]))))
709+
}
710+
569711
fn single(child: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
570712
Arc::new(CoalescePartitionsExec::new(child))
571713
}
@@ -581,6 +723,47 @@ mod tests {
581723
Arc::new(NetworkCoalesceExec::try_new(input, producer_tasks, 1).unwrap())
582724
}
583725

726+
fn hash_join(
727+
left: Arc<dyn ExecutionPlan>,
728+
right: Arc<dyn ExecutionPlan>,
729+
) -> Arc<dyn ExecutionPlan> {
730+
Arc::new(
731+
HashJoinExec::try_new(
732+
left,
733+
right,
734+
join_on(),
735+
None,
736+
&JoinType::Inner,
737+
None,
738+
PartitionMode::CollectLeft,
739+
NullEquality::NullEqualsNothing,
740+
false,
741+
)
742+
.unwrap(),
743+
)
744+
}
745+
746+
fn nested_loop_join(
747+
left: Arc<dyn ExecutionPlan>,
748+
right: Arc<dyn ExecutionPlan>,
749+
) -> Arc<dyn ExecutionPlan> {
750+
Arc::new(NestedLoopJoinExec::try_new(left, right, None, &JoinType::Inner, None).unwrap())
751+
}
752+
753+
fn cross_join(
754+
left: Arc<dyn ExecutionPlan>,
755+
right: Arc<dyn ExecutionPlan>,
756+
) -> Arc<dyn ExecutionPlan> {
757+
Arc::new(CrossJoinExec::new(left, right))
758+
}
759+
760+
fn join_on() -> Vec<(Arc<dyn PhysicalExpr>, Arc<dyn PhysicalExpr>)> {
761+
vec![(
762+
Arc::new(Column::new("l", 0)) as Arc<dyn PhysicalExpr>,
763+
Arc::new(Column::new("r", 0)) as Arc<dyn PhysicalExpr>,
764+
)]
765+
}
766+
584767
fn remote_network_boundary() -> Arc<dyn ExecutionPlan> {
585768
network_boundary(leaf(), 1)
586769
.as_network_boundary()
@@ -630,6 +813,12 @@ mod tests {
630813
"CIU"
631814
} else if p.is::<NetworkCoalesceExec>() {
632815
"Network"
816+
} else if p.is::<HashJoinExec>() {
817+
"HashJoin"
818+
} else if p.is::<NestedLoopJoinExec>() {
819+
"NestedLoopJoin"
820+
} else if p.is::<CrossJoinExec>() {
821+
"CrossJoin"
633822
} else {
634823
"?"
635824
}
@@ -664,6 +853,29 @@ mod tests {
664853
lines.join("\n")
665854
}
666855

856+
fn trace_driver_path(root: &Arc<dyn ExecutionPlan>) -> String {
857+
trace_driver_path_with(root, |_| TreeNodeRecursion::Continue)
858+
}
859+
860+
fn trace_driver_path_with<F: FnMut(&Arc<dyn ExecutionPlan>) -> TreeNodeRecursion>(
861+
root: &Arc<dyn ExecutionPlan>,
862+
mut decide: F,
863+
) -> String {
864+
let mut lines = vec![];
865+
root.apply_driver_path(|p| {
866+
let rec = decide(p);
867+
let suffix = match rec {
868+
TreeNodeRecursion::Continue => "",
869+
TreeNodeRecursion::Jump => " [->jump]",
870+
TreeNodeRecursion::Stop => " [->stop]",
871+
};
872+
lines.push(format!("{}{suffix}", plan_label(p)));
873+
Ok(rec)
874+
})
875+
.unwrap();
876+
lines.join("\n")
877+
}
878+
667879
fn trace_dt_ctx_down(root: Arc<dyn ExecutionPlan>, dt_ctx: DistributedTaskContext) -> String {
668880
trace_dt_ctx_down_with(root, dt_ctx, |_| TreeNodeRecursion::Continue)
669881
}

0 commit comments

Comments
 (0)