feat: ArrowReader processes FileScanTasks concurrently

Author: sdd

crates/iceberg/src/arrow/reader.rs

@@ -17,17 +17,26 @@
 
 //! Parquet file data reader
 
-use crate::error::Result;
+use crate::arrow::{arrow_schema_to_schema, get_arrow_datum};
+use crate::expr::visitors::bound_predicate_visitor::{visit, BoundPredicateVisitor};
+use crate::expr::{BoundPredicate, BoundReference};
+use crate::io::{FileIO, FileMetadata, FileRead};
+use crate::scan::FileScanTask;
+use crate::scan::{ArrowRecordBatchStream, FileScanTaskStream};
+use crate::spec::{Datum, SchemaRef};
+use crate::Result;
+use crate::{Error, ErrorKind};
 use arrow_arith::boolean::{and, is_not_null, is_null, not, or};
 use arrow_array::{ArrayRef, BooleanArray, RecordBatch};
 use arrow_ord::cmp::{eq, gt, gt_eq, lt, lt_eq, neq};
 use arrow_schema::{ArrowError, DataType, SchemaRef as ArrowSchemaRef};
-use async_stream::try_stream;
 use bytes::Bytes;
 use fnv::FnvHashSet;
+use futures::channel::mpsc::{channel, Sender};
 use futures::future::BoxFuture;
 use futures::stream::StreamExt;
 use futures::{try_join, TryFutureExt};
+use futures::{SinkExt, TryStreamExt};
 use parquet::arrow::arrow_reader::{ArrowPredicateFn, RowFilter};
 use parquet::arrow::async_reader::{AsyncFileReader, MetadataLoader};
 use parquet::arrow::{ParquetRecordBatchStreamBuilder, ProjectionMask, PARQUET_FIELD_ID_META_KEY};
@@ -37,22 +46,18 @@ use std::collections::{HashMap, HashSet};
 use std::ops::Range;
 use std::str::FromStr;
 use std::sync::Arc;
+use tokio::spawn;
 
-use crate::arrow::{arrow_schema_to_schema, get_arrow_datum};
-use crate::expr::visitors::bound_predicate_visitor::{visit, BoundPredicateVisitor};
-use crate::expr::{BoundPredicate, BoundReference};
-use crate::io::{FileIO, FileMetadata, FileRead};
-use crate::scan::{ArrowRecordBatchStream, FileScanTaskStream};
-use crate::spec::{Datum, SchemaRef};
-use crate::{Error, ErrorKind};
+const CHANNEL_BUFFER_SIZE: usize = 10;
+const CONCURRENCY_LIMIT_TASKS: usize = 10;
 
 /// Builder to create ArrowReader
 pub struct ArrowReaderBuilder {
     batch_size: Option<usize>,
     field_ids: Vec<usize>,
     file_io: FileIO,
     schema: SchemaRef,
-    predicates: Option<BoundPredicate>,
+    predicate: Option<BoundPredicate>,
 }
 
 impl ArrowReaderBuilder {
@@ -63,7 +68,7 @@ impl ArrowReaderBuilder {
             field_ids: vec![],
             file_io,
             schema,
-            predicates: None,
+            predicate: None,
         }
     }
 
@@ -82,7 +87,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 +98,7 @@ impl ArrowReaderBuilder {
             field_ids: self.field_ids,
             schema: self.schema,
             file_io: self.file_io,
-            predicates: self.predicates,
+            predicate: self.predicate,
         }
     }
 }
@@ -105,64 +110,140 @@ pub struct ArrowReader {
     #[allow(dead_code)]
     schema: SchemaRef,
     file_io: FileIO,
-    predicates: Option<BoundPredicate>,
+    predicate: Option<BoundPredicate>,
 }
 
 impl ArrowReader {
     /// Take a stream of FileScanTasks and reads all the files.
     /// Returns a stream of Arrow RecordBatches containing the data from the files
-    pub fn read(self, mut tasks: FileScanTaskStream) -> crate::Result<ArrowRecordBatchStream> {
-        let file_io = self.file_io.clone();
+    pub fn read(self, tasks: FileScanTaskStream) -> Result<ArrowRecordBatchStream> {
+        let (sender, receiver) = channel(CHANNEL_BUFFER_SIZE);
 
         // Collect Parquet column indices from field ids
         let mut collector = CollectFieldIdVisitor {
             field_ids: HashSet::default(),
         };
-        if let Some(predicates) = &self.predicates {
+        if let Some(predicates) = &self.predicate {
             visit(&mut collector, predicates)?;
         }
 
-        Ok(try_stream! {
-            while let Some(Ok(task)) = tasks.next().await {
-                let parquet_file = file_io
-                    .new_input(task.data().data_file().file_path())?;
-                let (parquet_metadata, parquet_reader) = try_join!(parquet_file.metadata(), parquet_file.reader())?;
-                let arrow_file_reader = ArrowFileReader::new(parquet_metadata, parquet_reader);
+        let tasks = tasks.map(move |task| self.build_file_scan_task_context(task, sender.clone()));
 
-                let mut batch_stream_builder = ParquetRecordBatchStreamBuilder::new(arrow_file_reader)
-                    .await?;
+        spawn(async move {
+            tasks
+                .try_for_each_concurrent(CONCURRENCY_LIMIT_TASKS, Self::process_file_scan_task)
+                .await
+        });
 
-                let parquet_schema = batch_stream_builder.parquet_schema();
-                let arrow_schema = batch_stream_builder.schema();
-                let projection_mask = self.get_arrow_projection_mask(parquet_schema, arrow_schema)?;
-                batch_stream_builder = batch_stream_builder.with_projection(projection_mask);
+        Ok(receiver.boxed())
+    }
 
-                let parquet_schema = batch_stream_builder.parquet_schema();
-                let row_filter = self.get_row_filter(parquet_schema, &collector)?;
+    fn build_file_scan_task_context(
+        &self,
+        task: Result<FileScanTask>,
+        sender: Sender<Result<RecordBatch>>,
+    ) -> Result<FileScanTaskContext> {
+        Ok(FileScanTaskContext::new(
+            task?,
+            self.file_io.clone(),
+            sender,
+            self.batch_size,
+            self.field_ids.clone(),
+            self.schema.clone(),
+            self.predicate.clone(),
+        ))
+    }
+
+    async fn process_file_scan_task(mut context: FileScanTaskContext) -> Result<()> {
+        let file_scan_task = context.take_task();
 
-                if let Some(row_filter) = row_filter {
-                    batch_stream_builder = batch_stream_builder.with_row_filter(row_filter);
-                }
+        // Collect Parquet column indices from field ids
+        let mut collector = CollectFieldIdVisitor {
+            field_ids: HashSet::default(),
+        };
+        if let Some(predicate) = &context.predicate {
+            visit(&mut collector, predicate)?;
+        }
 
-                if let Some(batch_size) = self.batch_size {
-                    batch_stream_builder = batch_stream_builder.with_batch_size(batch_size);
-                }
+        let parquet_file = context
+            .file_io
+            .new_input(file_scan_task.data().data_file().file_path())?;
+        let (parquet_metadata, parquet_reader) =
+            try_join!(parquet_file.metadata(), parquet_file.reader())?;
+        let arrow_file_reader = ArrowFileReader::new(parquet_metadata, parquet_reader);
 
-                let mut batch_stream = batch_stream_builder.build()?;
+        let mut batch_stream_builder =
+            ParquetRecordBatchStreamBuilder::new(arrow_file_reader).await?;
 
-                while let Some(batch) = batch_stream.next().await {
-                    yield batch?;
-                }
-            }
+        let parquet_schema = batch_stream_builder.parquet_schema();
+        let arrow_schema = batch_stream_builder.schema();
+
+        let projection_mask = context.get_arrow_projection_mask(parquet_schema, arrow_schema)?;
+        batch_stream_builder = batch_stream_builder.with_projection(projection_mask);
+
+        let parquet_schema = batch_stream_builder.parquet_schema();
+        let row_filter = context.get_row_filter(parquet_schema, &collector)?;
+
+        if let Some(row_filter) = row_filter {
+            batch_stream_builder = batch_stream_builder.with_row_filter(row_filter);
+        }
+
+        if let Some(batch_size) = context.batch_size {
+            batch_stream_builder = batch_stream_builder.with_batch_size(batch_size);
+        }
+
+        let mut batch_stream = batch_stream_builder.build()?;
+
+        while let Some(batch) = batch_stream.next().await {
+            context.sender.send(Ok(batch?)).await?;
+        }
+
+        Ok(())
+    }
+}
+
+struct FileScanTaskContext {
+    file_scan_task: Option<FileScanTask>,
+    file_io: FileIO,
+    sender: Sender<Result<RecordBatch>>,
+    batch_size: Option<usize>,
+    field_ids: Vec<usize>,
+    schema: SchemaRef,
+    predicate: Option<BoundPredicate>,
+}
+
+impl FileScanTaskContext {
+    fn new(
+        file_scan_task: FileScanTask,
+        file_io: FileIO,
+        sender: Sender<Result<RecordBatch>>,
+        batch_size: Option<usize>,
+        field_ids: Vec<usize>,
+        schema: SchemaRef,
+        predicate: Option<BoundPredicate>,
+    ) -> Self {
+        FileScanTaskContext {
+            file_scan_task: Some(file_scan_task),
+            file_io,
+            sender,
+            batch_size,
+            field_ids,
+            schema,
+            predicate,
         }
-        .boxed())
+    }
+
+    fn take_task(&mut self) -> FileScanTask {
+        let mut result = None;
+        std::mem::swap(&mut self.file_scan_task, &mut result);
+        result.unwrap()
     }
 
     fn get_arrow_projection_mask(
         &self,
         parquet_schema: &SchemaDescriptor,
         arrow_schema: &ArrowSchemaRef,
-    ) -> crate::Result<ProjectionMask> {
+    ) -> Result<ProjectionMask> {
         if self.field_ids.is_empty() {
             Ok(ProjectionMask::all())
         } else {
@@ -232,7 +313,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 +336,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 {