Skip to content

feat: plan eager Iceberg scans with source-owned files - #803

Open
alexanderbianchi wants to merge 8 commits into
mainfrom
codex/iceberg-assigned-scans
Open

alexanderbianchi wants to merge 8 commits into
mainfrom
codex/iceberg-assigned-scans

Conversation

@alexanderbianchi

@alexanderbianchi alexanderbianchi commented Oct 26, 2026 •

Copy link
Copy Markdown
Collaborator

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

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. 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, this retains native FileScanTasks 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.

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. 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, then readable routing tests. 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.

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 passed (58 tests and one doctest), with snapshot updates disabled. All-target Clippy with warnings denied and formatting also passed.

Comment thread iceberg/src/test_utils/harness.rs Outdated
Comment thread iceberg/src/codec.rs Outdated
Comment thread iceberg/src/data_source.rs Outdated
Comment thread iceberg/src/distributed_scale_up_handler.rs
Comment thread iceberg/src/data_source.rs Outdated
Comment thread iceberg/src/file_partitions.rs Outdated
Comment thread iceberg/src/iceberg_ext.rs
@alexanderbianchi alexanderbianchi changed the title feat: carry eager Iceberg assignments in task scans feat: plan eager Iceberg scans with source-owned files Oct 26, 2026
@alexanderbianchi
alexanderbianchi force-pushed the codex/iceberg-assigned-scans branch from a07bff5 to 093942f Compare October 26, 2026 20:21
@alexanderbianchi
alexanderbianchi changed the base branch from alex.bianchi/iceberg-selected-column-bytes to main October 26, 2026 20:22
@alexanderbianchi
alexanderbianchi force-pushed the codex/iceberg-assigned-scans branch from 093942f to fcf6f1e Compare October 26, 2026 20:31
@alexanderbianchi
alexanderbianchi force-pushed the codex/iceberg-assigned-scans branch from fcf6f1e to 9b37c33 Compare October 26, 2026 20:41
@alexanderbianchi
alexanderbianchi force-pushed the codex/iceberg-assigned-scans branch 3 times, most recently from 74e158a to cb1a7e8 Compare October 26, 2026 21:16

@gabotechs gabotechs left a comment

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.

👍 The approach looks good, my suggestions are mostly around organization, testing and some questions.

Glad to see you've found a good way of representing the eager path with existing tooling!

Comment thread iceberg/src/common/mod.rs Outdated
pub(crate) use error::{df_err, iceberg_err};

/// Visit distributed tasks before their local partitions. Counts must be nonzero.
pub(crate) fn file_partition(index: usize, task_count: usize, partitions_per_task: usize) -> usize {

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.

We typically avoid to commit anything that is not pure use plumbing to mod.rs files.

Even if it's trivial, I'd just place this in a separate file for consistency.

Comment thread iceberg/src/table_provider/scan.rs Outdated
) -> Result<Arc<dyn ExecutionPlan>> {
let output_schema = match opts.projection {
None => schema.clone(),
Some(projection) => Arc::new(schema.project(projection).unwrap()),

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.

Suggested change
Some(projection) => Arc::new(schema.project(projection).unwrap()),
Some(projection) => Arc::new(schema.project(projection)?),

Comment thread iceberg/src/table_provider/scan.rs Outdated
partitions[file_partition(index, 1, partitioning.partition_count())].push(index);
}
return Ok(DataSourceExec::from_data_source(
AssignedIcebergDataSource {

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.

🤔 I'm thinking about the naming, what about renaming them:

  • EagerIcebergDataSource
  • LazyIcebergDataSource
    ?

Those are the terms we use typically while talking about it.

} else {
None
};
if IcebergConfig::from_task_context(&context).eager_file_planning {

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.

I wonder how this would look like once apache/datafusion-iceberg#67 lands.

I imagine we no longer want to maintain here any extra param determining whether we want the discovery to be lazy or not, we'd just follow whatever we decide in apache/datafusion-iceberg#67 instead.

Should we be shipping apache/datafusion-iceberg#67 first? or do you think it does not matter much.

Comment on lines +82 to +86
let mut statistics = stats_from_snapshot(table_snapshot.clone(), &scan.schema)?;
if let Some(column_stats) = column_stats {
statistics.column_statistics = column_stats;
}
let mut builder = table.scan();

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.

If we've done eager file planning, the stats should not be coming from the table snapshot no? I imagine they'll not take pruning into account.

Is this what you had in mind for #787? probably worth a TODO comment linking the other PR if that's the case.

Comment thread iceberg/tests/assigned_file_routing.rs Outdated
file.bytes()
.fold(0usize, |sum, byte| sum.wrapping_add(usize::from(byte)))
}) % urls.len();
let response = match ev.dialer.dial(urls[index].clone()).await {

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.

The ok_or_some_err!() macro can help you reduce some boilerplate here.

Comment thread iceberg/tests/assigned_file_routing.rs Outdated
use datafusion_proto::protobuf::PhysicalPlanNode;

#[tokio::test]
async fn routing_inspects_assigned_files_in_the_worker_facing_plan() -> Result<()> {

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.

These tests are a bit hard to read. Integration tests serve also as live documentation, humans should be able to take a quick glance at a test and immediately know what's happening.

Typically we heavy really on insta and plan snapshots for this: they do not only capture behavior, they are also very easy to review by humans.

I remember giving similar feedback when the task routing API was introduced, and this was the result:

https://github.com/datafusion-contrib/datafusion-distributed/blob/14beca0048b939366e2583522402b604285c6585/tests/task_estimator_test.rs

Comment thread iceberg/tests/assigned_file_routing.rs Outdated
Comment on lines +149 to +151
let proto = PhysicalPlanNode::try_from_physical_plan(Arc::clone(plan), &codec)?;
let decoded = proto.try_into_physical_plan(ev.task_ctx, &codec)?;
let restored = iceberg_source(&decoded).unwrap();

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.

Why this proto roundtrip here? this test overall look very far from how someone would use the project in their codebase.

The challenge with having so much synthetic code during integration testing is that they stop resembling real scenarios that people would typically find in their codebases, and a good part of the coverage does not go to testing real scenarios, it goes to testing the synthetic test preparation code instead.

Same as above, if you find ways of expressing these test with something closer how people would use the project, that'd be great.

Comment thread iceberg/tests/file_planning.rs Outdated
use prost::Message;

#[tokio::test]
async fn eager_and_lazy_scans_return_the_same_rows_and_statistics() -> Result<()> {

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.

🤔 is this something we want to assert? I'd expect them to not return the same stats, as eager file planning stats should be better.

Comment thread iceberg/tests/file_planning.rs Outdated
Comment on lines +43 to +45
let lazy = harness(false, true).await?;
let eager = harness(true, true).await?;
let query = "SELECT pickup_date, COUNT(*) AS trips FROM taxi \

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.

Part of the charm of the IcebergTestHarness is that you should not need helpers for building the harness, and the harness configuration is explicit during the build process.

For example, here, I'm unable to understand what this test is doing without navigating to the harness() function and reading it first.

For example, with something like this:

        let eager = IcebergTestHarness::builder()
            .with_workers(2)
            .with_column_stats_enabled()
            .with_eager_file_planning()
            .build()
            .await?;

        let lazy = IcebergTestHarness::builder()
            .with_workers(2)
            .with_column_stats_enabled()
            .build()
            .await?;

It's very easy for a human to know what's the testing setup without doing any code jump. It's more verbose, but the verbosity here plays in the reviewers favor, as the cognitive price it takes to read that is less than if they need to jump through the code understanding what the helper function does.

@gabotechs gabotechs left a comment

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.

I see most of my comments have been address or are on the right track, so +1 in this now. Pretty epic work here 👍

Comment thread iceberg/src/lib.rs Outdated
Comment on lines +38 to +40

// Preserve the existing names for lazy scans and their codec.
pub use lazy::{LazyIcebergCodec as IcebergCodec, LazyIcebergDataSource as IcebergDataSource};

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.

It should be fine to just break the API here, and don't think we need to alias this.

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