toutane commented on code in PR #2671:
URL: https://github.com/apache/iceberg-rust/pull/2671#discussion_r3703807458
##########
crates/integrations/datafusion/src/physical_plan/scan.rs:
##########
@@ -138,18 +137,60 @@ impl ExecutionPlan for IcebergTableScan {
fn execute(
&self,
- _partition: usize,
+ partition: usize,
_context: Arc<TaskContext>,
) -> DFResult<SendableRecordBatchStream> {
- let fut = get_batch_stream(
- self.table.clone(),
- self.snapshot_id,
- self.projection.clone(),
- self.predicates.clone(),
- );
- let stream = futures::stream::once(fut).try_flatten();
-
- // Apply limit if specified
+ let stream: Pin<Box<dyn Stream<Item = DFResult<RecordBatch>> + Send>>
= match &self
+ .file_task_groups
+ {
+ Some(file_task_groups) => {
+ let Some(file_task_group) =
file_task_groups.get(partition).cloned() else {
+ return
Err(datafusion::common::DataFusionError::Internal(format!(
+ "IcebergTableScan partition {partition} does not
exist; scan has {} partitions",
+ file_task_groups.len()
+ )));
+ };
+
+ let tasks: FileScanTaskStream = Box::pin(futures::stream::iter(
+ (0..file_task_group.len()).map(move |idx|
Ok(file_task_group[idx].clone())),
+ ));
+ let stream = build_table_scan(&self.table, &self.scan_config)?
+ .arrow_reader_builder()
+ // Eager planning lets DataFusion drive scan concurrency
via output
+ // partitions. Match DataFusion's FileStream model, where
each
+ // output partition owns one ScanState; keep one data file
in
+ // flight per output partition here.
+ //
https://github.com/apache/datafusion/blob/ad8e7b7f2babe3fcddc3a4f9b5cd1ac0d1b16ad9/datafusion/datasource/src/file_stream/scan_state.rs#L42-L43
+ .with_data_file_concurrency_limit(1)
+ .build()
+ // TODO: Avoid cloning FileScanTasks here once ArrowReader
can accept shared tasks.
+ .read(tasks)
Review Comment:
Hey @mbutrovich, thanks for the thorough review and for staying with this PR.
Commit 9fe0442 addresses it by threading the planning `TableScan` through to
`execute()`.
`scan_planning.rs` gains an `EagerScanPlan` struct holding the
`Arc<TableScan>` that planned the tasks, alongside those tasks grouped per
output partition. `plan_file_task_groups` is renamed `plan_eager_scan` and
returns that struct instead of a bare `Vec<Vec<FileScanTask>>`. `TableScan`
isn't `Clone`, but `arrow_reader_builder(&self)` and `plan_files(&self)` both
take `&self`, so an `Arc` is enough and no change to the `iceberg` crate is
needed.
`IcebergTableScan` stores it as `eager_plan: Option<EagerScanPlan>`,
replacing the previous `Option<Vec<Arc<[FileScanTask]>>>`. A single `Option`
keeps the "tasks planned = the scan that planned them" invariant in the type
instead of spread across two fields; `None` still selects the lazy path.
To answer your question directly: it is the `TableScan` instance that is
reused, not the `ArrowReaderBuilder`. Each `execute(partition)` still calls
`arrow_reader_builder()` on that shared scan and builds its own reader - that
part is only a handful of `Arc` clones.
On the gap you flagged, both paths still go through `build_table_scan` and
`TableScan::arrow_reader_builder()`. `build_table_scan` has two callers -
`plan_eager_scan` at plan time, and the lazy arm of `execute()` at execute time
- so it stays the single source for snapshot, projection, predicate and reader
settings. The eager reader is now derived from the very instance that planned
the tasks, which makes that equivalence structural rather than a calling
convention. `test_multi_partition_scan_matches_single_partition_results` still
pins it behaviorally across lazy, eager with 1 partition, and eager with 4
partitions.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
