Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

Appearance settings

Define worker plan rewrites in terms of events - #568

#568
Open
gabotechs wants to merge 1 commit into
gabrielmusat/routing-eventsdatafusion-contrib/datafusion-distributed:gabrielmusat/routing-eventsfrom
gabrielmusat/worker-plan-rewrite-eventsdatafusion-contrib/datafusion-distributed:gabrielmusat/worker-plan-rewrite-eventsCopy head branch name to clipboard
Open

Define worker plan rewrites in terms of events#568
gabotechs wants to merge 1 commit into
gabrielmusat/routing-eventsdatafusion-contrib/datafusion-distributed:gabrielmusat/routing-eventsfrom
gabrielmusat/worker-plan-rewrite-eventsdatafusion-contrib/datafusion-distributed:gabrielmusat/worker-plan-rewrite-eventsCopy head branch name to clipboard

Conversation

@gabotechs

@gabotechs gabotechs commented Jul 27, 2026

Copy link
Copy Markdown
Collaborator

Stack

  1. Introduce distributed planning lifecycle event handlers for replacing TaskEstimator #564
  2. Define routing in terms of events #567
  3. Define worker plan rewrites in terms of events #568 <- you are here
  4. Make desired task-count handlers async #569
  5. Add dynamic stage built events #570

Replaces the worker-side OnPlanHook with composable WorkerPlanRewriteHandlers. Register handlers in each worker's session-state builder; coordinator registration is not propagated to workers.

ctx.builder
    .with_distributed_worker_plan_rewrite_handler(|event| {
        Ok(WorkerPlanRewriteEventResponse::new(event.plan))
    })
    .build()

@gabotechs
gabotechs force-pushed the gabrielmusat/routing-events branch from f2af26f to 062cc6c Compare July 27, 2026 08:37
@gabotechs
gabotechs force-pushed the gabrielmusat/worker-plan-rewrite-events branch from 019c8da to d111ed0 Compare July 27, 2026 08:37
@gabotechs
gabotechs force-pushed the gabrielmusat/routing-events branch from 062cc6c to 928c4db Compare July 27, 2026 08:52
@gabotechs
gabotechs force-pushed the gabrielmusat/worker-plan-rewrite-events branch 2 times, most recently from ca9c4f9 to ba2960f Compare July 27, 2026 09:05
@gabotechs
gabotechs marked this pull request as ready for review July 27, 2026 09:54
@gabotechs
gabotechs force-pushed the gabrielmusat/routing-events branch from a5a9f24 to b9636f0 Compare July 28, 2026 09:20
@gabotechs
gabotechs force-pushed the gabrielmusat/worker-plan-rewrite-events branch from ba2960f to f9fde4f Compare July 28, 2026 09:21
gabotechs added a commit that referenced this pull request Jul 28, 2026
… 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);
```
@gabotechs
gabotechs force-pushed the gabrielmusat/routing-events branch from b9636f0 to 814f286 Compare July 28, 2026 15:14
@gabotechs
gabotechs force-pushed the gabrielmusat/worker-plan-rewrite-events branch from f9fde4f to 335a90b Compare July 28, 2026 15:14
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant

Morty Proxy This is a proxified and sanitized view of the page, visit original site.