feat(datafusion): Add opt-in eager file scan planning with output partitioning#2671
feat(datafusion): Add opt-in eager file scan planning with output partitioning#2671toutane wants to merge 5 commits into
Conversation
|
👋 cc @mbutrovich @timsaucer |
mbutrovich
left a comment
There was a problem hiding this comment.
First pass, thanks @toutane!
| @@ -146,13 +139,55 @@ impl ExecutionPlan for IcebergTableScan { | |||
| _partition: usize, | |||
There was a problem hiding this comment.
execute's partition argument is read at line 146 (file_task_groups.get(_partition)) and interpolated into the error message at line 148 (partition {_partition} does not exist). The leading underscore signals "unused" to readers and lint, and it leaks into user-facing error text. Rename to partition. (_context at line 140 is still legitimately unused; leave it.)
There was a problem hiding this comment.
Both IcebergTableProvider::scan (:158, used at :178) and IcebergStaticTableProvider::scan (:343, used at :352) pass _state into create_scan_plan, which reads it via enable_eager_scan_planning(state). Rename both to state. Note: _state at :365 (insert_into is genuinely unused; leave it.)
| let target_partitions = target_partitions.max(1); | ||
|
|
||
| let mut groups: Vec<Vec<FileScanTask>> = vec![Vec::new(); target_partitions]; | ||
| for (i, task) in tasks.into_iter().enumerate() { | ||
| groups[i % target_partitions].push(task); | ||
| } | ||
|
|
||
| groups.retain(|group| !group.is_empty()); | ||
| groups |
There was a problem hiding this comment.
When target_partitions > tasks.len() (which the tests deliberately hit with data_file_count + 1), this allocates extra empty vecs and then walks the whole vec to remove them. The result is always min(target_partitions, tasks.len()) non-empty groups, so clamp up front and the retain disappears:
// tasks is non-empty here (early return above handles the empty case)
let target_partitions = target_partitions.max(1).min(tasks.len());
let mut groups: Vec<Vec<FileScanTask>> = vec![Vec::new(); target_partitions];
for (i, task) in tasks.into_iter().enumerate() {
groups[i % target_partitions].push(task);
}
groups| let task_count = file_task_group.len(); | ||
| let tasks: FileScanTaskStream = Box::pin(futures::stream::iter( | ||
| (0..task_count).map(move |idx| Ok(file_task_group[idx].clone())), | ||
| )); |
There was a problem hiding this comment.
task_count is used once; inline it as (0..file_task_group.len()). Heads-up for anyone tempted to go further: file_task_group.iter().cloned().map(Ok) does not compile here, because the stream must be 'static and iter() would borrow the local Arc<[FileScanTask]>. The move-plus-index form is the allocation-free way to clone lazily out of the Arc, so keep that shape; just remove the extra binding.
|
Thanks @mbutrovich, addressed all four comments in d99603f |
ef774e8 to
6a4f1a1
Compare
6a4f1a1 to
be6d53e
Compare
Which issue does this PR close?
What changes are included in this PR?
This PR adds an opt-in eager scan planning path for the DataFusion integration.
When
iceberg.enable_eager_scan_planningis enabled,IcebergTableProvider::scan()plansFileScanTasks during physical planning, groups them across DataFusiontarget_partitions, and exposes the resulting output partition count asUnknownPartitioning(N). Eachexecute(partition)call then reads only the task group assigned to that output partition throughArrowReaderBuilder.The default behavior is unchanged: eager planning is disabled by default, so scans keep the existing lazy single-partition planning path unless explicitly enabled.
This is a narrower split from #2298, focused on file-level eager planning and output partition count reporting. Some of the broader design discussions and review feedback happened in #2298; this PR keeps only the first scoped step and intentionally does not include hash partitioning, size-aware bin-packing, or row-group/sub-file planning.
Trade-offs:
TableProvider::scan(), so it is kept opt-inFileScanTasklevel only. A table with one large file will not benefit from this changeUnknownPartitioning(N), not hash partitioning. This exposes the number of output partitions without claiming stronger partitioning semanticsFollow-up work:
file_size_in_bytes(Plan file scan task according scan file size. #128)Are these changes tested?
Yes. Integration tests cover:
iceberg.enable_eager_scan_planning