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 |
## What changes
Add opt-in eager Iceberg file discovery with
`iceberg.eager_file_planning=true`. Physical planning reads manifests,
applies Iceberg's pruning and delete matching, and retains the selected
`FileScanTask`s on an `EagerIcebergDataSource`. Distributed planning
assigns those files to task-local partitions before worker placement.
Each worker receives only its assigned descriptors in its execution plan
and reads them through the existing Iceberg reader.
Lazy discovery remains the default. Both modes use `DataSourceExec`,
shared reader configuration, and Iceberg's `FileIO`; eager planning does
not read Parquet data.
### Lazy
Planning uses snapshot statistics. During execution, the coordinator
reads manifests and progressively supplies file tasks to workers through
a remote work-unit feed.
This overlaps discovery with reading and can avoid some discovery when
execution stops early. Initial planning and worker placement happen
without the complete selected-file set.
### Eager
Planning completes discovery and retains the selected tasks before
execution starts. The source owns the assignment that routing inspects
and workers execute.
This adds startup latency, retained metadata, and larger worker-plan
payloads. In exchange, selected files are available for placement and
future task-sizing, balancing, and statistics improvements. This matters
especially when partition pruning reduces a large table to a handful of
files. Fully consumed scans generally require the same discovery I/O in
either mode.
```text
Lazy: plan → remote feed reference → coordinator discovers and dispatches → reader
Eager: discover → source owns files → assign task-local groups → route
→ serialize each task's groups → worker restores groups → reader
```
## Why replace #773
#773 moves discovery earlier but retains the files behind the work-unit
feed. Before routing, DFD converts native feeds into remote references,
hiding the retained descriptors from the router. Keeping that design
would require another way to preserve access to the files across the
conversion.
This PR instead lands source-owned eager execution directly against
`main`, following [Gabriel's suggestion to retain routing-relevant
information on the data
source](#797 (comment)).
It replaces #773 as the proposed foundation rather than merging eager
feeds and immediately refactoring them.
- **One authoritative assignment:** task-local sources share descriptor
storage on the coordinator. Routing borrows the source's assigned files;
serialization sends those same groups; workers restore them without
reassignment. There is no separate routing metadata copy.
- **Existing routing API:** like DFD's `FileScanConfig` path, the
scale-up handler creates task-specific `DistributedLeafExec` variants.
`RouteTaskEvent::task_specialized_plan` exposes the selected variant
before placement. No second plan view or work-unit-feed extension is
needed.
- **Shared reading, separate delivery:** lazy scans keep their feed
lifecycle; eager scans carry assigned files. Both produce file-task
streams for the same reader implementation. Independent codecs preserve
the existing lazy payload. Coordinator and worker codec registrations
must match.
- **Iceberg-native I/O:** the structure resembles file scans without
forcing Iceberg into `FileScanConfig` or supplying a placeholder
object-store URL.
- **Bounded retention:** descriptors are charged to the query memory
pool, with no arbitrary file-count cap. Abandoned eager planning cancels
its discovery future.
Lazy remains available for progressive startup and early termination.
Keeping it means maintaining two discovery/delivery lifecycles; this PR
isolates that difference rather than making eager execution depend on a
feed.
## Relationship to upstream datafusion-iceberg #67
Like [#67](apache/datafusion-iceberg#67), this
retains native `FileScanTask`s with a private memory-reservation owner,
exposes borrowed accessors on the source, and feeds local partition
streams into the Iceberg reader. Moving ownership out of DFD's feed
brings the eager representation closer to upstream.
DFD adds distributed assignment, placement, and reconstruction of
assigned groups from worker plans. The implementations are not
interchangeable yet: #67 uses one source with optional eager planning
and local lazy discovery; this PR uses separate eager/lazy source types
and retains DFD's remote feeds for lazy scans. Upstream would need to
support reconstructing an already assigned scan without rediscovery or
repartitioning before it could replace this implementation directly.
## Source and scheduling boundary
The eager source represents local file groups; it does not choose
distributed task assignments. `with_file_groups` derives a source using
positions in `planned_files()` (retained-file order), preserving
supplied group order and sharing descriptors and their memory
reservation. It rejects invalid or repeated positions and resets
metrics/statistics for the resulting scan.
The scale-up handler chooses task assignments through this public API.
The codec uses public schema, file, and partition accessors and
reconstructs received native descriptors with `from_file_groups`,
without rediscovery. `PlannedFiles` remains private. Reader
configuration and shared delete-loader behavior still need
reconciliation with upstream #67; this change separates scheduling
without changing reader behavior.
## Scope and follow-ups
This includes eager discovery, descriptor retention, task-first
assignment, routing access, and transport. It preserves current snapshot
statistics and task-size estimates; task-local eager variants and
decoded plans report unknown statistics instead of applying
whole-snapshot aggregates to a subset.
Selected-file sizing (#787), byte balancing (#788), and selected-file
statistics (#789) remain separate. Affinity policy can use the existing
routing hook outside this crate. Better selected-file metrics still need
iceberg-rust support to avoid rereading manifests. This PR makes no
performance claim.
## How to review
Start with [the eager-planning feature
commit](d6c8b6a).
The first four commits extract existing statistics, reader
configuration, lazy scan construction, and task-first partition
calculation.
The next two commits address review feedback separately: [module moves
and source
names](bd68327),
then [readable routing
tests](bad9344).
In the final tree, follow `table_provider/scan.rs` →
`eager/data_source.rs` → the scale-up handler → `eager/codec.rs`. Lazy
sources and codecs live under `lazy/`, with shared reading in
`common/scan_config.rs`. The old `IcebergDataSource` and `IcebergCodec`
exports are removed; consumers must use `LazyIcebergDataSource` and
`LazyIcebergCodec`, as documented in the upgrade guide.
The latest [source/scheduling separation
commit](96f55ed)
adds these public construction/regrouping operations and moves task
assignment into the scale-up handler.
## Validation
Coverage includes eager/lazy result parity, per-task distribution
snapshots, placement after pruning, isolated UNION branches, repeated
scan execution, and query-memory limits. Routing snapshots show the
assigned files, selected worker, and query results; normal distributed
execution exercises transport. Codec tests reject missing configuration
and malformed groups/descriptors.
The full Iceberg suite previously passed (61 tests and one doctest).
After the grouping-test cleanup, all library tests, grouping tests, and
routing tests passed with snapshot updates disabled. All-target Clippy
with warnings denied and formatting also passed.
Public grouping tests snapshot partition groups and actual dates/counts
read after regrouping and reconstruction. Invalid cases are labeled
explicitly; reservation lifetime is tested separately. The codec
validation test uses `encode_config` through public source accessors
without a synthetic plan roundtrip. Existing routing/distribution
snapshots remain unchanged.
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.