Skip to content

Commit b8fc79d

Browse files
feat(datafusion): parallel file scanning with eager task bucketing (#43)
Port of apache#2298 onto the spiceai-0.9.0 fork. IcebergTableProvider::scan() now plans files eagerly and distributes FileScanTasks into min(target_partitions, n_files) buckets, one bucket per DataFusion partition, so file reads are scheduled concurrently instead of streaming through a single UnknownPartitioning(1) partition. When the table is identity-partitioned (single spec, supported column types, partition columns projected) the scan declares Partitioning::Hash so downstream joins/aggregates can skip a RepartitionExec. - TableScan::to_arrow_from_tasks: replay pre-collected FileScanTasks through the Arrow reader; preserves the spice fork's ArrowReaderBuilder (file_io, runtime) signature and row-selection config. - IcebergTableScan gains new_with_tasks (eager) alongside new (lazy, used by IcebergStaticTableProvider); execute(i) streams buckets[i]. Constructors made pub; with_new_children now errors on children. - New table/bucketing.rs: identity-hash bucketing via REPARTITION_RANDOM_STATE + create_hashes, fallback to data_file_path. - Spice limit pushdown preserved: with_limit threaded into the planning builder and build_table_scan. - Drop the unused convert_filters_to_predicate re-export.
1 parent b69ab70 commit b8fc79d

11 files changed

Lines changed: 868 additions & 95 deletions

File tree

crates/iceberg/src/scan/mod.rs

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -472,6 +472,21 @@ impl TableScan {
472472

473473
/// Returns an [`ArrowRecordBatchStream`].
474474
pub async fn to_arrow(&self) -> Result<ArrowRecordBatchStream> {
475+
self.to_arrow_from_tasks(self.plan_files().await?)
476+
}
477+
478+
/// Like [`TableScan::to_arrow`], but accepts a caller-supplied
479+
/// [`FileScanTask`] stream instead of running [`TableScan::plan_files`]
480+
/// internally.
481+
///
482+
/// # Correctness
483+
///
484+
/// Tasks must come from a [`TableScan`] with the same projection and
485+
/// filter as `self`: predicates are baked into each task at planning
486+
/// time and are not re-applied here. Reader-side configuration
487+
/// (concurrency, batch size, row-group filtering, row selection) is
488+
/// taken from `self` and may differ from the planning scan.
489+
pub fn to_arrow_from_tasks(&self, tasks: FileScanTaskStream) -> Result<ArrowRecordBatchStream> {
475490
let mut arrow_reader_builder =
476491
ArrowReaderBuilder::new(self.file_io.clone(), self.runtime.clone())
477492
.with_data_file_concurrency_limit(self.concurrency_limit_data_files)
@@ -484,7 +499,7 @@ impl TableScan {
484499

485500
arrow_reader_builder
486501
.build()
487-
.read(self.plan_files().await?)
502+
.read(tasks)
488503
.map(|result| result.stream())
489504
}
490505

crates/integrations/datafusion/src/physical_plan/mod.rs

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,5 @@ pub(crate) mod write;
2626

2727
pub(crate) const DATA_FILES_COL_NAME: &str = "data_files";
2828

29-
pub use expr_to_predicate::convert_filters_to_predicate;
3029
pub use project::project_with_partition;
3130
pub use scan::IcebergTableScan;

0 commit comments

Comments
 (0)