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

Author: sdd

crates/iceberg/src/arrow/reader.rs

@@ -52,7 +52,7 @@ pub struct ArrowReaderBuilder {
     field_ids: Vec<usize>,
     file_io: FileIO,
     schema: SchemaRef,
-    predicates: Option<BoundPredicate>,
+    predicate: Option<BoundPredicate>,
 }
 
 impl ArrowReaderBuilder {
@@ -63,7 +63,7 @@ impl ArrowReaderBuilder {
             field_ids: vec![],
             file_io,
             schema,
-            predicates: None,
+            predicate: None,
         }
     }
 
@@ -82,7 +82,7 @@ impl ArrowReaderBuilder {
 
     /// Sets the predicates to apply to the scan.
     pub fn with_predicates(mut self, predicates: BoundPredicate) -> Self {
-        self.predicates = Some(predicates);
+        self.predicate = Some(predicates);
         self
     }
 
@@ -93,7 +93,7 @@ impl ArrowReaderBuilder {
             field_ids: self.field_ids,
             schema: self.schema,
             file_io: self.file_io,
-            predicates: self.predicates,
+            predicate: self.predicate,
         }
     }
 }
@@ -105,7 +105,7 @@ pub struct ArrowReader {
     #[allow(dead_code)]
     schema: SchemaRef,
     file_io: FileIO,
-    predicates: Option<BoundPredicate>,
+    predicate: Option<BoundPredicate>,
 }
 
 impl ArrowReader {
@@ -118,7 +118,7 @@ impl ArrowReader {
         let mut collector = CollectFieldIdVisitor {
             field_ids: HashSet::default(),
         };
-        if let Some(predicates) = &self.predicates {
+        if let Some(predicates) = &self.predicate {
             visit(&mut collector, predicates)?;
         }
 
@@ -162,7 +162,7 @@ impl ArrowReader {
         &self,
         parquet_schema: &SchemaDescriptor,
         arrow_schema: &ArrowSchemaRef,
-    ) -> crate::Result<ProjectionMask> {
+    ) -> Result<ProjectionMask> {
         if self.field_ids.is_empty() {
             Ok(ProjectionMask::all())
         } else {
@@ -232,7 +232,7 @@ impl ArrowReader {
         parquet_schema: &SchemaDescriptor,
         collector: &CollectFieldIdVisitor,
     ) -> Result<Option<RowFilter>> {
-        if let Some(predicates) = &self.predicates {
+        if let Some(predicate) = &self.predicate {
             let field_id_map = build_field_id_map(parquet_schema)?;
 
             // Collect Parquet column indices from field ids.
@@ -255,7 +255,7 @@ impl ArrowReader {
             // After collecting required leaf column indices used in the predicate,
             // creates the projection mask for the Arrow predicates.
             let projection_mask = ProjectionMask::leaves(parquet_schema, column_indices.clone());
-            let predicate_func = visit(&mut converter, predicates)?;
+            let predicate_func = visit(&mut converter, predicate)?;
             let arrow_predicate = ArrowPredicateFn::new(projection_mask, predicate_func);
             Ok(Some(RowFilter::new(vec![Box::new(arrow_predicate)])))
         } else {

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,19 +24,24 @@ use crate::expr::visitors::manifest_evaluator::ManifestEvaluator;
 use crate::expr::{Bind, BoundPredicate, Predicate};
 use crate::io::FileIO;
 use crate::spec::{
-    DataContentType, ManifestContentType, 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 serde::{Deserialize, Serialize};
 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 SUPPORTED_MANIFEST_FILE_CONTENT_TYPES: [ManifestContentType; 1] = [ManifestContentType::Data];
 
 /// A stream of [`FileScanTask`].
 pub type FileScanTaskStream = BoxStream<'static, Result<FileScanTask>>;
@@ -207,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 partition_spec_id = entry.partition_spec_id;
-
-                let partition_filter = partition_filter_cache.get(
-                    partition_spec_id,
-                    &context,
-                )?;
+        let manifest_list = context
+            .snapshot
+            .load_manifest_list(&context.file_io, &context.table_metadata)
+            .await?;
 
-                if let Some(partition_filter) = partition_filter {
-                    let manifest_evaluator = manifest_evaluator_cache.get(
-                        partition_spec_id,
-                        partition_filter,
-                    );
+        let (sender, receiver) = channel(CHANNEL_BUFFER_SIZE);
 
-                    if !manifest_evaluator.eval(entry)? {
-                        continue;
-                    }
-                }
+        spawn(async move {
+            let _ = ConcurrentFileScanStreamContext::new(context, sender)
+                .run(manifest_list)
+                .await;
+        });
 
-                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_file_path: manifest_entry.data_file().file_path().to_string(),
-                                start: 0,
-                                length: manifest_entry.file_size_in_bytes(),
-                            });
-                            yield scan_task?;
-                        }
-                    }
-                }
-            }
-        }
-        .boxed())
+        return Ok(receiver.boxed());
     }
 
     /// Returns an [`ArrowRecordBatchStream`].
@@ -334,13 +283,167 @@ 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, mut sender, context) = manifest_and_file_io_and_sender;
+
+        let manifest = manifest_file.load_manifest(&file_io).await?;
+        for manifest_entry in manifest.entries() {
+            if !manifest_entry.is_alive() {
+                continue;
+            }
+
+            Self::reject_unsupported_manifest_entry_content_types(manifest_entry)?;
+
+            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
+
+            sender
+                .send(Ok(Self::manifest_entry_to_file_scan_task(manifest_entry)))
+                .await?;
+        }
+
+        Ok(())
+    }
+
+    fn reject_unsupported_manifest_entry_content_types(
+        manifest_entry: &ManifestEntryRef,
+    ) -> Result<()> {
+        if matches!(manifest_entry.content_type(), DataContentType::Data) {
+            Ok(())
+        } else {
+            Err(Error::new(
+                ErrorKind::FeatureUnsupported,
+                format!(
+                    "Files of type '{:?}' are not supported yet.",
+                    manifest_entry.content_type()
+                ),
+            ))
+        }
+    }
 
-    /// 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_file_path: manifest_entry.file_path().to_string(),
+            start: 0,
+            length: manifest_entry.file_size_in_bytes(),
         }
-        false
     }
 }