Skip to content

Commit ff09ec5

Browse files
feat(datafusion): parallel file scanning with eager task bucketing (#44)
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 a24ad5d commit ff09ec5

9 files changed

Lines changed: 867 additions & 93 deletions

File tree

crates/iceberg/src/scan/mod.rs

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

443443
/// Returns an [`ArrowRecordBatchStream`].
444444
pub async fn to_arrow(&self) -> Result<ArrowRecordBatchStream> {
445+
self.to_arrow_from_tasks(self.plan_files().await?)
446+
}
447+
448+
/// Like [`TableScan::to_arrow`], but accepts a caller-supplied
449+
/// [`FileScanTask`] stream instead of running [`TableScan::plan_files`]
450+
/// internally.
451+
///
452+
/// # Correctness
453+
///
454+
/// Tasks must come from a [`TableScan`] with the same projection and
455+
/// filter as `self`: predicates are baked into each task at planning
456+
/// time and are not re-applied here. Reader-side configuration
457+
/// (concurrency, batch size, row-group filtering, row selection) is
458+
/// taken from `self` and may differ from the planning scan.
459+
pub fn to_arrow_from_tasks(&self, tasks: FileScanTaskStream) -> Result<ArrowRecordBatchStream> {
445460
let mut arrow_reader_builder = ArrowReaderBuilder::new(self.file_io.clone())
446461
.with_data_file_concurrency_limit(self.concurrency_limit_data_files)
447462
.with_row_group_filtering_enabled(self.row_group_filtering_enabled)
@@ -451,7 +466,7 @@ impl TableScan {
451466
arrow_reader_builder = arrow_reader_builder.with_batch_size(batch_size);
452467
}
453468

454-
arrow_reader_builder.build().read(self.plan_files().await?)
469+
arrow_reader_builder.build().read(tasks)
455470
}
456471

457472
/// Returns a reference to the column names of the table scan.

0 commit comments

Comments
 (0)