Add dynamic stage built events - #570
Conversation
ca40839 to
30bbc84
Compare
fdb1562 to
4ccf26f
Compare
30bbc84 to
9ff5732
Compare
4ccf26f to
d8fdb19
Compare
6846880 to
a6ac3d4
Compare
b0f2fe8 to
700d9ca
Compare
… TaskEstimator (#564) ## Stack 1. #564 <- you are here 2. #567 3. #568 4. #569 5. #570 --- Introduces an event-driven API for reacting to distributed planning lifecycle events, and replaces the previous TaskEstimator methods with it. ## Motivation - We need to extend existing methods with new capabilities without imposing those changes on unrelated implementations. For example #569. - We need to add new event handlers without turning `TaskEstimator` into a generic bucket for every planner lifecycle hook. For example #570. - We could be expressing current behavior in terms of just event handler implementations, instead of inlining it in the code, for example: #567. - Converge into a single pattern all the user extension points, instead of having different shape of APIs for them. For example: #568 `TaskEstimator` was a convenient initial extension point, but its methods run at different points in planning, observe different context, and have different result semantics. Task-count estimation is evaluated while traversing the plan; leaf specialization only runs after a stage's final task count is known; routing runs later on the coordinator. Keeping these operations in one trait couples otherwise independent APIs. This event-specific pattern is easier to extend without breaking API consumers: a capability can evolve within the event that owns it instead of requiring every `TaskEstimator` implementation to adopt a broader trait. It is also easier to document and expose to users. Dedicated event types give us focused API pages and code documentation, and provide a quick index of the project’s extension points by the planner lifecycle events where users can hook their code. ## What changed Previously, one trait covered task-count estimation, leaf specialization, and task routing: ```rust trait TaskEstimator { fn task_estimation(&self, plan: &Arc<dyn ExecutionPlan>, cfg: &ConfigOptions) -> Option<TaskEstimation>; fn scale_up_leaf_node(&self, plan: &Arc<dyn ExecutionPlan>, task_count: usize, cfg: &ConfigOptions) -> Result<Option<Arc<dyn ExecutionPlan>>>; fn route_tasks(&self, ctx: &TaskRoutingContext<'_>) -> Result<Option<Vec<Url>>>; } ``` The old API registered one implementation containing every hook: ```rust SessionStateBuilder::new() .with_distributed_task_estimator(MyTaskEstimator); ``` Those phases now have dedicated event handlers, which can be ordinary functions and are registered independently: ```rust fn desired_task_count( _event: DesiredTaskCountEvent, ) -> Option<DesiredTaskCountEventResponse> { None } fn scale_up_leaf_node( _event: ScaleUpLeafNodeEvent, ) -> Option<Result<ScaleUpLeafNodeEventResponse>> { None } fn route_tasks( _event: RouteTasksEvent, ) -> Option<Result<RouteTasksEventResponse>> { None } SessionStateBuilder::new() .with_distributed_desired_task_count_handler(desired_task_count) .with_distributed_scale_up_leaf_node_handler(scale_up_leaf_node) .with_distributed_route_tasks_handler(route_tasks); ```
a6ac3d4 to
a5e0b27
Compare
700d9ca to
7fbb1b4
Compare
## Stack 1. #564 2. #567 <- you are here 3. #568 4. #569 5. #570 --- Moves routing policy into ordered `RouteTasksHandler`s. Built-ins now cover child co-location, coordinator co-location, and randomized fallback, while custom handlers can override them. ```rust config.with_distributed_route_tasks_handler(|event| { Some(Ok(RouteTasksEventResponse::new( vec![worker_url.clone(); event.task_count], ))) }); ``` --------- Co-authored-by: Jayant Shrivastava <jshrivastava03@gmail.com>
## Stack 1. #564 2. #567 3. #568 <- you are here 4. #569 5. #570 --- Replaces the worker-side `OnPlanHook` with composable `WorkerPlanRewriteHandler`s. Register handlers in each worker's session-state builder; coordinator registration is not propagated to workers. ```rust ctx.builder .with_distributed_worker_plan_rewrite_handler(|event| { Ok(WorkerPlanRewriteEventResponse::new(event.plan)) }) .build() ```
465e8a0 to
b4a96e5
Compare
## Stack 1. #564 2. #567 3. #568 4. #569 <- you are here 5. #570 --- Makes only desired task-count handlers asynchronous, so implementations can await metadata before returning a hint. The event still borrows `SessionConfig`, avoiding configuration copies. ```rust #[async_trait] impl DesiredTaskCountHandler for MyHandler { async fn handle(&self, event: DesiredTaskCountEvent<'_>) -> Option<Result<DesiredTaskCountEventResponse>> { Some(Ok(DesiredTaskCountEventResponse::desired(8))) } } ```
7fbb1b4 to
1ddbdb1
Compare
|
benchmarks run tpch/sf10 --instance-type c5n.2xlarge --nodes 1 |
|
benchmarks run tpch/sf1 |
1 similar comment
|
benchmarks run tpch/sf1 |
|
Queued benchmark job 1: |
|
Running |
|
Benchmark job 1 failed for |
|
benchmarks run tpch/sf1 --nodes 1 |
|
Requested by this comment. Run metadata
Queue: 1s · Dataset validation: 0s · Total: 4m 36s Base: === Comparing tpch/sf1 results 'datafusion-distributed-ca19679c6928' [prev] with 'datafusion-distributed-1ddbdb1378de' [new] === TOTAL: prev=10080 ms, new=9213 ms, diff=1.09 faster ✔ Show full query output q1: prev= 384 ms, new= 342 ms, diff=1.12 faster ✔
q2: prev= 885 ms, new= 625 ms, diff=1.42 faster ✔
q3: prev= 663 ms, new= 425 ms, diff=1.56 faster ✅
q4: prev= 273 ms, new= 229 ms, diff=1.19 faster ✔
q5: prev= 754 ms, new= 655 ms, diff=1.15 faster ✔
q6: prev= 178 ms, new= 206 ms, diff=1.16 slower ✖
q7: prev= 545 ms, new= 612 ms, diff=1.12 slower ✖
q8: prev= 787 ms, new= 813 ms, diff=1.03 slower ✖
q9: prev= 717 ms, new= 669 ms, diff=1.07 faster ✔
q10: prev= 422 ms, new= 417 ms, diff=1.01 faster ✔
q11: prev= 335 ms, new= 307 ms, diff=1.09 faster ✔
q12: prev= 257 ms, new= 239 ms, diff=1.08 faster ✔
q13: prev= 273 ms, new= 259 ms, diff=1.05 faster ✔
q14: prev= 265 ms, new= 227 ms, diff=1.17 faster ✔
q15: prev= 434 ms, new= 377 ms, diff=1.15 faster ✔
q16: prev= 169 ms, new= 174 ms, diff=1.03 slower ✖
q17: prev= 607 ms, new= 507 ms, diff=1.20 faster ✔
q18: prev= 515 ms, new= 481 ms, diff=1.07 faster ✔
q19: prev= 324 ms, new= 319 ms, diff=1.02 faster ✔
q20: prev= 420 ms, new= 413 ms, diff=1.02 faster ✔
q21: prev= 682 ms, new= 688 ms, diff=1.01 slower ✖
q22: prev= 191 ms, new= 229 ms, diff=1.20 slower ✖
|
|
benchmarks run tpch/sf100 |
|
Requested by this comment. Run metadata
Queue: 1s · Dataset validation: 0s · Total: 17m 18s Base: === Comparing tpch/sf100 results 'datafusion-distributed-ca19679c6928' [prev] with 'datafusion-distributed-1ddbdb1378de' [new] === TOTAL: prev=73162 ms, new=67858 ms, diff=1.08 faster ✔ Show full query output q1: prev=2611 ms, new=2538 ms, diff=1.03 faster ✔
q2: prev=1720 ms, new=1557 ms, diff=1.10 faster ✔
q3: prev=2616 ms, new=2637 ms, diff=1.01 slower ✖
q4: prev=1197 ms, new=1177 ms, diff=1.02 faster ✔
q5: prev=3921 ms, new=4110 ms, diff=1.05 slower ✖
q6: prev=1081 ms, new=1079 ms, diff=1.00 faster ✔
q7: prev=4660 ms, new=4686 ms, diff=1.01 slower ✖
q8: prev=4632 ms, new=4772 ms, diff=1.03 slower ✖
q9: prev=11604 ms, new=6305 ms, diff=1.84 faster ✅
q10: prev=6277 ms, new=5818 ms, diff=1.08 faster ✔
q11: prev=1132 ms, new=1097 ms, diff=1.03 faster ✔
q12: prev=1929 ms, new=1743 ms, diff=1.11 faster ✔
q13: prev=1827 ms, new=1812 ms, diff=1.01 faster ✔
q14: prev=1588 ms, new=1635 ms, diff=1.03 slower ✖
q15: prev=3223 ms, new=3249 ms, diff=1.01 slower ✖
q16: prev= 812 ms, new= 813 ms, diff=1.00 slower ✖
q17: prev=4782 ms, new=4825 ms, diff=1.01 slower ✖
q18: prev=5701 ms, new=5638 ms, diff=1.01 faster ✔
q19: prev=1601 ms, new=1763 ms, diff=1.10 slower ✖
q20: prev=2374 ms, new=2442 ms, diff=1.03 slower ✖
q21: prev=6918 ms, new=7112 ms, diff=1.03 slower ✖
q22: prev= 956 ms, new=1050 ms, diff=1.10 slower ✖
|
|
@codex review |
shinzoxD
left a comment
There was a problem hiding this comment.
The rewrite path needs one more invariant/coverage pass before this public hook is safe to use.
A handler may return a semantics-preserving wrapper around the original stage plan, which is allowed by the documented “optimize it” contract. After the assignment at prepare_dynamic_plan.rs:80, however, ProducerHead::insert_sampler only inserts a sampler when the replacement root itself is a RepartitionExec or BroadcastExec. A wrapper therefore leaves the stage without any SamplerExec, and the later gather_runtime_statistics call fails with Missing SamplerExec for shuffle/broadcast stages.
Please either validate and clearly enforce the required stage-head/topology contract on DynamicStageBuiltEventResponse, or make sampler insertion robust to supported rewrites. This new public behavior should also have an integration test that installs the handler and exercises an actual rewrite (plus the decline/error paths), since the current suite never invokes the new hook.
|
benchmarks run tpch/sf1 --nodes 1 |
|
Requested by this comment. Run metadata
Queue: 1s · Dataset validation: 0s · Total: 4m 41s Base: === Comparing tpch/sf1 results 'datafusion-distributed-ca19679c6928' [prev] with 'datafusion-distributed-1ddbdb1378de' [new] === TOTAL: prev=11001 ms, new=9954 ms, diff=1.11 faster ✅ Show full query output q1: prev= 342 ms, new= 353 ms, diff=1.03 slower ✖
q2: prev= 709 ms, new= 637 ms, diff=1.11 faster ✔
q3: prev= 570 ms, new= 392 ms, diff=1.45 faster ✔
q4: prev= 253 ms, new= 228 ms, diff=1.11 faster ✔
q5: prev= 695 ms, new= 605 ms, diff=1.15 faster ✔
q6: prev= 347 ms, new= 180 ms, diff=1.93 faster ✅
q7: prev= 746 ms, new= 603 ms, diff=1.24 faster ✔
q8: prev=1159 ms, new= 873 ms, diff=1.33 faster ✔
q9: prev= 907 ms, new= 732 ms, diff=1.24 faster ✔
q10: prev= 519 ms, new= 534 ms, diff=1.03 slower ✖
q11: prev= 313 ms, new= 316 ms, diff=1.01 slower ✖
q12: prev= 262 ms, new= 292 ms, diff=1.11 slower ✖
q13: prev= 269 ms, new= 298 ms, diff=1.11 slower ✖
q14: prev= 264 ms, new= 281 ms, diff=1.06 slower ✖
q15: prev= 400 ms, new= 408 ms, diff=1.02 slower ✖
q16: prev= 188 ms, new= 203 ms, diff=1.08 slower ✖
q17: prev= 616 ms, new= 607 ms, diff=1.01 faster ✔
q18: prev= 599 ms, new= 608 ms, diff=1.02 slower ✖
q19: prev= 385 ms, new= 368 ms, diff=1.05 faster ✔
q20: prev= 455 ms, new= 447 ms, diff=1.02 faster ✔
q21: prev= 771 ms, new= 775 ms, diff=1.01 slower ✖
q22: prev= 232 ms, new= 214 ms, diff=1.08 faster ✔
|
gabotechs
left a comment
There was a problem hiding this comment.
Test review from Codex — no action needed.
| /// Estimated cost associated to the fragment of the plan. Calculated based on the amount of data | ||
| /// that will flow through each node and the static cpu, memory and network complexity of the | ||
| /// different nodes based on the algorithms that will run inside them. | ||
| #[derive(Default, Debug, Clone, Copy)] |
There was a problem hiding this comment.
Test inline comment from Codex — no action needed.
|
benchmarks run tpch/sf1 --nodes 1 |
|
Requested by this comment. Benchmark job 28 failed for |
|
benchmarks run tpch/sf1 --nodes 1 |
|
Requested by this comment. Benchmark job 29 failed for |
|
benchmarks run tpch/sf1 --nodes 1 |
|
Requested by this comment. Benchmark job 30 failed for |
|
benchmarks run tpch/sf1 --nodes 1 |
|
Requested by this comment. Run metadata
Queue: 1s · Dataset validation: 0s · Total: 4m 16s Base: === Comparing tpch/sf1 results 'datafusion-benchmark-base' [prev] with 'datafusion-benchmark-head' [new] === TOTAL: prev=9726 ms, new=9731 ms, diff=1.00 slower ✖ Show full query output q1: prev= 320 ms, new= 302 ms, diff=1.06 faster ✔
q2: prev= 628 ms, new= 595 ms, diff=1.06 faster ✔
q3: prev= 442 ms, new= 432 ms, diff=1.02 faster ✔
q4: prev= 268 ms, new= 287 ms, diff=1.07 slower ✖
q5: prev= 697 ms, new= 614 ms, diff=1.14 faster ✔
q6: prev= 222 ms, new= 171 ms, diff=1.30 faster ✔
q7: prev= 669 ms, new= 573 ms, diff=1.17 faster ✔
q8: prev= 865 ms, new= 869 ms, diff=1.00 slower ✖
q9: prev= 701 ms, new= 689 ms, diff=1.02 faster ✔
q10: prev= 472 ms, new= 449 ms, diff=1.05 faster ✔
q11: prev= 321 ms, new= 339 ms, diff=1.06 slower ✖
q12: prev= 250 ms, new= 327 ms, diff=1.31 slower ✖
q13: prev= 271 ms, new= 366 ms, diff=1.35 slower ✖
q14: prev= 243 ms, new= 299 ms, diff=1.23 slower ✖
q15: prev= 388 ms, new= 492 ms, diff=1.27 slower ✖
q16: prev= 191 ms, new= 196 ms, diff=1.03 slower ✖
q17: prev= 548 ms, new= 505 ms, diff=1.09 faster ✔
q18: prev= 546 ms, new= 506 ms, diff=1.08 faster ✔
q19: prev= 342 ms, new= 326 ms, diff=1.05 faster ✔
q20: prev= 451 ms, new= 425 ms, diff=1.06 faster ✔
q21: prev= 718 ms, new= 806 ms, diff=1.12 slower ✖
q22: prev= 173 ms, new= 163 ms, diff=1.06 faster ✔
|
|
benchmarks run tpch/sf1 --nodes 1 |
|
Requested by this comment. Run metadata
Queue: 1s · Dataset validation: 0s · Total: 4m 46s Capacity: 1 === Comparing tpch/sf1 results 'datafusion-benchmark-base' [prev] with 'datafusion-benchmark-head' [new] === TOTAL: prev=9295 ms, new=9025 ms, diff=1.03 faster ✔ Show full query output q1: prev= 316 ms, new= 291 ms, diff=1.09 faster ✔
q2: prev= 574 ms, new= 543 ms, diff=1.06 faster ✔
q3: prev= 387 ms, new= 360 ms, diff=1.07 faster ✔
q4: prev= 233 ms, new= 219 ms, diff=1.06 faster ✔
q5: prev= 575 ms, new= 576 ms, diff=1.00 slower ✖
q6: prev= 184 ms, new= 153 ms, diff=1.20 faster ✔
q7: prev= 575 ms, new= 564 ms, diff=1.02 faster ✔
q8: prev= 815 ms, new= 864 ms, diff=1.06 slower ✖
q9: prev= 690 ms, new= 687 ms, diff=1.00 faster ✔
q10: prev= 470 ms, new= 425 ms, diff=1.11 faster ✔
q11: prev= 301 ms, new= 307 ms, diff=1.02 slower ✖
q12: prev= 265 ms, new= 252 ms, diff=1.05 faster ✔
q13: prev= 272 ms, new= 246 ms, diff=1.11 faster ✔
q14: prev= 253 ms, new= 228 ms, diff=1.11 faster ✔
q15: prev= 460 ms, new= 405 ms, diff=1.14 faster ✔
q16: prev= 188 ms, new= 180 ms, diff=1.04 faster ✔
q17: prev= 505 ms, new= 513 ms, diff=1.02 slower ✖
q18: prev= 547 ms, new= 537 ms, diff=1.02 faster ✔
q19: prev= 341 ms, new= 327 ms, diff=1.04 faster ✔
q20: prev= 432 ms, new= 427 ms, diff=1.01 faster ✔
q21: prev= 715 ms, new= 705 ms, diff=1.01 faster ✔
q22: prev= 197 ms, new= 216 ms, diff=1.10 slower ✖
|
1ddbdb1 to
4ffe219
Compare
|
benchmarks run tpch/sf1 |
|
Requested by this comment. Run metadata
Queue: 1s · Dataset validation: 0s · Total: 9m 57s Capacity: 12 === Comparing tpch/sf1 results 'datafusion-benchmark-base' [prev] with 'datafusion-benchmark-head' [new] === TOTAL: prev=7825 ms, new=8148 ms, diff=1.04 slower ✖ Show full query output q1: prev= 387 ms, new= 303 ms, diff=1.28 faster ✔
q2: prev= 563 ms, new= 596 ms, diff=1.06 slower ✖
q3: prev= 387 ms, new= 554 ms, diff=1.43 slower ✖
q4: prev= 180 ms, new= 215 ms, diff=1.19 slower ✖
q5: prev= 412 ms, new= 502 ms, diff=1.22 slower ✖
q6: prev= 100 ms, new= 103 ms, diff=1.03 slower ✖
q7: prev= 429 ms, new= 482 ms, diff=1.12 slower ✖
q8: prev= 632 ms, new= 613 ms, diff=1.03 faster ✔
q9: prev= 540 ms, new= 563 ms, diff=1.04 slower ✖
q10: prev= 401 ms, new= 406 ms, diff=1.01 slower ✖
q11: prev= 307 ms, new= 324 ms, diff=1.06 slower ✖
q12: prev= 222 ms, new= 203 ms, diff=1.09 faster ✔
q13: prev= 258 ms, new= 221 ms, diff=1.17 faster ✔
q14: prev= 184 ms, new= 202 ms, diff=1.10 slower ✖
q15: prev= 240 ms, new= 210 ms, diff=1.14 faster ✔
q16: prev= 195 ms, new= 171 ms, diff=1.14 faster ✔
q17: prev= 351 ms, new= 394 ms, diff=1.12 slower ✖
q18: prev= 457 ms, new= 439 ms, diff=1.04 faster ✔
q19: prev= 255 ms, new= 230 ms, diff=1.11 faster ✔
q20: prev= 480 ms, new= 546 ms, diff=1.14 slower ✖
q21: prev= 636 ms, new= 655 ms, diff=1.03 slower ✖
q22: prev= 209 ms, new= 216 ms, diff=1.03 slower ✖
|
|
benchmarks run tpch/sf1 |
|
Requested by this comment. Benchmark resultsCompared: Base === Comparing tpch/sf1 results 'datafusion-benchmark-base' [prev] with 'datafusion-benchmark-head' [new] === TOTAL: prev=8169 ms, new=8194 ms, diff=1.00 slower ✖ Show full query output q1: prev= 413 ms, new= 265 ms, diff=1.56 faster ✅
q2: prev= 563 ms, new= 639 ms, diff=1.13 slower ✖
q3: prev= 318 ms, new= 362 ms, diff=1.14 slower ✖
q4: prev= 319 ms, new= 393 ms, diff=1.23 slower ✖
q5: prev= 447 ms, new= 478 ms, diff=1.07 slower ✖
q6: prev= 137 ms, new= 96 ms, diff=1.43 faster ✔
q7: prev= 600 ms, new= 563 ms, diff=1.07 faster ✔
q8: prev= 626 ms, new= 608 ms, diff=1.03 faster ✔
q9: prev= 574 ms, new= 632 ms, diff=1.10 slower ✖
q10: prev= 407 ms, new= 398 ms, diff=1.02 faster ✔
q11: prev= 314 ms, new= 308 ms, diff=1.02 faster ✔
q12: prev= 201 ms, new= 184 ms, diff=1.09 faster ✔
q13: prev= 261 ms, new= 240 ms, diff=1.09 faster ✔
q14: prev= 203 ms, new= 248 ms, diff=1.22 slower ✖
q15: prev= 241 ms, new= 273 ms, diff=1.13 slower ✖
q16: prev= 202 ms, new= 168 ms, diff=1.20 faster ✔
q17: prev= 358 ms, new= 331 ms, diff=1.08 faster ✔
q18: prev= 442 ms, new= 457 ms, diff=1.03 slower ✖
q19: prev= 233 ms, new= 260 ms, diff=1.12 slower ✖
q20: prev= 426 ms, new= 432 ms, diff=1.01 slower ✖
q21: prev= 704 ms, new= 651 ms, diff=1.08 faster ✔
q22: prev= 180 ms, new= 208 ms, diff=1.16 slower ✖
Verification and run detailsJob
Workload: Capacity: 12 Other timings: Queue 1s · Dataset validation 0s · Total 5m 0s |
Stack
Adds
DynamicStageBuiltHandler, called whenever dynamic planning builds a stage. A handler can inspect cost and the outgoing plan, rewrite it, or abort execution.