toutane commented on code in PR #2671:
URL: https://github.com/apache/iceberg-rust/pull/2671#discussion_r3681266502
##########
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)
+ .map_err(to_datafusion_error)?
+ .stream()
+ .map_err(to_datafusion_error);
+
+ Box::pin(stream)
+ }
+ None => {
+ let table = self.table.clone();
+ let scan_config = self.scan_config.clone();
+ let fut = async move {
+ let table_scan = build_table_scan(&table, &scan_config)?;
+ let stream = table_scan
+ .to_arrow()
+ .await
+ .map_err(to_datafusion_error)?
+ .map_err(to_datafusion_error);
+ Ok::<_, datafusion::common::DataFusionError>(stream)
+ };
+
+ Box::pin(futures::stream::once(fut).try_flatten())
+ }
+ };
+
+ // Apply a scan-partition bound if specified. In eager planning this
is only
Review Comment:
- 5ae5544b7f12c1399bad9a855085349ca8c99b90: add a targeted eager
multi-partition `SELECT ... LIMIT k` integration test that verifies the scan
has multiple partitions, the optimized plan keeps a global bound above the
scan, and the collected result contains exactly `k` valid rows .
- 1082c50de3bf9ef01b765d4422b9e7e0ef0f1c6d: add a separate `ORDER BY foo1
LIMIT k` test to cover the deterministic ordered case end to end.
--
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]
