Make scan::execute return a lazy iterator

acceptance/Cargo.toml

@@ -20,6 +20,7 @@ delta_kernel = { path = "../kernel", features = [
   "developer-visibility",
 ] }
 futures = "0.3"
+itertools = "0.13"
 object_store = { workspace = true }
 parquet = { workspace = true }
 serde = { version = "1", features = ["derive"] }

acceptance/src/data.rs

@@ -7,6 +7,7 @@ use arrow_select::{concat::concat_batches, filter::filter_record_batch, take::ta
 
 use delta_kernel::{engine::arrow_data::ArrowEngineData, DeltaResult, Engine, Error, Table};
 use futures::{stream::TryStreamExt, StreamExt};
+use itertools::Itertools;
 use object_store::{local::LocalFileSystem, ObjectStore};
 use parquet::arrow::async_reader::{ParquetObjectReader, ParquetRecordBatchStreamBuilder};
 
@@ -127,9 +128,9 @@ pub async fn assert_scan_data(engine: Arc<dyn Engine>, test_case: &TestCaseInfo)
     let mut schema = None;
     let batches: Vec<RecordBatch> = scan
         .execute(engine)?
-        .into_iter()
-        .map(|res| {
-            let data = res.raw_data.unwrap();
+        .map(|scan_result| -> DeltaResult<_> {
+            let scan_result = scan_result?;
+            let data = scan_result.raw_data?;
             let record_batch: RecordBatch = data
                 .into_any()
                 .downcast::<ArrowEngineData>()
@@ -138,13 +139,13 @@ pub async fn assert_scan_data(engine: Arc<dyn Engine>, test_case: &TestCaseInfo)
             if schema.is_none() {
                 schema = Some(record_batch.schema());
             }
-            if let Some(mask) = res.mask {
-                filter_record_batch(&record_batch, &mask.into()).unwrap()
+            if let Some(mask) = scan_result.mask {
+                Ok(filter_record_batch(&record_batch, &mask.into())?)
             } else {
-                record_batch
+                Ok(record_batch)
             }
         })
-        .collect();
+        .try_collect()?;
     let all_data = concat_batches(&schema.unwrap(), batches.iter()).map_err(Error::from)?;
     let all_data = sort_record_batch(all_data)?;
     let golden = read_golden(test_case.root_dir(), None).await?;

kernel/examples/read-table-single-threaded/src/main.rs

@@ -13,6 +13,7 @@ use delta_kernel::schema::Schema;
 use delta_kernel::{DeltaResult, Engine, Table};
 
 use clap::{Parser, ValueEnum};
+use itertools::Itertools;
 
 /// An example program that dumps out the data of a delta table. Struct and Map types are not
 /// supported.
@@ -92,7 +93,6 @@ fn try_main() -> DeltaResult<()> {
     let read_schema_opt = cli
         .columns
         .map(|cols| {
-            use itertools::Itertools;
             let table_schema = snapshot.schema();
             let selected_fields = cols
                 .iter()
@@ -113,26 +113,28 @@ fn try_main() -> DeltaResult<()> {
         .with_schema_opt(read_schema_opt)
         .build()?;
 
-    let mut batches = vec![];
-    for res in scan.execute(engine.as_ref())?.into_iter() {
-        let data = res.raw_data?;
-        let record_batch: RecordBatch = data
-            .into_any()
-            .downcast::<ArrowEngineData>()
-            .map_err(|_| delta_kernel::Error::EngineDataType("ArrowEngineData".to_string()))?
-            .into();
-        let batch = if let Some(mut mask) = res.mask {
-            let extra_rows = record_batch.num_rows() - mask.len();
-            if extra_rows > 0 {
-                // we need to extend the mask here in case it's too short
-                mask.extend(std::iter::repeat(true).take(extra_rows));
+    let batches: Vec<RecordBatch> = scan
+        .execute(engine.as_ref())?
+        .map(|scan_result| -> DeltaResult<_> {
+            let scan_result = scan_result?;
+            let data = scan_result.raw_data?;
+            let record_batch: RecordBatch = data
+                .into_any()
+                .downcast::<ArrowEngineData>()
+                .map_err(|_| delta_kernel::Error::EngineDataType("ArrowEngineData".to_string()))?
+                .into();
+            if let Some(mut mask) = scan_result.mask {
+                let extra_rows = record_batch.num_rows() - mask.len();
+                if extra_rows > 0 {
+                    // we need to extend the mask here in case it's too short
+                    mask.extend(std::iter::repeat(true).take(extra_rows));
+                }
+                Ok(filter_record_batch(&record_batch, &mask.into())?)
+            } else {
+                Ok(record_batch)
             }
-            filter_record_batch(&record_batch, &mask.into())?
-        } else {
-            record_batch
-        };
-        batches.push(batch);
-    }
+        })
+        .try_collect()?;
     print_batches(&batches)?;
     Ok(())
 }

kernel/src/scan/mod.rs

@@ -227,14 +227,17 @@ impl Scan {
     }
 
     /// Perform an "all in one" scan. This will use the provided `engine` to read and
-    /// process all the data for the query. Each [`ScanResult`] in the resultant vector encapsulates
+    /// process all the data for the query. Each [`ScanResult`] in the resultant iterator encapsulates
     /// the raw data and an optional boolean vector built from the deletion vector if it was
     /// present. See the documentation for [`ScanResult`] for more details. Generally
     /// connectors/engines will want to use [`Scan::scan_data`] so they can have more control over
     /// the execution of the scan.
-    // This calls [`Scan::files`] to get a set of `Add` actions for the scan, and then uses the
+    // This calls [`Scan::scan_data`] to get an iterator of `ScanData` actions for the scan, and then uses the
     // `engine`'s [`crate::ParquetHandler`] to read the actual table data.
-    pub fn execute(&self, engine: &dyn Engine) -> DeltaResult<Vec<ScanResult>> {
+    pub fn execute<'a>(
+        &'a self,
+        engine: &'a dyn Engine,
+    ) -> DeltaResult<impl Iterator<Item = DeltaResult<ScanResult>> + 'a> {
         struct ScanFile {
             path: String,
             size: i64,
@@ -264,15 +267,18 @@ impl Scan {
 
         let global_state = Arc::new(self.global_scan_state());
         let scan_data = self.scan_data(engine)?;
-        let mut scan_files = vec![];
-        for data in scan_data {
-            let (data, vec) = data?;
-            scan_files =
-                state::visit_scan_files(data.as_ref(), &vec, scan_files, scan_data_callback)?;
-        }
-        scan_files
-            .into_iter()
-            .map(|scan_file| -> DeltaResult<_> {
+        let scan_files_iter = scan_data
+            .map(|res| {
+                let (data, vec) = res?;
+                let scan_files = vec![];
+                state::visit_scan_files(data.as_ref(), &vec, scan_files, scan_data_callback)
+            })
+            // Iterator<DeltaResult<Vec<ScanFile>>> to Iterator<DeltaResult<ScanFile>>
+            .flatten_ok();
+
+        let result = scan_files_iter
+            .map(move |scan_file| -> DeltaResult<_> {
+                let scan_file = scan_file?;
                 let file_path = self.snapshot.table_root.join(&scan_file.path)?;
                 let mut selection_vector = scan_file
                     .dv_info
@@ -288,7 +294,7 @@ impl Scan {
                     None,
                 )?;
                 let gs = global_state.clone(); // Arc clone
-                Ok(read_result_iter.into_iter().map(move |read_result| {
+                Ok(read_result_iter.map(move |read_result| -> DeltaResult<_> {
                     let read_result = read_result?;
                     // to transform the physical data into the correct logical form
                     let logical = transform_to_logical_internal(
@@ -314,8 +320,11 @@ impl Scan {
                     Ok(result)
                 }))
             })
+            // Iterator<DeltaResult<Iterator<DeltaResult<ScanResult>>>> to Iterator<DeltaResult<DeltaResult<ScanResult>>>
             .flatten_ok()
-            .try_collect()?
+            // Iterator<DeltaResult<DeltaResult<ScanResult>>> to Iterator<DeltaResult<ScanResult>>
+            .map(|x| x?);
+        Ok(result)
     }
 }
 
@@ -633,7 +642,7 @@ mod tests {
         let table = Table::new(url);
         let snapshot = table.snapshot(&engine, None).unwrap();
         let scan = snapshot.into_scan_builder().build().unwrap();
-        let files = scan.execute(&engine).unwrap();
+        let files: Vec<ScanResult> = scan.execute(&engine).unwrap().try_collect().unwrap();
 
         assert_eq!(files.len(), 1);
         let num_rows = files[0].raw_data.as_ref().unwrap().length();

kernel/tests/dv.rs

@@ -1,12 +1,32 @@
 //! Read a small table with/without deletion vectors.
 //! Must run at the root of the crate
+use std::ops::Add;
 use std::path::PathBuf;
 
 use delta_kernel::engine::sync::SyncEngine;
-use delta_kernel::Table;
+use delta_kernel::scan::ScanResult;
+use delta_kernel::{DeltaResult, Table};
 
+use itertools::Itertools;
 use test_log::test;
 
+fn count_total_scan_rows(
+    scan_result_iter: impl Iterator<Item = DeltaResult<ScanResult>>,
+) -> DeltaResult<usize> {
+    scan_result_iter
+        .map(|scan_result| {
+            let scan_result = scan_result?;
+            let data = scan_result.raw_data?;
+            // NOTE: The mask only suppresses rows for which it is both present and false.
+            let deleted_rows = scan_result
+                .mask
+                .as_ref()
+                .map_or(0, |mask| mask.iter().filter(|&&m| !m).count());
+            Ok(data.length() - deleted_rows)
+        })
+        .fold_ok(0, Add::add)
+}
+
 #[test]
 fn dv_table() -> Result<(), Box<dyn std::error::Error>> {
     let path = std::fs::canonicalize(PathBuf::from("./tests/data/table-with-dv-small/"))?;
@@ -18,16 +38,7 @@ fn dv_table() -> Result<(), Box<dyn std::error::Error>> {
     let scan = snapshot.into_scan_builder().build()?;
 
     let stream = scan.execute(&engine)?;
-    let mut total_rows = 0;
-    for res in stream {
-        let data = res.raw_data?;
-        let rows = data.length();
-        for i in 0..rows {
-            if res.mask.as_ref().map_or(true, |mask| mask[i]) {
-                total_rows += 1;
-            }
-        }
-    }
+    let total_rows = count_total_scan_rows(stream)?;
     assert_eq!(total_rows, 8);
     Ok(())
 }
@@ -43,16 +54,7 @@ fn non_dv_table() -> Result<(), Box<dyn std::error::Error>> {
     let scan = snapshot.into_scan_builder().build()?;
 
     let stream = scan.execute(&engine)?;
-    let mut total_rows = 0;
-    for res in stream {
-        let data = res.raw_data?;
-        let rows = data.length();
-        for i in 0..rows {
-            if res.mask.as_ref().map_or(true, |mask| mask[i]) {
-                total_rows += 1;
-            }
-        }
-    }
+    let total_rows = count_total_scan_rows(stream)?;
     assert_eq!(total_rows, 10);
     Ok(())
 }

kernel/tests/golden_tables.rs

@@ -7,6 +7,7 @@ use arrow::{compute::filter_record_batch, record_batch::RecordBatch};
 use arrow_ord::sort::{lexsort_to_indices, SortColumn};
 use arrow_schema::Schema;
 use arrow_select::{concat::concat_batches, take::take};
+use itertools::Itertools;
 use paste::paste;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
@@ -153,37 +154,35 @@ async fn latest_snapshot_test(
     table: Table,
     expected_path: Option<PathBuf>,
 ) -> Result<(), Box<dyn std::error::Error>> {
-    let snapshot = table.snapshot(&engine, None).unwrap();
+    let snapshot = table.snapshot(&engine, None)?;
 
-    let scan = snapshot.into_scan_builder().build().unwrap();
-    let scan_res = scan.execute(&engine).unwrap();
+    let scan = snapshot.into_scan_builder().build()?;
+    let scan_res = scan.execute(&engine)?;
     let batches: Vec<RecordBatch> = scan_res
-        .into_iter()
-        .map(|sr| {
-            let data = sr.raw_data.unwrap();
-            let record_batch = to_arrow(data).unwrap();
-            if let Some(mut mask) = sr.mask {
+        .map(|scan_result| -> DeltaResult<_> {
+            let scan_result = scan_result?;
+            let data = scan_result.raw_data?;
+            let record_batch = to_arrow(data)?;
+            if let Some(mut mask) = scan_result.mask {
                 let extra_rows = record_batch.num_rows() - mask.len();
                 if extra_rows > 0 {
                     // we need to extend the mask here in case it's too short
                     mask.extend(std::iter::repeat(true).take(extra_rows));
                 }
-                filter_record_batch(&record_batch, &mask.into()).unwrap()
+                Ok(filter_record_batch(&record_batch, &mask.into())?)
             } else {
-                record_batch
+                Ok(record_batch)
             }
         })
-        .collect();
+        .try_collect()?;
 
-    let expected = read_expected(&expected_path.expect("expect an expected dir"))
-        .await
-        .unwrap();
+    let expected = read_expected(&expected_path.expect("expect an expected dir")).await?;
 
-    let schema: Arc<Schema> = Arc::new(scan.schema().as_ref().try_into().unwrap());
+    let schema: Arc<Schema> = Arc::new(scan.schema().as_ref().try_into()?);
 
-    let result = concat_batches(&schema, &batches).unwrap();
-    let result = sort_record_batch(result).unwrap();
-    let expected = sort_record_batch(expected).unwrap();
+    let result = concat_batches(&schema, &batches)?;
+    let result = sort_record_batch(result)?;
+    let expected = sort_record_batch(expected)?;
     assert_columns_match(result.columns(), expected.columns());
     assert_schema_fields_match(expected.schema().as_ref(), result.schema().as_ref());
     assert!(