Repository navigation
feat: expose local task plans to routing handlers - #797
alexanderbianchi wants to merge 1 commit into
Conversation
| // Keep native providers available to routing. Only the worker-facing view | ||
| // replaces them with remote feeds; retained file metadata stays shared. | ||
| let local_plan = transformed.data; |
There was a problem hiding this comment.
This here is not really a local plan. The plan here has already been specialized to a task by trimming down irrelevant UNION branches, specializing DistributedLeafExec children for the task index, etc...
It'd be more correct to capture the local plan right before the first transform_down_with_dt_ctx call (~line 475)
| /// Task-specialized coordinator view, before native feeds become remote handles. | ||
| /// Use this view for inspecting connector metadata, not for execution. Native | ||
| /// providers may be shared with other tasks; do not consume their work streams. | ||
| /// Use [`Self::visit_local_task_plan`] to account for isolated UNION branches. | ||
| pub local_task_plan: &'a Arc<dyn ExecutionPlan>, |
There was a problem hiding this comment.
This is not a bad idea, I could see this being useful, although it might complicate the API a bit more. I'll throw some ideas in case any of them stick:
-
Store any information that should survive a WorkUnitFeed remote->local transition in the containing DataSource instead of the WorkUnitFeed. With this, you'd not care about whether there's a WUF or not in the DataSource, you just care about whatever field in the DataSource impl you chose to hold the routing-relevant data.
-
Add a new optional
.to_task_specialized_extension(&self, df_ctx: DistributedTaskContext) -> Box<dyn Any>method to theWorkUnitFeedProvidertrait. That way, the planner will call thisto_task_specialized_extensionmethod dumping any arbitraryBox<dyn Any>somewhere into theWorkUnitFeedstructure, maybe in theRemoteOrLocalProvideror something?
#[derive(Debug, Clone)]
pub enum RemoteOrLocalProvider<T: WorkUnitFeedProvider> {
Local(T),
+ Remote(RemoteFeedProvider, Option<Box<dyn Any + Send + Sync>>),
}Pretty drafty, not suggesting this specifically, just an idea.
Have you tried 1)? it sounds like the simpler approach compatible with the capabilities we have today.
Otherwise, your idea about exposing both the task_specialized_plan and the original_plan is not bad.
Summary of here i'm at with the POC's right nowEager planning makes file metadata available before execution, so we can use it for worker placement. Today that metadata sits behind the native work-unit feed, which becomes a remote reference before routing runs. We need to preserve access to the task’s assigned files after that conversion. We can retain separate metadata on the datasource, feed wrapper, or make the assigned files the execution input itself. Whatever solution, I want one source of truth for the files+metadata and a natural lifecycle that doesn't hide information somewhere to get around the current flow.
ConclusionPending, but leaning right now towards something like #799, closer to @gabotechs's first recommendation without forcing the data source into the upstream |
The existing routing API lets a handler inspect a task’s plan, choose a worker, and establish a connection. Whether that is sufficient depends on where the connector stores the task’s assigned work.
Files embedded in the scan plan
A connector can split its files into separate scan plans during distributed planning:
The routing handler receives the corresponding task plan. It can inspect the files, derive cache keys from their storage locations, and choose a worker:
The existing API is sufficient because the assigned files remain directly accessible in the plan.
Files supplied through a work-unit feed
Another connector can retain file descriptors on the coordinator and deliver them to workers through a feed:
The assignment algorithm determines which files each task receives. The worker’s scan contains a reference to that feed rather than an embedded file list:
Before invoking routing, DFD currently converts the scan into this worker-facing representation. The conversion preserves the feed ID but omits the native provider and retained files.
The routing handler therefore receives:
It can identify the feed, but the routing API provides no lookup from that ID to the task’s assigned files. Those files still exist on the coordinator; they are inaccessible through the supplied plan.
Expose the local plan alongside the worker-facing plan
This change gives routing access to both representations:
The handler inspects the local scan using its task index and count. A connector can expose the assigned files through a borrowed iterator using the same assignment logic as delivery.
This lets applications implement file-aware placement regardless of whether files are embedded in the worker’s plan or supplied later through a feed. The routing API exposes the information; the application decides the placement policy.
Implementation
RouteTaskEvent::local_task_planretains native providers alongside the existing worker-facingtask_specialized_plan.visit_local_task_plansupplies scan-local task indices and counts, including isolated UNION branches.assigned_filesborrows descriptors using the same assignment logic as delivery. Greedy inspection still computes temporary destination indices.Stacked on #789. This explores the local-plan alternative to #790; it does not add a built-in affinity router.
Validation
RouteTaskEventdirectly.