Refactor concurrency of incremental streams (#21)

Author: gbrgr

crates/iceberg/src/arrow/incremental.rs

@@ -34,7 +34,7 @@ use crate::metadata_columns::{RESERVED_FIELD_ID_UNDERSCORE_POS, row_pos_field};
 use crate::runtime::spawn;
 use crate::scan::ArrowRecordBatchStream;
 use crate::scan::incremental::{
-    AppendedFileScanTask, IncrementalFileScanTask, IncrementalFileScanTaskStream,
+    AppendedFileScanTask, DeleteScanTask, IncrementalFileScanTaskStreams,
 };
 use crate::{Error, ErrorKind, Result};
 
@@ -57,120 +57,6 @@ pub type CombinedIncrementalBatchRecordStream =
 /// Stream type for obtaining a separate stream of appended and deleted record batches.
 pub type UnzippedIncrementalBatchRecordStream = (ArrowRecordBatchStream, ArrowRecordBatchStream);
 
-impl StreamsInto<ArrowReader, CombinedIncrementalBatchRecordStream>
-    for IncrementalFileScanTaskStream
-{
-    /// Takes a stream of `IncrementalFileScanTasks` and reads all the files. Returns a
-    /// stream of Arrow `RecordBatch`es containing the data from the files.
-    fn stream(self, reader: ArrowReader) -> Result<CombinedIncrementalBatchRecordStream> {
-        let (appends, deletes) =
-            StreamsInto::<ArrowReader, UnzippedIncrementalBatchRecordStream>::stream(self, reader)?;
-
-        let left = appends.map(|res| res.map(|batch| (IncrementalBatchType::Append, batch)));
-        let right = deletes.map(|res| res.map(|batch| (IncrementalBatchType::Delete, batch)));
-
-        Ok(Box::pin(select(left, right)) as CombinedIncrementalBatchRecordStream)
-    }
-}
-
-impl StreamsInto<ArrowReader, UnzippedIncrementalBatchRecordStream>
-    for IncrementalFileScanTaskStream
-{
-    /// Takes a stream of `IncrementalFileScanTasks` and reads all the files. Returns two
-    /// separate streams of Arrow `RecordBatch`es containing appended data and deleted records.
-    fn stream(self, reader: ArrowReader) -> Result<UnzippedIncrementalBatchRecordStream> {
-        let (appends_tx, appends_rx) =
-            channel::<Result<RecordBatch>>(reader.concurrency_limit_data_files);
-        let (deletes_tx, deletes_rx) =
-            channel::<Result<RecordBatch>>(reader.concurrency_limit_data_files);
-
-        let batch_size = reader.batch_size;
-
-        spawn(async move {
-            let _ = self
-                .try_for_each_concurrent(reader.concurrency_limit_data_files, |task| {
-                    let file_io = reader.file_io.clone();
-                    match task {
-                        IncrementalFileScanTask::Append(append_task) => {
-                            let appends_tx = appends_tx.clone();
-                            Box::pin(async move {
-                                spawn(async move {
-                                    let record_batch_stream = process_incremental_append_task(
-                                        append_task,
-                                        batch_size,
-                                        file_io,
-                                    )
-                                    .await;
-
-                                    process_record_batch_stream(
-                                        record_batch_stream,
-                                        appends_tx,
-                                        "failed to read appended record batch",
-                                    )
-                                    .await;
-                                });
-                                Ok(())
-                            })
-                                as Pin<Box<dyn futures::Future<Output = Result<()>> + Send>>
-                        }
-                        IncrementalFileScanTask::Delete(deleted_file_task) => {
-                            let deletes_tx = deletes_tx.clone();
-                            Box::pin(async move {
-                                spawn(async move {
-                                    let file_path = deleted_file_task.data_file_path().to_string();
-                                    let total_records =
-                                        deleted_file_task.base.record_count.unwrap_or(0);
-
-                                    let record_batch_stream = process_incremental_deleted_file_task(
-                                        file_path,
-                                        total_records,
-                                        batch_size,
-                                    );
-
-                                    process_record_batch_stream(
-                                        record_batch_stream,
-                                        deletes_tx,
-                                        "failed to read deleted file record batch",
-                                    )
-                                    .await;
-                                });
-                                Ok(())
-                            })
-                                as Pin<Box<dyn futures::Future<Output = Result<()>> + Send>>
-                        }
-                        IncrementalFileScanTask::PositionalDeletes(file_path, delete_vector) => {
-                            let deletes_tx = deletes_tx.clone();
-                            Box::pin(async move {
-                                spawn(async move {
-                                    let record_batch_stream = process_incremental_delete_task(
-                                        file_path,
-                                        delete_vector,
-                                        batch_size,
-                                    );
-
-                                    process_record_batch_stream(
-                                        record_batch_stream,
-                                        deletes_tx,
-                                        "failed to read deleted record batch",
-                                    )
-                                    .await;
-                                });
-                                Ok(())
-                            })
-                                as Pin<Box<dyn futures::Future<Output = Result<()>> + Send>>
-                        }
-                    }
-                })
-                .await;
-        });
-
-        Ok((
-            Box::pin(appends_rx) as ArrowRecordBatchStream,
-            Box::pin(deletes_rx) as ArrowRecordBatchStream,
-        ))
-    }
-}
-
 async fn process_incremental_append_task(
     task: AppendedFileScanTask,
     batch_size: Option<usize>,
@@ -330,3 +216,117 @@ fn process_incremental_deleted_file_task(
 
     Ok(Box::pin(stream) as ArrowRecordBatchStream)
 }
+
+impl StreamsInto<ArrowReader, CombinedIncrementalBatchRecordStream>
+    for IncrementalFileScanTaskStreams
+{
+    /// Takes separate streams of appended and deleted file scan tasks and reads all the files.
+    /// Returns a combined stream of Arrow `RecordBatch`es containing the data from the files.
+    fn stream(self, reader: ArrowReader) -> Result<CombinedIncrementalBatchRecordStream> {
+        let (appends, deletes) =
+            StreamsInto::<ArrowReader, UnzippedIncrementalBatchRecordStream>::stream(self, reader)?;
+
+        let left = appends.map(|res| res.map(|batch| (IncrementalBatchType::Append, batch)));
+        let right = deletes.map(|res| res.map(|batch| (IncrementalBatchType::Delete, batch)));
+
+        Ok(Box::pin(select(left, right)) as CombinedIncrementalBatchRecordStream)
+    }
+}
+
+impl StreamsInto<ArrowReader, UnzippedIncrementalBatchRecordStream>
+    for IncrementalFileScanTaskStreams
+{
+    /// Takes separate streams of appended and deleted file scan tasks and reads all the files.
+    /// Returns two separate streams of Arrow `RecordBatch`es containing appended data and deleted records.
+    fn stream(self, reader: ArrowReader) -> Result<UnzippedIncrementalBatchRecordStream> {
+        let (appends_tx, appends_rx) =
+            channel::<Result<RecordBatch>>(reader.concurrency_limit_data_files);
+        let (deletes_tx, deletes_rx) =
+            channel::<Result<RecordBatch>>(reader.concurrency_limit_data_files);
+
+        let batch_size = reader.batch_size;
+
+        let (append_stream, delete_stream) = self;
+
+        // Process append tasks
+        let file_io = reader.file_io.clone();
+        spawn(async move {
+            let _ = append_stream
+                .try_for_each_concurrent(reader.concurrency_limit_data_files, |append_task| {
+                    let file_io = file_io.clone();
+                    let appends_tx = appends_tx.clone();
+                    async move {
+                        spawn(async move {
+                            let record_batch_stream =
+                                process_incremental_append_task(append_task, batch_size, file_io)
+                                    .await;
+
+                            process_record_batch_stream(
+                                record_batch_stream,
+                                appends_tx,
+                                "failed to read appended record batch",
+                            )
+                            .await;
+                        });
+                        Ok(())
+                    }
+                })
+                .await;
+        });
+
+        // Process delete tasks
+        spawn(async move {
+            let _ = delete_stream
+                .try_for_each_concurrent(reader.concurrency_limit_data_files, |delete_task| {
+                    let deletes_tx = deletes_tx.clone();
+                    async move {
+                        match delete_task {
+                            DeleteScanTask::DeletedFile(deleted_file_task) => {
+                                spawn(async move {
+                                    let file_path = deleted_file_task.data_file_path().to_string();
+                                    let total_records =
+                                        deleted_file_task.base.record_count.unwrap_or(0);
+
+                                    let record_batch_stream = process_incremental_deleted_file_task(
+                                        file_path,
+                                        total_records,
+                                        batch_size,
+                                    );
+
+                                    process_record_batch_stream(
+                                        record_batch_stream,
+                                        deletes_tx,
+                                        "failed to read deleted file record batch",
+                                    )
+                                    .await;
+                                });
+                            }
+                            DeleteScanTask::PositionalDeletes(file_path, delete_vector) => {
+                                spawn(async move {
+                                    let record_batch_stream = process_incremental_delete_task(
+                                        file_path,
+                                        delete_vector,
+                                        batch_size,
+                                    );
+
+                                    process_record_batch_stream(
+                                        record_batch_stream,
+                                        deletes_tx,
+                                        "failed to read deleted record batch",
+                                    )
+                                    .await;
+                                });
+                            }
+                        }
+                        Ok(())
+                    }
+                })
+                .await;
+        });
+
+        Ok((
+            Box::pin(appends_rx) as ArrowRecordBatchStream,
+            Box::pin(deletes_rx) as ArrowRecordBatchStream,
+        ))
+    }
+}

crates/iceberg/src/scan/incremental/mod.rs

@@ -426,7 +426,8 @@ impl IncrementalTableScan {
     }
 
     /// Plans the files to be read in this incremental table scan.
-    pub async fn plan_files(&self) -> Result<IncrementalFileScanTaskStream> {
+    /// Returns separate streams for appended and deleted records.
+    pub async fn plan_files(&self) -> Result<IncrementalFileScanTaskStreams> {
         let concurrency_limit_manifest_files = self.concurrency_limit_manifest_files;
         let concurrency_limit_manifest_entries = self.concurrency_limit_manifest_entries;
 
@@ -584,21 +585,22 @@ impl IncrementalTableScan {
             }
         }
 
-        // Build final task list with net changes
+        // Build final task lists with net changes
         // We filter out tasks for files that appear in both sets (cancelled out)
-        let mut final_tasks: Vec<IncrementalFileScanTask> = Vec::new();
+        let mut final_append_tasks: Vec<AppendedFileScanTask> = Vec::new();
+        let mut final_delete_tasks: Vec<DeleteScanTask> = Vec::new();
 
         // Add net append tasks (only files not in deleted_files)
         for append_task in append_tasks {
             if !deleted_files.contains(append_task.data_file_path()) {
-                final_tasks.push(IncrementalFileScanTask::Append(append_task));
+                final_append_tasks.push(append_task);
             }
         }
 
         // Add net delete tasks (only files not in appended_files)
         for delete_task in delete_tasks {
             if !appended_files.contains(delete_task.data_file_path()) {
-                final_tasks.push(IncrementalFileScanTask::Delete(delete_task));
+                final_delete_tasks.push(DeleteScanTask::DeletedFile(delete_task));
             }
         }
 
@@ -641,14 +643,13 @@ impl IncrementalTableScan {
                         .with_source(e)
                 })?;
 
-            let positional_delete_task =
-                IncrementalFileScanTask::PositionalDeletes(path, delete_vector_inner);
-            final_tasks.push(positional_delete_task);
+            final_delete_tasks.push(DeleteScanTask::PositionalDeletes(path, delete_vector_inner));
         }
 
-        // We actually would not need a stream here, but we can keep it compatible with
-        // other scan types.
-        Ok(futures::stream::iter(final_tasks).map(Ok).boxed())
+        let append_stream = futures::stream::iter(final_append_tasks).map(Ok).boxed();
+        let delete_stream = futures::stream::iter(final_delete_tasks).map(Ok).boxed();
+
+        Ok((append_stream, delete_stream))
     }
 
     /// Returns an [`CombinedIncrementalBatchRecordStream`] for this incremental table scan.

crates/iceberg/src/scan/incremental/task.rs

@@ -111,8 +111,21 @@ impl DeletedFileScanTask {
     }
 }
 
-/// The stream of incremental file scan tasks.
-pub type IncrementalFileScanTaskStream = BoxStream<'static, Result<IncrementalFileScanTask>>;
+/// The streams of appended and deleted file scan tasks.
+pub type IncrementalFileScanTaskStreams = (
+    BoxStream<'static, Result<AppendedFileScanTask>>,
+    BoxStream<'static, Result<DeleteScanTask>>,
+);
+
+/// A delete scan task, which can be a deleted data file or positional deletes.
+#[derive(Debug, Clone)]
+pub enum DeleteScanTask {
+    /// A deleted data file.
+    DeletedFile(DeletedFileScanTask),
+    /// Positional deletes (deleted records of a data file). First argument is the file path,
+    /// second the delete vector.
+    PositionalDeletes(String, DeleteVector),
+}
 
 /// An incremental file scan task, which can be an appended data file, deleted data file,
 /// or positional deletes.