diff --git a/docs/source/advanced/04-work-unit-feeds.md b/docs/source/advanced/04-work-unit-feeds.md index 837c0a67..957bbce1 100644 --- a/docs/source/advanced/04-work-unit-feeds.md +++ b/docs/source/advanced/04-work-unit-feeds.md @@ -7,8 +7,8 @@ fully known at planning time. ## When to use a work unit feed Distributed DataFusion already has a mechanism for splitting a leaf node's work across tasks: the -[`TaskEstimator`](../user-guide/04-distribute-custom-plan.md). With a `TaskEstimator`, all of the work is decided **at planning -time** — `scale_up_leaf_node` produces `N` per-task plan variants, each pre-loaded with the slice of +[desired task-count and leaf-scale handlers](../user-guide/04-distribute-custom-plan.md). With these handlers, all of the work is decided **at planning +time** — `ScaleUpLeafNodeHandler` produces `N` per-task plan variants, each pre-loaded with the slice of data it is responsible for (e.g., a group of files). This is the right tool whenever the units of work are known before execution begins. @@ -24,18 +24,19 @@ work-unit descriptors from your provider and streams each one — over gRPC — owns that partition. The worker turns each descriptor into rows. Slow partitions don't block fast ones, and back-pressure is handled per partition. -| | `TaskEstimator` | Work unit feed | +| | Event handlers | Work unit feed | |-------------------------|---------------------------------------------|-----------------------------------------------| | When work is known | Planning time | Runtime | | How work is distributed | Pre-built per-task plan variants | Streamed per-partition from the coordinator | | Good for | Files, ranges, anything enumerable up front | Paginated APIs, queues, progressive discovery | -The two mechanisms are complementary. A feed-backed leaf still provides a `TaskEstimator` to tell the -planner how many tasks to use; the feed only governs *what work flows into each partition at runtime*. +The two mechanisms are complementary. A feed-backed leaf still provides a +`DesiredTaskCountHandler` to tell the planner how many tasks to use; the feed only +governs *what work flows into each partition at runtime*. -> **A `TaskEstimator` is always required for a feed-backed leaf.** The feed decides *what* work each +> **A desired task-count handler is always required for a feed-backed leaf.** The feed decides *what* work each > partition receives at runtime, but something still has to tell Distributed DataFusion the *desired task -> count* for the node — that is the `TaskEstimator`'s job (`task_estimation`). Without it the leaf defaults +> count* for the node — that is the `DesiredTaskCountHandler`'s job. Without it the leaf defaults > to a single task. See [Distribute a custom `ExecutionPlan`](../user-guide/04-distribute-custom-plan.md). ## How it works @@ -199,7 +200,8 @@ let state = SessionStateBuilder::new() .with_distributed_worker_resolver(/* ... */) .with_distributed_planner() .with_distributed_user_codec(MyExecCodec) // so workers can deserialize the node - .with_distributed_task_estimator(MyTaskEstimator) // how many tasks the leaf gets + .with_distributed_desired_task_count_handler(my_desired_task_count) // how many tasks the leaf gets + .with_distributed_scale_up_leaf_node_handler(my_scale_up_leaf_node) // per-task plan variants .with_distributed_work_unit_feed(|exec: &MyExec| Some(&exec.feed)) .build(); ``` @@ -218,7 +220,8 @@ distributed node, plus the feed registration: `feed.feed(partition, ctx)?` to produce `RecordBatch`es. 3. A `PhysicalExtensionCodec` that serializes the node, encoding the feed handle with `to_proto()` / `from_proto()`. -4. A `TaskEstimator` so the planner knows how many tasks the leaf stage should use. +4. A `DesiredTaskCountHandler` so the planner knows how many tasks the leaf stage should use, and a + `ScaleUpLeafNodeHandler` when the leaf needs per-task plan variants. There is a complete, runnable example in the `examples/` folder: diff --git a/docs/source/advanced/05-custom-distributed-plans.md b/docs/source/advanced/05-custom-distributed-plans.md index b483e373..4ea38696 100644 --- a/docs/source/advanced/05-custom-distributed-plans.md +++ b/docs/source/advanced/05-custom-distributed-plans.md @@ -73,14 +73,14 @@ many tasks into fewer). ## Leaf data splitting still happens automatically -Even when you inject the boundaries yourself, the distributed planner runs the registered -[`TaskEstimator`](../user-guide/04-distribute-custom-plan.md) over each stage's leaves and calls `scale_up_leaf_node` with the -stage's task count. So a parquet `DataSourceExec` is wrapped in a `DistributedLeafExec` (with one -per-task file-group variant) by the default file-scan estimator, and a custom leaf is split by whatever -`TaskEstimator` you registered for it — exactly as in the automatic path. You only place the boundaries; -the leaves are scaled for you. - -This means a custom leaf node still needs its `TaskEstimator` (and its `scale_up_leaf_node` / +Even when you inject the boundaries yourself, the distributed planner evaluates +the registered [desired task-count handler](../user-guide/04-distribute-custom-plan.md) and then calls the +registered `ScaleUpLeafNodeHandler` with each stage's final task count. So a parquet +`DataSourceExec` is wrapped in a `DistributedLeafExec` (with one per-task file-group variant) by the +default file-scan handlers, and a custom leaf is split by its registered handlers — exactly as in the +automatic path. You only place the boundaries; the leaves are scaled for you. + +This means a custom leaf node still needs its desired task-count and leaf-scale handlers (or `DistributedTaskContext`-based dispatch) registered, just as it would for automatic planning — you do **not** need to hand-build `DistributedLeafExec` in your boundary-injection rule. @@ -110,5 +110,5 @@ There is a complete, runnable example in the `examples/` folder: - [custom_distributed_partial_reduction_tree.rs](https://github.com/datafusion-contrib/datafusion-distributed/blob/main/examples/custom_distributed_partial_reduction_tree.rs) — a `PhysicalOptimizerRule` that rewrites a `GROUP BY` aggregation over a parquet table into the tree above (`Partial → NetworkCoalesce → PartialReduce → NetworkCoalesce → Final`). It only injects the - boundaries; the planner's `TaskEstimator` splits the parquet leaf across the leaf-stage tasks + boundaries; the planner's default event handlers split the parquet leaf across the leaf-stage tasks automatically. diff --git a/docs/source/advanced/06-worker-routing.md b/docs/source/advanced/06-worker-routing.md index c2b73cf0..1d0b32f3 100644 --- a/docs/source/advanced/06-worker-routing.md +++ b/docs/source/advanced/06-worker-routing.md @@ -5,25 +5,28 @@ round-robin. When a task's data has a *home* — a worker that already holds it cache or on local disk — you can send the task **there** instead, so it reads locally instead of pulling data over the network. -Routing is the third method of the -[`TaskEstimator`](../user-guide/04-distribute-custom-plan.md) trait, `route_tasks`. It receives -a `TaskRoutingContext` (the head plan of the stage, the task count, and the active +Routing is handled by a registered `RouteTasksHandler`. It receives a +`RouteTasksEvent` (the head plan of the stage, the task count, and the active `TaskContext`) and returns one worker URL per task, in task order: ```rust -fn route_tasks(&self, routing_ctx: &TaskRoutingContext<'_>) -> Result>>; +fn route_tasks(event: RouteTasksEvent) -> Option>; ``` -- `Ok(Some(urls))` — task `i` is sent to `urls[i]`. -- `Ok(None)` — the default; keep the round-robin behaviour. +- `Some(Ok(RouteTasksEventResponse::new(urls)))` — task `i` is sent to `urls[i]`. +- `None` — defer to the next handler; the built-in fallback keeps the round-robin behaviour. -Because `route_tasks` is part of `TaskEstimator`, you implement it on the same -estimator you register with `with_distributed_task_estimator`. +Register it on the coordinating session builder: -Routing pairs naturally with `scale_up_leaf_node`: that decides *what* data task -`i` reads, and `route_tasks` decides *where* it runs. If your estimator returned a -`DistributedLeafExec`, its `variants()` are in task order too, so you can line up -each task's data with the worker that should serve it. +```rust +SessionStateBuilder::new() + .with_distributed_route_tasks_handler(route_tasks); +``` + +Routing pairs naturally with `ScaleUpLeafNodeHandler`: that decides *what* data +task `i` reads, and `RouteTasksHandler` decides *where* it runs. If the leaf-scale +handler returns a `DistributedLeafExec`, its `variants()` are in task order too, +so you can line up each task's data with the worker that should serve it. For a complete, runnable walkthrough — parquet files consistently routed to workers by rendezvous hashing of the file path, so each worker can serve them from diff --git a/docs/source/advanced/08-adaptive-query-execution.md b/docs/source/advanced/08-adaptive-query-execution.md index 6875f48e..4e5e4c80 100644 --- a/docs/source/advanced/08-adaptive-query-execution.md +++ b/docs/source/advanced/08-adaptive-query-execution.md @@ -40,9 +40,9 @@ The following settings affect AQE decisions: AQE depends on both planning-time statistics and execution-time metrics. A custom leaf should provide all of the following: -- A registered `TaskEstimator` that implements `scale_up_leaf_node` for the task - count selected by AQE. Without an estimator, the planner treats an unknown - leaf as limited to one task. +- Registered desired task-count and leaf-scale handlers. The desired handler supplies a task-count + hint; the leaf-scale handler specializes the leaf for the count selected by AQE. Without a desired + handler, the planner treats an unknown leaf as limited to one task. - A useful implementation of `ExecutionPlan::partition_statistics` in the custom data source. The more accurate the statistics are, the better the decisions AQE can make. diff --git a/docs/source/learn/01-concepts.md b/docs/source/learn/01-concepts.md index e2d1f175..a87a812d 100644 --- a/docs/source/learn/01-concepts.md +++ b/docs/source/learn/01-concepts.md @@ -61,9 +61,12 @@ Different organizations have different networking requirements—from Kubernetes deployments to cloud provider solutions. This trait allows Distributed DataFusion to adapt to various scenarios. -## [TaskEstimator](https://github.com/datafusion-contrib/datafusion-distributed/blob/main/src/distributed_planner/task_estimator.rs) +## [Planner event handlers](https://github.com/datafusion-contrib/datafusion-distributed/tree/main/src/events) -Estimates the number of tasks required in the leaf stage of a distributed query. +Event handlers let applications participate in distinct distributed-planning lifecycle phases: +`DesiredTaskCountHandler` supplies task-count hints, `ScaleUpLeafNodeHandler` +specializes leaves after the final count is known, and `RouteTasksHandler` assigns +task slots to workers. The number of tasks each stage has is determined from bottom to top. This means that leaf stages will decide how many tasks they need to execute based on the @@ -84,8 +87,9 @@ is needed. For custom leaf nodes that need to dispatch work themselves, `DistributedTaskContext` exposes `task_index` and `task_count` so execution logic can select the appropriate data subset. For example, task 0 of 3 might -return the first third of rows, task 2 the last third, and so on. See the -`TaskEstimator` documentation for guidance on which approach to use. +return the first third of rows, task 2 the last third, and so on. See +[Distribute a custom `ExecutionPlan`](../user-guide/04-distribute-custom-plan.md) +for guidance on which approach to use. ## [ChannelResolver](https://github.com/datafusion-contrib/datafusion-distributed/blob/main/src/protocol/channel_resolver.rs) diff --git a/docs/source/learn/02-how-a-distributed-plan-is-built.md b/docs/source/learn/02-how-a-distributed-plan-is-built.md index f65f77c0..913c070a 100644 --- a/docs/source/learn/02-how-a-distributed-plan-is-built.md +++ b/docs/source/learn/02-how-a-distributed-plan-is-built.md @@ -79,12 +79,17 @@ The first step is to split the leaf node into different tasks: Each task will handle a different non-overlapping piece of data. -The number of tasks that will be used for executing leaf nodes is determined by a `TaskEstimator` implementation. -A default implementation exists for file-based `DataSourceExec` nodes. However, since `DataSourceExec` can be -customized to represent any data source, users with custom implementations should also provide a corresponding -`TaskEstimator`. +The number of tasks that will be used for executing leaf nodes is determined by +`DesiredTaskCountHandler` implementations. Default handlers exist for file-based +`DataSourceExec` nodes. However, since `DataSourceExec` can be customized to represent +any data source, users with custom implementations should also provide corresponding +desired task-count and leaf-scale handlers. -In the case above, a `TaskEstimator` decided to use four tasks for the leaf node. Note that even if we are distributing +the data across different tasks, each task will also distribute its data across partitions using the vanilla DataFusion +partitioning mechanism. A partition is a split of data processed by a single thread on a single machine, whereas a +In the case above, a desired task-count handler decided to use four tasks for the leaf node. Note that even if we are distributing +the data across different tasks, each task will also distribute its data across partitions using the vanilla DataFusion +partitioning mechanism. A partition is a split of data processed by a single thread on a single machine, whereas a the data across different tasks, each task will also distribute its data across partitions using the vanilla DataFusion partitioning mechanism. A partition is a split of data processed by a single thread on a single machine, whereas a task is a split of data processed by an entire machine within a cluster. @@ -122,7 +127,7 @@ Flight, and each `NetworkShuffleExec` instance will know from which partitions a Note how this means that we have just built the first stage, as the first network boundary was introduced. We are now in the process of building the second stage, and note how it has just two tasks. -If the number of tasks in a leaf stage is driven by the hints given by `TaskEstimator`s, the number of tasks in upper +If the number of tasks in a leaf stage is driven by the hints given by desired task-count handlers, the number of tasks in upper stages is driven by the nodes in between that reduce or increase the cardinality of the data. In this case, the leaf stage is performing a partial aggregation before sending data to the next stage, so we can diff --git a/docs/source/learn/03-how-adaptive-query-execution-works.md b/docs/source/learn/03-how-adaptive-query-execution-works.md index a13dae26..98d01cd3 100644 --- a/docs/source/learn/03-how-adaptive-query-execution-works.md +++ b/docs/source/learn/03-how-adaptive-query-execution-works.md @@ -40,8 +40,8 @@ The coordinator walks this plan from the leaves upward. When it reaches the first point where a network boundary is required, it has enough information to close the plan below that boundary as the first stage. Because this stage contains a data-source leaf, its cost is calculated from -`ExecutionPlan::partition_statistics`. The cost model and the registered -`TaskEstimator` are then used to choose its task count: +`ExecutionPlan::partition_statistics`. The cost model and the registered desired +task-count handlers are then used to choose its task count: ```text SortPreservingMergeExec diff --git a/docs/source/user-guide/04-distribute-custom-plan.md b/docs/source/user-guide/04-distribute-custom-plan.md index 17c06369..dd6caf50 100644 --- a/docs/source/user-guide/04-distribute-custom-plan.md +++ b/docs/source/user-guide/04-distribute-custom-plan.md @@ -101,73 +101,63 @@ node is encoded on one side and decoded on the other. ## 2. Choose how many tasks to use -Implement `TaskEstimator` and register it with -`.with_distributed_task_estimator(...)`. Its first method, `task_estimation`, -tells the planner how many tasks the stage containing your leaf should run on. +Register a `DesiredTaskCountHandler` with +`.with_distributed_desired_task_count_handler(...)`. It tells the planner how +many tasks the stage containing your leaf should run on. For a sharded scan, one task per shard is a natural choice: ```rust -impl TaskEstimator for ShardedScanEstimator { - fn task_estimation( - &self, - plan: &Arc, - _cfg: &ConfigOptions, - ) -> Option { - // Only estimate for our own node; returning None lets other estimators try. - let scan = plan.downcast_ref::()?; - // One task per shard — the planner caps this at the number of workers. - Some(TaskEstimation::desired(scan.shards.len())) - } - - // scale_up_leaf_node: see the next section. +fn sharded_scan_desired_task_count( + event: DesiredTaskCountEvent, +) -> Option { + // Only handle our own node; returning None lets other handlers try. + let scan = event.plan.downcast_ref::()?; + // One task per shard — the planner caps this at the number of workers. + Some(DesiredTaskCountEventResponse::desired(scan.shards.len())) } ``` What the return value means: -- `TaskEstimation::desired(n)` — a **soft** hint. The planner may land on a +- `DesiredTaskCountEventResponse::desired(n)` — a **soft** hint. The planner may land on a different number: within a stage the largest `desired` wins, and the count is capped at the number of available workers. -- `TaskEstimation::maximum(n)` — a **hard** cap. `maximum(1)` means "this node +- `DesiredTaskCountEventResponse::maximum(n)` — a **hard** cap. `maximum(1)` means "this node cannot be distributed." -- `None` — defer to the other registered estimators (and finally the built-in - file-scan estimator). +- `None` — defer to the other registered handlers (and finally the built-in + file-scan handler). To send each task to a specific worker instead of the default round-robin, see [Routing tasks to workers](../advanced/06-worker-routing.md). ## 3. Split the work across tasks -Once the final task count is settled, the planner calls -`scale_up_leaf_node(plan, task_count, cfg)` on your estimator. This is where you -divide the leaf's work into `task_count` **non-overlapping** pieces — here, by -handing each task its own subset of shards. +Once the final task count is settled, the planner calls the registered +`ScaleUpLeafNodeHandler`. This is where you divide the leaf's work into +`event.task_count` **non-overlapping** pieces — here, by handing each task its +own subset of shards. The recommended approach is to return a `DistributedLeafExec` wrapping one **variant** of your node per task: ```rust -fn scale_up_leaf_node( - &self, - plan: &Arc, - task_count: usize, - _cfg: &ConfigOptions, -) -> Result>> { - let Some(scan) = plan.downcast_ref::() else { - return Ok(None); - }; - - // Spread the shards across `task_count` tasks, one variant each. - let mut per_task: Vec> = vec![Vec::new(); task_count]; +fn sharded_scan_scale_up_leaf_node( + event: ScaleUpLeafNodeEvent, +) -> Option> { + let scan = event.plan.downcast_ref::()?; + + // Spread the shards across `event.task_count` tasks, one variant each. + let mut per_task: Vec> = vec![Vec::new(); event.task_count]; for (i, shard) in scan.shards.iter().enumerate() { - per_task[i % task_count].push(shard.clone()); + per_task[i % event.task_count].push(shard.clone()); } let variants = per_task.into_iter().map(|shards| { Arc::new(ShardedScanExec::new(shards, scan.schema())) as Arc }); - let leaf = DistributedLeafExec::try_new(Arc::clone(plan), variants)?; - Ok(Some(Arc::new(leaf))) + let leaf = DistributedLeafExec::try_new(Arc::clone(event.plan), variants) + .map(|leaf| ScaleUpLeafNodeEventResponse::new(Arc::new(leaf))); + Some(leaf) } ``` @@ -205,13 +195,14 @@ let state = SessionStateBuilder::new() .with_distributed_worker_resolver(resolver) .with_distributed_planner() .with_distributed_user_codec(ShardedScanCodec) - .with_distributed_task_estimator(ShardedScanEstimator) + .with_distributed_desired_task_count_handler(sharded_scan_desired_task_count) + .with_distributed_scale_up_leaf_node_handler(sharded_scan_scale_up_leaf_node) .build(); ``` Each worker only needs the codec — the planner, worker resolver, and task -estimator are coordinating-context concerns; workers just decode and run the plan -variants they receive: +event handlers are coordinating-context concerns; workers just decode and run +the plan variants they receive: ```rust let worker = Worker::from_session_builder(|ctx: WorkerQueryContext| async move { @@ -223,7 +214,7 @@ let worker = Worker::from_session_builder(|ctx: WorkerQueryContext| async move { ``` For a complete, runnable program that follows this same pattern — a custom leaf -split across tasks, with its own codec and `TaskEstimator` — see +split across tasks, with its own codec and event handlers — see [`custom_execution_plan.rs`](https://github.com/datafusion-contrib/datafusion-distributed/blob/main/examples/custom_execution_plan.rs), which distributes a `numbers(start, end)` source.