Skip to content

feat: expose local task plans to routing handlers - #797

Draft
alexanderbianchi wants to merge 1 commit into
alex.bianchi/iceberg-selected-column-bytesfrom
codex/local-task-routing
Draft

alexanderbianchi wants to merge 1 commit into
alex.bianchi/iceberg-selected-column-bytesfrom
codex/local-task-routing

Conversation

@alexanderbianchi

@alexanderbianchi alexanderbianchi commented Oct 22, 2026 •

Copy link
Copy Markdown
Collaborator

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:

Task 0: Scan { files: [A, C] }
Task 1: Scan { files: [B, D] }

The routing handler receives the corresponding task plan. It can inspect the files, derive cache keys from their storage locations, and choose a worker:

Task 0’s scan → files A, C → preferred worker → dial(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:

Coordinator’s scan
  └─ Local feed
       └─ files [A, B, C, D]

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:

Worker’s scan
  └─ Remote feed { id: abc123 }

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:

Task 0’s scan → feed abc123

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:

Local task plan                         Worker-facing task plan
  native feed + retained files            remote feed ID
             │                                   │
             ▼                                   ▼
      routing inspection                    dialer sends

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_plan retains native providers alongside the existing worker-facing task_specialized_plan.
  • visit_local_task_plan supplies scan-local task indices and counts, including isolated UNION branches.
  • Iceberg’s assigned_files borrows descriptors using the same assignment logic as delivery. Greedy inspection still computes temporary destination indices.
  • The dialer continues to send the worker-facing plan. Feed IDs and the worker protocol are unchanged.

Stacked on #789. This explores the local-plan alternative to #790; it does not add a built-in affinity router.

Validation

  • 54 tests passed across Iceberg file planning/local routing and existing task routing, work-unit feeds, and dynamic filtering.
  • Tests cover inspection without consuming feeds, remote worker-facing feeds, UNION task contexts, and agreement between inspected and delivered files.
  • Formatting and targeted Clippy checks passed with warnings denied.
  • Includes a migration note for callers constructing RouteTaskEvent directly.

Comment on lines +501 to +503
// 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;

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Comment thread src/events/route_tasks.rs
Comment on lines +30 to +34
/// 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>,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

  1. 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.

  2. Add a new optional .to_task_specialized_extension(&self, df_ctx: DistributedTaskContext) -> Box<dyn Any> method to the WorkUnitFeedProvider trait. That way, the planner will call this to_task_specialized_extension method dumping any arbitrary Box<dyn Any> somewhere into the WorkUnitFeed structure, maybe in the RemoteOrLocalProvider or 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.

@alexanderbianchi alexanderbianchi Oct 26, 2026 •

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#799 The first idea is similar to what I'm trying here. But I'm winding up down a bit of a rabbit hole, and have #801 as well which completely isolated the eager path. Currently experimenting.

@alexanderbianchi

Copy link
Copy Markdown
Collaborator Author

Summary of here i'm at with the POC's right now

Eager 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.

  • feat: expose local task plans to routing handlers #797: Adds local_task_plan to RouteTaskEvent during task preparation. The native feed remains the source of truth; routing inspects it through the additional plan view, while workers receive remote feed references.
  • Gabriel 1: Adds shared assignment metadata to the data source and preserves it through conversion. Planning, routing, and feed execution can share one authoritative assignment, provided the field references that assignment rather than maintaining a separate summary.
  • Gabriel 2: Adds a task-specific extension to the remote feed wrapper during conversion. That extension can retain the native provider’s shared assignment for routing, while execution continues retrieving files through the feed.
  • feat: embed eager Iceberg assignments in scan inputs #799: Adds ScanInput::AssignedFiles to IcebergDataSource. Eager planning produces files, task specialization assigns them, and those same descriptors flow through routing and serialization into execution.
  • feat: represent eager Iceberg scans with FileScanConfig #801: Places eager assignments in FileScanConfig.file_groups, with Iceberg descriptors attached to each file. Those groups flow through task specialization, routing, and execution, but the standard path/size fields must agree with the descriptors, and the storage URL remains a placeholder.

Conclusion

Pending, but leaning right now towards something like #799, closer to @gabotechs's first recommendation without forcing the data source into the upstream FileSource

This branch has not been deployed

No deployments
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.

2 participants