refactor: TableScan's plan_files method now operates on manifest files and manifest data files concurrently rather than sequentially

Author: sdd

crates/iceberg/src/error.rs

@@ -333,6 +333,12 @@ define_from_err!(
 
 define_from_err!(std::io::Error, ErrorKind::Unexpected, "IO Operation failed");
 
+define_from_err!(
+    futures::channel::mpsc::SendError,
+    ErrorKind::Unexpected,
+    "Failed to send a message to a channel"
+);
+
 /// Helper macro to check arguments.
 ///
 ///

crates/iceberg/src/scan.rs

@@ -24,18 +24,24 @@ use crate::expr::visitors::manifest_evaluator::ManifestEvaluator;
 use crate::expr::{Bind, BoundPredicate, Predicate};
 use crate::io::FileIO;
 use crate::spec::{
-    DataContentType, ManifestContentType, ManifestEntryRef, ManifestFile, Schema, SchemaRef,
-    SnapshotRef, TableMetadataRef,
+    DataContentType, ManifestContentType, ManifestEntryRef, ManifestFile, ManifestList, Schema,
+    SchemaRef, SnapshotRef, TableMetadataRef,
 };
 use crate::table::Table;
 use crate::{Error, ErrorKind, Result};
 use arrow_array::RecordBatch;
-use async_stream::try_stream;
+use futures::channel::mpsc::{channel, Sender};
 use futures::stream::BoxStream;
-use futures::StreamExt;
+use futures::{SinkExt, StreamExt, TryStreamExt};
 use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 use std::sync::Arc;
+use tokio::spawn;
+
+const CHANNEL_BUFFER_SIZE: usize = 10;
+const CONCURRENCY_LIMIT_MANIFEST_FILES: usize = 10;
+const CONCURRENCY_LIMIT_MANIFEST_ENTRIES: usize = 10;
+const SUPPORTED_MANIFEST_FILE_CONTENT_TYPES: [ManifestContentType; 1] = [ManifestContentType::Data];
 
 /// A stream of [`FileScanTask`].
 pub type FileScanTaskStream = BoxStream<'static, Result<FileScanTask>>;
@@ -206,76 +212,20 @@ impl TableScan {
             self.case_sensitive,
         )?;
 
-        let mut partition_filter_cache = PartitionFilterCache::new();
-        let mut manifest_evaluator_cache = ManifestEvaluatorCache::new();
-
-        Ok(try_stream! {
-            let manifest_list = context
-                .snapshot
-                .load_manifest_list(&context.file_io, &context.table_metadata)
-                .await?;
-
-            for entry in manifest_list.entries() {
-                if !Self::content_type_is_data(entry) {
-                    continue;
-                }
+        let (sender, receiver) = channel(CHANNEL_BUFFER_SIZE);
 
-                let partition_spec_id = entry.partition_spec_id;
+        let manifest_list = context
+            .snapshot
+            .load_manifest_list(&context.file_io, &context.table_metadata)
+            .await?;
 
-                let partition_filter = partition_filter_cache.get(
-                    partition_spec_id,
-                    &context,
-                )?;
+        spawn(async move {
+            let _ = ConcurrentFileScanStreamContext::new(context, sender)
+                .run(manifest_list)
+                .await;
+        });
 
-                if let Some(partition_filter) = partition_filter {
-                    let manifest_evaluator = manifest_evaluator_cache.get(
-                        partition_spec_id,
-                        partition_filter,
-                    );
-
-                    if !manifest_evaluator.eval(entry)? {
-                        continue;
-                    }
-                }
-
-                let manifest = entry.load_manifest(&context.file_io).await?;
-                let mut manifest_entries_stream =
-                    futures::stream::iter(manifest.entries().iter().filter(|e| e.is_alive()));
-
-                while let Some(manifest_entry) = manifest_entries_stream.next().await {
-                    // TODO: Apply ExpressionEvaluator
-
-                    if let Some(bound_predicate) = context.bound_filter() {
-                        // reject any manifest entries whose data file's metrics don't match the filter.
-                        if !InclusiveMetricsEvaluator::eval(
-                            bound_predicate,
-                            manifest_entry.data_file(),
-                            false
-                        )? {
-                            continue;
-                        }
-                    }
-
-                    match manifest_entry.content_type() {
-                        DataContentType::EqualityDeletes | DataContentType::PositionDeletes => {
-                            yield Err(Error::new(
-                                ErrorKind::FeatureUnsupported,
-                                "Delete files are not supported yet.",
-                            ))?;
-                        }
-                        DataContentType::Data => {
-                            let scan_task: Result<FileScanTask> = Ok(FileScanTask {
-                                data_manifest_entry: manifest_entry.clone(),
-                                start: 0,
-                                length: manifest_entry.file_size_in_bytes(),
-                            });
-                            yield scan_task?;
-                        }
-                    }
-                }
-            }
-        }
-        .boxed())
+        return Ok(receiver.boxed());
     }
 
     /// Returns an [`ArrowRecordBatchStream`].
@@ -333,13 +283,172 @@ impl TableScan {
 
         arrow_reader_builder.build().read(self.plan_files().await?)
     }
+}
+
+#[derive(Debug)]
+struct ConcurrentFileScanStreamContext {
+    context: Arc<FileScanStreamContext>,
+    sender: Sender<Result<FileScanTask>>,
+    manifest_evaluator_cache: ManifestEvaluatorCache,
+    partition_filter_cache: PartitionFilterCache,
+}
+
+impl ConcurrentFileScanStreamContext {
+    fn new(context: FileScanStreamContext, sender: Sender<Result<FileScanTask>>) -> Self {
+        ConcurrentFileScanStreamContext {
+            context: Arc::new(context),
+            sender,
+            manifest_evaluator_cache: ManifestEvaluatorCache::new(),
+            partition_filter_cache: PartitionFilterCache::new(),
+        }
+    }
+
+    async fn run(&mut self, manifest_list: ManifestList) -> Result<()> {
+        let file_io = self.context.file_io.clone();
+        let sender = self.sender.clone();
+        let context = self.context.clone();
+
+        // This whole Vec-and-for-loop approach feels sub-optimally structured.
+        // I've tried structuring this in multiple ways but run into
+        // issues with ownership. Ideally I'd like to structure this
+        // with a functional programming approach: extracting
+        // sections 1, 2, and 3 out into different methods on Self,
+        // and then use some iterator combinators to orchestrate it all.
+        // Section 1 is pretty trivially refactorable into a static method
+        // that can be used in a closure that can be used with Iterator::filter.
+        // Similarly, section 3 seems easily factor-able into a method that can
+        // be passed into Iterator::map.
+        // Section 2 turns out trickier - we want to exit the entire `run` method early
+        // if the eval fails, and filter out any manifest_files from the iterator / stream
+        // if the eval succeeds but returns true. We bump into ownership issues due
+        // to needing to pass mut self as the caches need to be able to mutate.
+        // Section 3 runs into ownership issues when trying to refactor its closure to be
+        // a static or non-static method.
+
+        // 1
+        let filtered_manifest_files = manifest_list
+            .entries()
+            .iter()
+            .filter(Self::reject_unsupported_content_types);
+
+        // 2
+        let mut filtered_manifest_files2 = vec![];
+        for manifest_file in filtered_manifest_files {
+            if !self.apply_evaluator(manifest_file)? {
+                continue;
+            }
+
+            filtered_manifest_files2.push(manifest_file);
+        }
+
+        // 3
+        let filtered_manifest_files = filtered_manifest_files2.into_iter().map(|manifest_file| {
+            Ok((
+                manifest_file,
+                file_io.clone(),
+                sender.clone(),
+                context.clone(),
+            ))
+        });
+
+        futures::stream::iter(filtered_manifest_files)
+            .try_for_each_concurrent(
+                CONCURRENCY_LIMIT_MANIFEST_FILES,
+                Self::process_manifest_file,
+            )
+            .await
+    }
+
+    fn reject_unsupported_content_types(manifest_file: &&ManifestFile) -> bool {
+        SUPPORTED_MANIFEST_FILE_CONTENT_TYPES.contains(&manifest_file.content)
+    }
+
+    fn apply_evaluator(&mut self, manifest_file: &ManifestFile) -> Result<bool> {
+        let partition_spec_id = manifest_file.partition_spec_id;
+
+        let partition_filter = self
+            .partition_filter_cache
+            .get(partition_spec_id, &self.context)?;
+
+        if let Some(partition_filter) = partition_filter {
+            let manifest_evaluator = self
+                .manifest_evaluator_cache
+                .get(partition_spec_id, partition_filter);
+
+            if !manifest_evaluator.eval(manifest_file)? {
+                return Ok(false);
+            }
+        }
+
+        Ok(true)
+    }
+
+    async fn process_manifest_file(
+        manifest_and_file_io_and_sender: (
+            &ManifestFile,
+            FileIO,
+            Sender<Result<FileScanTask>>,
+            Arc<FileScanStreamContext>,
+        ),
+    ) -> Result<()> {
+        let (manifest_file, file_io, sender, context) = manifest_and_file_io_and_sender;
+
+        let manifest = manifest_file.load_manifest(&file_io).await?;
+
+        let manifest_entries = manifest
+            .entries()
+            .iter()
+            .filter(|x| x.is_alive())
+            .map(|manifest_entry| Ok((manifest_entry, sender.clone(), context.clone())));
+
+        futures::stream::iter(manifest_entries)
+            .try_for_each_concurrent(
+                CONCURRENCY_LIMIT_MANIFEST_ENTRIES,
+                Self::process_manifest_entry,
+            )
+            .await
+    }
+
+    async fn process_manifest_entry(
+        manifest_entry_and_sender: (
+            &ManifestEntryRef,
+            Sender<Result<FileScanTask>>,
+            Arc<FileScanStreamContext>,
+        ),
+    ) -> Result<()> {
+        let (manifest_entry, mut sender, context) = manifest_entry_and_sender;
+
+        if !matches!(manifest_entry.content_type(), DataContentType::Data) {
+            return Err(Error::new(
+                ErrorKind::FeatureUnsupported,
+                format!(
+                    "Files of type '{:?}' are not supported yet.",
+                    manifest_entry.content_type()
+                ),
+            ));
+        }
+
+        if let Some(bound_predicate) = context.bound_filter() {
+            // reject any manifest entries whose data file's metrics don't match the filter.
+            if !InclusiveMetricsEvaluator::eval(bound_predicate, manifest_entry.data_file(), false)?
+            {
+                return Ok(());
+            }
+        }
+
+        // TODO: Apply ExpressionEvaluator
+
+        Ok(sender
+            .send(Ok(Self::manifest_entry_to_file_scan_task(manifest_entry)))
+            .await?)
+    }
 
-    /// Checks whether the [`ManifestContentType`] is `Data` or not.
-    fn content_type_is_data(entry: &ManifestFile) -> bool {
-        if let ManifestContentType::Data = entry.content {
-            return true;
+    fn manifest_entry_to_file_scan_task(manifest_entry: &ManifestEntryRef) -> FileScanTask {
+        FileScanTask {
+            data_manifest_entry: manifest_entry.clone(),
+            start: 0,
+            length: manifest_entry.file_size_in_bytes(),
         }
-        false
     }
 }