# Upgrading from 2.0.0 to 3.0.0
This guide covers the source and behavioural changes merged after the
`v2.0.0` tag. The largest change is that the old `TaskEstimator` extension point
has been split into focused planning events. The other source-breaking area is
the `ChannelResolver` transport boundary.
## 1. `WorkerChannel` and `ChannelResolver`
The default gRPC resolver and client-construction helper are no longer root
exports. Move those imports under `datafusion_distributed::grpc`:
```rust
// 2.0
use datafusion_distributed::{
BoxCloneSyncChannel, DefaultChannelResolver, WorkerServiceClient,
create_worker_client,
};
// 3.0
use datafusion_distributed::grpc::{
BoxCloneSyncChannel, DefaultChannelResolver, create_worker_client,
};
```
The generated `WorkerServiceClient` and `WorkerServiceServer` types are no
longer public API. Use `Worker::into_worker_server()` for the server.
`ChannelResolver::get_worker_client_for_url` no longer returns the generated
Tonic `WorkerServiceClient<BoxCloneSyncChannel>`. It returns
`Box<dyn WorkerChannel>`:
```rust
// 2.0
async fn get_worker_client_for_url(
&self,
url: &Url,
) -> Result<WorkerServiceClient<BoxCloneSyncChannel>, DataFusionError>;
// 3.0
async fn get_worker_client_for_url(
&self,
url: &Url,
) -> Result<Box<dyn WorkerChannel>, DataFusionError>;
```
For a resolver using the built-in gRPC transport, pass the
`grpc::BoxCloneSyncChannel` produced by its existing connection and caching
logic to `grpc::create_worker_client`. The helper already returns
`Box<dyn WorkerChannel>`:
```rust
use datafusion_distributed::{WorkerChannel, grpc};
let client: Box<dyn WorkerChannel> = grpc::create_worker_client(channel);
```
## 2. Move `TaskEstimator` methods to event handlers
`TaskEstimator`, `TaskEstimation`, `TaskRoutingContext`, and
`DistributedExt::{with,set}_distributed_task_estimator` have been removed.
`DistributedConfig::with_task_estimator` has also been removed. Register the
corresponding handlers on the `SessionConfig`, `SessionState`, or
`SessionContext` that owns distributed planning; they are not fields of
`DistributedConfig`.
They are replaced by independently registered handlers:
| `task_estimation` | `DesiredTaskCountHandler` | `DesiredTaskCountEventResponse` |
| `scale_up_leaf_node` | `ScaleUpLeafNodeHandler` | `ScaleUpLeafNodeEventResponse` |
| `route_tasks` | `RouteTasksHandler` | `RouteTasksEventResponse` |
Handlers receive an event containing the plan and the context appropriate to
that lifecycle phase. Return `None` when the handler does not recognize the
plan, so the next registered handler (and then built-ins) can handle it.
Here is a direct migration for an estimator that implemented all three hooks:
```rust
// 2.0
impl TaskEstimator for MyEstimator {
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>>> { /* ... */ }
}
let state = SessionStateBuilder::new()
.with_distributed_task_estimator(MyEstimator)
.build();
```
```rust
// 3.0
use datafusion::common::Result;
use datafusion_distributed::{
DesiredTaskCountEvent, DesiredTaskCountEventResponse,
DistributedExt, RouteTasksEvent, RouteTasksEventResponse,
ScaleUpLeafNodeEvent, ScaleUpLeafNodeEventResponse,
};
fn desired_task_count(
event: DesiredTaskCountEvent<'_>,
) -> Option<Result<DesiredTaskCountEventResponse>> {
// event.plan replaces `plan`; event.session_config replaces `cfg`.
let my_leaf = event.plan.downcast_ref::<MyLeaf>()?;
Some(Ok(DesiredTaskCountEventResponse::desired(my_leaf.task_count())))
}
fn scale_up_leaf(
event: ScaleUpLeafNodeEvent<'_>,
) -> Option<Result<ScaleUpLeafNodeEventResponse>> {
let my_leaf = event.plan.downcast_ref::<MyLeaf>()?;
let plan = split_leaf(my_leaf, event.task_count, event.session_config)?;
Some(Ok(ScaleUpLeafNodeEventResponse::new(plan)))
}
fn route_tasks(
event: RouteTasksEvent<'_>,
) -> Option<Result<RouteTasksEventResponse>> {
// Resolve available URLs here if the old code used ctx.available_urls.
let urls = route_stage(event.plan, event.task_count, &event.task_ctx)?;
Some(Ok(RouteTasksEventResponse::new(urls)))
}
let state = SessionStateBuilder::new()
.with_distributed_desired_task_count_handler(desired_task_count)
.with_distributed_scale_up_leaf_node_handler(scale_up_leaf)
.with_distributed_route_tasks_handler(route_tasks)
.build();
```
Registration order matters, same as with multiple `TaskEstimator`s. User
handlers run in registration order before built-in fallbacks. For the
desired-count, leaf-scale, and routing events, the first applicable handler
supplies the result;
See [distributing a custom execution plan](../source/user-guide/04-distribute-custom-plan.md)
and [routing tasks](../source/advanced/06-worker-routing.md) for complete
current examples.
## 3. Move `add_on_plan_hook` to a worker plan rewrite handler
`Worker::add_on_plan_hook` was first changed to accept
`(plan, &SessionConfig)` and return `Result`, then removed. Register a
`WorkerPlanRewriteHandler` in the session builder used by each worker instead.
WARNING: this event handler should be registered on workers, not on the
coordinating context.
```rust
// 2.0
worker.add_on_plan_hook(|plan| rewrite(plan));
```
```rust
// 3.0
use datafusion_distributed::{
DistributedExt, WorkerPlanRewriteEvent, WorkerPlanRewriteEventResponse,
};
fn worker_plan_rewrite_handler(
event: WorkerPlanRewriteEvent
) -> Result<WorkerPlanRewriteEventResponse> {
let new_plan = rewrite(event.plan);
WorkerPlanRewriteEventResponse::new(new_plan)
}
.builder
.with_distributed_worker_plan_rewrite_handler(worker_plan_rewrite_handler)
.build())
});
```
Every rewrite handler runs in registration order and receives the preceding
handler's plan. It must preserve stage topology and the head plan's schema,
partitioning, ordering, boundedness, emission type, and output rows. An error
now aborts worker plan registration.
See [Worker plan rewrite handlers](../source/advanced/03-plan-hooks.md).
## Behavior and output changes
- **Broadcast joins are now enabled by default.** A `CollectLeft` hash join
broadcasts its build side to every consumer task instead of coalescing it to
one partition. If the 2.0 behaviour is required while validating the upgrade,
set `.with_distributed_broadcast_joins(false)?` on the coordinating session,
then remove it after evaluating the new plan and memory profile.
- **Distributed `EXPLAIN` / display output is intentionally more compact.**
Stage headers now look like `tasks=80, partitions=1280` rather than listing
every task's partition range. Per-task metrics are grouped, for example
`output_rows={0:132, 1:216}`, rather than suffixed metric names such as
`output_rows_0=132`. Update golden files and any parser that consumed the old
format.
- **`plan_bytes_sent` is a bytes metric.** Update metric-type assumptions in
exporters or assertions that treated it as a generic counter.
- **Final worker metrics are now collected after an execution stream is
dropped.** A query stopped early (for example by `LIMIT`) can now report
drop-finalized metrics that were previously absent. This is a correctness
improvement, but metric snapshots may change.
## Relevant merged PRs
- [#499: change the worker plan-hook contract](https://github.com/datafusion-contrib/datafusion-distributed/pull/499)
- [#503: finalize execution-stream metrics before collecting them](https://github.com/datafusion-contrib/datafusion-distributed/pull/503)
- [#506: represent
`plan_bytes_sent` as a bytes metric](https://github.com/datafusion-contrib/datafusion-distributed/pull/506)
- [#527: compact distributed plan display](https://github.com/datafusion-contrib/datafusion-distributed/pull/527)
- [#512: abstract the worker protocol from gRPC](https://github.com/datafusion-contrib/datafusion-distributed/pull/512)
- [#564: replace
`TaskEstimator` with planning lifecycle handlers](https://github.com/datafusion-contrib/datafusion-distributed/pull/564)
- [#567: route tasks through events](https://github.com/datafusion-contrib/datafusion-distributed/pull/567)
- [#568: replace worker plan hooks with rewrite handlers](https://github.com/datafusion-contrib/datafusion-distributed/pull/568)
- [#569: make desired-task-count handlers async](https://github.com/datafusion-contrib/datafusion-distributed/pull/569)
- [#574: enable broadcast joins by default](https://github.com/datafusion-contrib/datafusion-distributed/pull/574)