more tests

Author: alamb

datafusion/datasource/src/file_stream.rs

@@ -778,12 +778,46 @@ mod tests {
             .expect("error executing stream")
     }
 
+    /// Helper for morsel-driven `FileStream` tests that bundles the mock
+    /// `Morselizer` setup with the corresponding `FileScanConfig`.
+    struct MorselTest {
+        morselizer: MockMorselizer,
+        file_names: Vec<String>,
+    }
+
+    impl MorselTest {
+        fn new() -> Self {
+            Self {
+                morselizer: MockMorselizer::new(),
+                file_names: vec![],
+            }
+        }
+
+        fn with_file(mut self, path: impl Into<String>, planner: MockPlanner) -> Self {
+            let path = path.into();
+            self.morselizer = self.morselizer.with_file(path.clone(), planner);
+            self.file_names.push(path);
+            self
+        }
+
+        fn with_observer(mut self, observer: MorselObserver) -> Self {
+            self.morselizer = self.morselizer.with_observer(observer);
+            self
+        }
+
+        async fn run(self) -> Result<String> {
+            let file_names = self.file_names.iter().map(String::as_str).collect();
+            let config = test_config(file_names);
+            run_stream(self.morselizer, config).await
+        }
+    }
+
     /// Verifies the simplest morsel-driven flow: one planner produces one
     /// morsel immediately, and the morsel is then scanned to completion.
     #[tokio::test]
     async fn morsel_framework_single_morsel_no_io() -> Result<()> {
         let observer = MorselObserver::new();
-        let morselizer = MockMorselizer::new()
+        let test = MorselTest::new()
             .with_file(
                 "file1.parquet",
                 MockPlanner::builder()
@@ -794,8 +828,7 @@ mod tests {
             )
             .with_observer(observer.clone());
 
-        let config = test_config(vec!["file1.parquet"]);
-        insta::assert_snapshot!(run_stream(morselizer, config).await.unwrap(), @r"
+        insta::assert_snapshot!(test.run().await.unwrap(), @r"
         Batch: 42
         Done
         ");
@@ -820,9 +853,9 @@ mod tests {
     /// Verifies that a planner can block on one I/O phase, resume, and only
     /// then produce its morsel.
     #[tokio::test]
-    async fn morsel_framework_single_io() -> Result<()> {
+    async fn morsel_framework_single_morsel_io() -> Result<()> {
         let observer = MorselObserver::new();
-        let morselizer = MockMorselizer::new()
+        let test = MorselTest::new()
             .with_file(
                 "file1.parquet",
                 MockPlanner::builder()
@@ -834,8 +867,7 @@ mod tests {
             )
             .with_observer(observer.clone());
 
-        let config = test_config(vec!["file1.parquet"]);
-        insta::assert_snapshot!(run_stream(morselizer, config).await.unwrap(), @r"
+        insta::assert_snapshot!(test.run().await.unwrap(), @r"
         Batch: 42
         Done
         ");
@@ -870,7 +902,7 @@ mod tests {
     #[tokio::test]
     async fn morsel_framework_two_cpu_steps_before_morsel() -> Result<()> {
         let observer = MorselObserver::new();
-        let morselizer = MockMorselizer::new()
+        let test = MorselTest::new()
             .with_file(
                 "file1.parquet",
                 MockPlanner::builder()
@@ -883,8 +915,7 @@ mod tests {
             )
             .with_observer(observer.clone());
 
-        let config = test_config(vec!["file1.parquet"]);
-        insta::assert_snapshot!(run_stream(morselizer, config).await.unwrap(), @r"
+        insta::assert_snapshot!(test.run().await.unwrap(), @r"
         Batch: 42
         Done
         ");
@@ -911,7 +942,7 @@ mod tests {
     /// Verifies direct morsels returned from a planner are consumed before
     /// batches produced by any returned child planners.
     #[tokio::test]
-    async fn morsel_framework_direct_morsel_before_child_planner() -> Result<()> {
+    async fn morsel_framework_morsels_before_child_planner() -> Result<()> {
         let observer = MorselObserver::new();
         let child_planner = MockPlanner::builder()
             .with_id(PlannerId(1))
@@ -930,12 +961,11 @@ mod tests {
             .return_none()
             .build();
 
-        let morselizer = MockMorselizer::new()
+        let test = MorselTest::new()
             .with_file("file1.parquet", parent_planner)
             .with_observer(observer.clone());
 
-        let config = test_config(vec!["file1.parquet"]);
-        insta::assert_snapshot!(run_stream(morselizer, config).await.unwrap(), @r"
+        insta::assert_snapshot!(test.run().await.unwrap(), @r"
         Batch: 42
         Batch: 43
         Done
@@ -965,13 +995,90 @@ mod tests {
         Ok(())
     }
 
+    /// Verifies the non-ordered behavior for child planners: if the first child
+    /// planner blocks on I/O and the second can make progress immediately, the
+    /// second planner's batches are emitted first.
+    #[tokio::test]
+    async fn morsel_framework_child_planner_reorder() -> Result<()> {
+        let observer = MorselObserver::new();
+        let planner_1 = MockPlanner::builder()
+            .with_id(PlannerId(1))
+            .return_plan(ReturnPlanBuilder::new().with_io(IoFutureId(100), 2))
+            .return_morsel(MorselId(11), 41)
+            .return_none()
+            .build();
+        let planner_2 = MockPlanner::builder()
+            .with_id(PlannerId(2))
+            .return_plan(ReturnPlanBuilder::new().with_io(IoFutureId(101), 1)) // IO returns after 1 poll
+            .return_morsel(MorselId(12), 42)
+            .return_none()
+            .build();
+
+        let parent_planner = MockPlanner::builder()
+            .with_id(PlannerId(0))
+            .return_plan(
+                ReturnPlanBuilder::new()
+                    .with_planner(planner_1)
+                    .with_planner(planner_2),
+            )
+            .return_none()
+            .build();
+
+        let test = MorselTest::new()
+            .with_file("file1.parquet", parent_planner)
+            .with_observer(observer.clone());
+
+        insta::assert_snapshot!(test.run().await.unwrap(), @r"
+        Batch: 42
+        Batch: 41
+        Done
+        ");
+
+        assert_eq!(
+            observer.events(),
+            MorselEventsBuilder::new()
+                .morselize_file("file1.parquet")
+                .planner_created(PlannerId(0))
+                .planner_called(PlannerId(0))
+                .planner_produced_child(PlannerId(0), PlannerId(1))
+                .planner_produced_child(PlannerId(0), PlannerId(2))
+                .planner_called(PlannerId(0))
+                .planner_called(PlannerId(1))
+                .io_future_created(PlannerId(1), IoFutureId(100))
+                .planner_called(PlannerId(2))
+                .io_future_created(PlannerId(2), IoFutureId(101))
+                .io_future_polled(PlannerId(1), IoFutureId(100))
+                .io_future_polled(PlannerId(2), IoFutureId(101))
+                .io_future_polled(PlannerId(1), IoFutureId(100))
+                .io_future_polled(PlannerId(2), IoFutureId(101))
+                .io_future_resolved(PlannerId(2), IoFutureId(101))
+                .planner_called(PlannerId(2))
+                .morsel_produced(PlannerId(2), MorselId(12))
+                .planner_called(PlannerId(2))
+                .io_future_polled(PlannerId(1), IoFutureId(100))
+                .io_future_resolved(PlannerId(1), IoFutureId(100))
+                .planner_called(PlannerId(1))
+                .morsel_produced(PlannerId(1), MorselId(11))
+                .planner_called(PlannerId(1))
+                .morsel_stream_started(MorselId(12))
+                .morsel_stream_batch_produced(MorselId(12), BatchId(42))
+                .morsel_stream_finished(MorselId(12))
+                .morsel_stream_started(MorselId(11))
+                .morsel_stream_batch_produced(MorselId(11), BatchId(41))
+                .morsel_stream_finished(MorselId(11))
+                .build()
+        );
+
+        Ok(())
+    }
+
     /// Verifies that `FileStream` overlaps planner I/O across multiple files
     /// rather than waiting for the first file to finish before starting the
     /// second.
     #[tokio::test]
     async fn morsel_framework_two_files_overlapping_io() -> Result<()> {
         let observer = MorselObserver::new();
-        let morselizer = MockMorselizer::new()
+        let test = MorselTest::new()
             .with_file(
                 "file1.parquet",
                 MockPlanner::builder()
@@ -992,8 +1099,7 @@ mod tests {
             )
             .with_observer(observer.clone());
 
-        let config = test_config(vec!["file1.parquet", "file2.parquet"]);
-        insta::assert_snapshot!(run_stream(morselizer, config).await.unwrap(), @r"
+        insta::assert_snapshot!(test.run().await.unwrap(), @r"
         Batch: 42
         Batch: 43
         Done