Repository navigation
feat: plan eager Iceberg scans with source-owned files - #803
alexanderbianchi wants to merge 8 commits into
Conversation
a07bff5 to
093942f
Compare
093942f to
fcf6f1e
Compare
fcf6f1e to
9b37c33
Compare
74e158a to
cb1a7e8
Compare
cb1a7e8 to
d6c8b6a
Compare
gabotechs
left a comment
There was a problem hiding this comment.
👍 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!
| 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 { |
There was a problem hiding this comment.
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.
| ) -> Result<Arc<dyn ExecutionPlan>> { | ||
| let output_schema = match opts.projection { | ||
| None => schema.clone(), | ||
| Some(projection) => Arc::new(schema.project(projection).unwrap()), |
There was a problem hiding this comment.
| Some(projection) => Arc::new(schema.project(projection).unwrap()), | |
| Some(projection) => Arc::new(schema.project(projection)?), |
| partitions[file_partition(index, 1, partitioning.partition_count())].push(index); | ||
| } | ||
| return Ok(DataSourceExec::from_data_source( | ||
| AssignedIcebergDataSource { |
There was a problem hiding this comment.
🤔 I'm thinking about the naming, what about renaming them:
EagerIcebergDataSourceLazyIcebergDataSource
?
Those are the terms we use typically while talking about it.
| } else { | ||
| None | ||
| }; | ||
| if IcebergConfig::from_task_context(&context).eager_file_planning { |
There was a problem hiding this comment.
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.
| 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(); |
There was a problem hiding this comment.
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.
| file.bytes() | ||
| .fold(0usize, |sum, byte| sum.wrapping_add(usize::from(byte))) | ||
| }) % urls.len(); | ||
| let response = match ev.dialer.dial(urls[index].clone()).await { |
There was a problem hiding this comment.
The ok_or_some_err!() macro can help you reduce some boilerplate here.
| use datafusion_proto::protobuf::PhysicalPlanNode; | ||
|
|
||
| #[tokio::test] | ||
| async fn routing_inspects_assigned_files_in_the_worker_facing_plan() -> Result<()> { |
There was a problem hiding this comment.
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:
| 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(); |
There was a problem hiding this comment.
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.
| use prost::Message; | ||
|
|
||
| #[tokio::test] | ||
| async fn eager_and_lazy_scans_return_the_same_rows_and_statistics() -> Result<()> { |
There was a problem hiding this comment.
🤔 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.
| let lazy = harness(false, true).await?; | ||
| let eager = harness(true, true).await?; | ||
| let query = "SELECT pickup_date, COUNT(*) AS trips FROM taxi \ |
There was a problem hiding this comment.
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
left a comment
There was a problem hiding this comment.
I see most of my comments have been address or are on the right track, so +1 in this now. Pretty epic work here 👍
|
|
||
| // Preserve the existing names for lazy scans and their codec. | ||
| pub use lazy::{LazyIcebergCodec as IcebergCodec, LazyIcebergDataSource as IcebergDataSource}; |
There was a problem hiding this comment.
It should be fine to just break the API here, and don't think we need to alias this.
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 selectedFileScanTasks on anEagerIcebergDataSource. 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'sFileIO; 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.
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.FileScanConfigpath, the scale-up handler creates task-specificDistributedLeafExecvariants.RouteTaskEvent::task_specialized_planexposes the selected variant before placement. No second plan view or work-unit-feed extension is needed.FileScanConfigor supplying a placeholder object-store URL.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 underlazy/, with shared reading incommon/scan_config.rs. The oldIcebergDataSourceandIcebergCodecexports are removed; consumers must useLazyIcebergDataSourceandLazyIcebergCodec, 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.