sketch out testing framework

Author: alamb

Cargo.lock

@@ -73,6 +73,7 @@ zstd = { workspace = true, optional = true }
 [dev-dependencies]
 criterion = { workspace = true }
 tempfile = { workspace = true }
+insta = { workspace = true }
 
 # Note: add additional linter rules in lib.rs.
 # Rust does not support workspace + new linter rules in subcrates yet

datafusion/datasource/Cargo.toml

@@ -611,22 +611,26 @@ impl FileStreamMetrics {
 
 #[cfg(test)]
 mod tests {
-    use crate::PartitionedFile;
-    use crate::file_scan_config::FileScanConfigBuilder;
+    use crate::file_scan_config::{FileScanConfig, FileScanConfigBuilder};
+    use crate::morsel::test_utils::{
+        MockMorselSpec, MockMorselizer, MockPlannerSteps, MorselEvent, MorselObserver,
+        PlannerStep,
+    };
     use crate::tests::make_partition;
+    use crate::{PartitionedFile, TableSchema};
+    use arrow::datatypes::Int32Type;
     use datafusion_common::error::Result;
     use datafusion_execution::object_store::ObjectStoreUrl;
     use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
-    use futures::{FutureExt as _, StreamExt as _};
     use std::sync::Arc;
     use std::sync::atomic::{AtomicUsize, Ordering};
 
     use crate::file_stream::{FileOpenFuture, FileOpener, FileStream, OnError};
     use crate::test_util::MockSource;
-    use arrow::array::RecordBatch;
-    use arrow::datatypes::Schema;
-
+    use arrow::array::{Array, AsArray, RecordBatch};
+    use arrow::datatypes::{DataType, Field, Schema};
     use datafusion_common::{assert_batches_eq, exec_err, internal_err};
+    use futures::{FutureExt, StreamExt};
 
     /// Test `FileOpener` which will simulate errors during file opening or scanning
     #[derive(Default)]
@@ -741,7 +745,7 @@ mod tests {
 
             let on_error = self.on_error;
 
-            let table_schema = crate::table_schema::TableSchema::new(file_schema, vec![]);
+            let table_schema = TableSchema::new(file_schema, vec![]);
             let config = FileScanConfigBuilder::new(
                 ObjectStoreUrl::parse("test:///").unwrap(),
                 Arc::new(MockSource::new(table_schema)),
@@ -774,6 +778,113 @@ mod tests {
             .expect("error executing stream")
     }
 
+    #[tokio::test]
+    async fn morsel_framework_single_planner_single_morsel_no_io() -> Result<()> {
+        let observer = MorselObserver::new();
+        let morselizer = MockMorselizer::new()
+            .with_file(
+                "file1.parquet",
+                vec![MockPlannerSteps::new(
+                    0,
+                    vec![
+                        PlannerStep::ReturnPlan {
+                            morsels: vec![MockMorselSpec::single_batch(10, 42)],
+                            planners: vec![],
+                            io_future_id: None,
+                            io_polls: 0,
+                        },
+                        PlannerStep::ReturnNone,
+                    ],
+                )],
+            )
+            .with_observer(observer.clone());
+
+        let config = test_config(vec!["file1.parquet"]);
+        insta::assert_snapshot!(run_stream(morselizer, config).await.unwrap(), @r"
+        **** Batch: 42
+        **** Done
+        ");
+
+        assert_eq!(
+            observer.events(),
+            vec![
+                MorselEvent::MorselizeFile {
+                    path: "file1.parquet".to_string()
+                },
+                MorselEvent::PlannerCreated { planner_id: 0 },
+                MorselEvent::PlannerPlanCalled { planner_id: 0 },
+                MorselEvent::MorselProduced {
+                    planner_id: 0,
+                    morsel_id: 10
+                },
+                MorselEvent::PlannerPlanCalled { planner_id: 0 },
+                MorselEvent::MorselStreamStarted { morsel_id: 10 },
+                MorselEvent::MorselStreamBatchProduced {
+                    morsel_id: 10,
+                    batch_id: 42
+                },
+                MorselEvent::MorselStreamFinished { morsel_id: 10 },
+            ]
+        );
+
+        Ok(())
+    }
+
+    fn test_config(file_names: Vec<&str>) -> FileScanConfig {
+        let file_group = file_names
+            .into_iter()
+            .map(|name| PartitionedFile::new(name, 10))
+            .collect();
+        let table_schema = TableSchema::new(
+            Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)])),
+            vec![],
+        );
+        FileScanConfigBuilder::new(
+            ObjectStoreUrl::parse("test:///").unwrap(),
+            Arc::new(MockSource::new(table_schema)),
+        )
+        .with_file_group(file_group)
+        .build()
+    }
+
+    /// Creates a  [`FileStream`] for reading the specified config, reads all
+    /// record batches, and returns a stringified version of the results for
+    /// easy comparison in tests.
+    ///
+    /// If the FileStream returns an error during execution, formats that error
+    /// in the output stream, rather than returning an error
+    async fn run_stream(
+        morselizer: MockMorselizer,
+        config: FileScanConfig,
+    ) -> Result<String> {
+        let metrics_set = ExecutionPlanMetricsSet::new();
+        let mut stream = FileStream::new_with_morselizer(
+            &config,
+            0,
+            Box::new(morselizer),
+            &metrics_set,
+        )?;
+
+        let mut stream_contents = Vec::new();
+        while let Some(result) = stream.next().await {
+            match result {
+                Ok(batch) => {
+                    // each batch should have a single int32 column with batch id
+                    let col = batch.column(0).as_primitive::<Int32Type>();
+                    assert_eq!(col.len(), 1);
+                    assert!(col.is_valid(0));
+                    let batch_id = col.value(0);
+                    stream_contents.push(format!("**** Batch: {batch_id}"));
+                }
+                Err(e) => {
+                    stream_contents.push(format!("**** Error: {e}"));
+                }
+            }
+        }
+        stream_contents.push("**** Done".to_string());
+        Ok(stream_contents.join("\n"))
+    }
+
     #[tokio::test]
     async fn on_error_opening() -> Result<()> {
         let batches = FileStreamTest::new()

datafusion/datasource/src/file_stream.rs

@@ -23,6 +23,8 @@
 //! Evaluation Framework for the Many-Core Age](https://db.in.tum.de/~leis/papers/morsels.pdf)
 
 mod adapters;
+#[cfg(test)]
+pub(crate) mod test_utils;
 
 use crate::PartitionedFile;
 use arrow::array::RecordBatch;