Cleanup

Author: alamb

datafusion/datasource/src/file_stream.rs

@@ -663,8 +663,8 @@ impl FileStreamMetrics {
 mod tests {
     use crate::file_scan_config::{FileScanConfig, FileScanConfigBuilder};
     use crate::morsel::test_utils::{
-        IoFutureId, MockMorselSpec, MockMorselizer, MockPlanner, MorselId,
-        MorselObserver, PlannerId, ReturnPlanBuilder,
+        IoFutureId, MockMorselSpec, MockMorselizer, MockPlanner, MorselId, PlannerId,
+        ReturnPlanBuilder,
     };
     use crate::tests::make_partition;
     use crate::{PartitionedFile, TableSchema};
@@ -834,7 +834,6 @@ mod tests {
     struct MorselTest {
         morselizer: MockMorselizer,
         file_names: Vec<String>,
-        observer: Option<MorselObserver>,
         preserve_order: bool,
     }
 
@@ -843,7 +842,6 @@ mod tests {
             Self {
                 morselizer: MockMorselizer::new(),
                 file_names: vec![],
-                observer: None,
                 preserve_order: false,
             }
         }
@@ -855,24 +853,18 @@ mod tests {
             self
         }
 
-        fn reset_observer(mut self) -> Self {
-            let observer = MorselObserver::new();
-            self.morselizer = self.morselizer.with_observer(observer);
-            self
-        }
-
-        fn with_observer(mut self, observer: MorselObserver) -> Self {
-            self.morselizer = self.morselizer.with_observer(observer.clone());
-            self.observer = Some(observer);
-            self
-        }
-
         fn with_preserve_order(mut self, preserve_order: bool) -> Self {
             self.preserve_order = preserve_order;
             self
         }
 
         async fn run(self) -> Result<String> {
+            // handle to shared observer
+            let observer = self.morselizer.observer().clone();
+            // Clear any prior observer events before running the test, so the
+            // snapshot only includes events from this run.
+            observer.clear();
+
             let file_names = self.file_names.iter().map(String::as_str).collect();
             let config = test_config(file_names);
             let metrics_set = ExecutionPlanMetricsSet::new();
@@ -908,10 +900,8 @@ mod tests {
             // together. This makes scheduler changes much easier to review than
             // maintaining long hand-written event assertions separately.
             let mut parts = vec!["----- Output Stream -----".to_string(), output];
-            if let Some(observer) = self.observer {
                 parts.push("----- File Stream Events -----".to_string());
                 parts.push(observer.format_events());
-            }
             Ok(parts.join("\n"))
         }
     }
@@ -920,7 +910,6 @@ mod tests {
     /// 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 test = MorselTest::new()
             .with_file(
                 "file1.parquet",
@@ -929,8 +918,7 @@ mod tests {
                     .return_morsel(MorselId(10), 42)
                     .return_none()
                     .build(),
-            )
-            .with_observer(observer.clone());
+            );
 
         insta::assert_snapshot!(test.run().await.unwrap(), @r"
         ----- Output Stream -----
@@ -954,7 +942,6 @@ mod tests {
     /// then produce its morsel.
     #[tokio::test]
     async fn morsel_framework_single_morsel_io() -> Result<()> {
-        let observer = MorselObserver::new();
         let test = MorselTest::new()
             .with_file(
                 "file1.parquet",
@@ -964,8 +951,7 @@ mod tests {
                     .return_morsel(MorselId(10), 42)
                     .return_none()
                     .build(),
-            )
-            .with_observer(observer.clone());
+            );
 
         insta::assert_snapshot!(test.run().await.unwrap(), @r"
         ----- Output Stream -----
@@ -995,7 +981,6 @@ mod tests {
     /// the Parquet morsel planner.
     #[tokio::test]
     async fn morsel_framework_two_cpu_steps_before_morsel() -> Result<()> {
-        let observer = MorselObserver::new();
         let test = MorselTest::new()
             .with_file(
                 "file1.parquet",
@@ -1006,8 +991,7 @@ mod tests {
                     .return_morsel(MorselId(10), 42)
                     .return_none()
                     .build(),
-            )
-            .with_observer(observer.clone());
+            );
 
         insta::assert_snapshot!(test.run().await.unwrap(), @r"
         ----- Output Stream -----
@@ -1033,7 +1017,6 @@ mod tests {
     /// batches produced by any returned child planners.
     #[tokio::test]
     async fn morsel_framework_morsels_before_child_planner() -> Result<()> {
-        let observer = MorselObserver::new();
         let child_planner = MockPlanner::builder()
             .with_id(PlannerId(1))
             .return_morsel(MorselId(11), 43)
@@ -1052,8 +1035,7 @@ mod tests {
             .build();
 
         let test = MorselTest::new()
-            .with_file("file1.parquet", parent_planner)
-            .with_observer(observer.clone());
+            .with_file("file1.parquet", parent_planner);
 
         insta::assert_snapshot!(test.run().await.unwrap(), @r"
         ----- Output Stream -----
@@ -1086,7 +1068,6 @@ mod tests {
     /// 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))
             // Note IO required 2 polls
@@ -1113,8 +1094,7 @@ mod tests {
             .build();
 
         let test = MorselTest::new()
-            .with_file("file1.parquet", parent_planner)
-            .with_observer(observer.clone());
+            .with_file("file1.parquet", parent_planner);
 
         // Expect both futures to be polled, but second planner's (42) batch to be
         // produced first
@@ -1158,7 +1138,7 @@ mod tests {
         // Run same test using `with_preserve_order(true)`, but expect the first
         // planner's batch (41) to be produced before the second's (42), even
         // though the second planner's I/O resolves first.
-        let test = test.reset_observer().with_preserve_order(true);
+        let test = test.with_preserve_order(true);
 
         insta::assert_snapshot!(test.run().await.unwrap(), @r"
         ----- Output Stream -----
@@ -1174,27 +1154,27 @@ mod tests {
         planner_called: PlannerId(0)
         planner_called: PlannerId(1)
         io_future_created: PlannerId(1), IoFutureId(100)
+        io_future_polled: PlannerId(1), IoFutureId(100)
+        io_future_polled: PlannerId(1), IoFutureId(100)
+        io_future_polled: PlannerId(1), IoFutureId(100)
+        io_future_resolved: 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)
+        planner_called: PlannerId(2)
+        morsel_produced: PlannerId(2), MorselId(12)
+        planner_called: PlannerId(2)
+        morsel_stream_started: MorselId(12)
+        morsel_stream_batch_produced: MorselId(12), BatchId(42)
+        morsel_stream_finished: MorselId(12)
         ");
 
         Ok(())
@@ -1205,7 +1185,6 @@ mod tests {
     /// second.
     #[tokio::test]
     async fn morsel_framework_two_files_overlapping_io() -> Result<()> {
-        let observer = MorselObserver::new();
         let test = MorselTest::new()
             .with_file(
                 "file1.parquet",
@@ -1224,8 +1203,7 @@ mod tests {
                     .return_morsel(MorselId(11), 43)
                     .return_none()
                     .build(),
-            )
-            .with_observer(observer.clone());
+            );
 
         insta::assert_snapshot!(test.run().await.unwrap(), @r"
         ----- Output Stream -----

datafusion/datasource/src/morsel/test_utils.rs

@@ -109,10 +109,18 @@ impl MorselObserver {
         Self::default()
     }
 
+    /// Clear any previously buffered events
+    pub fn clear(&self) {
+        self.events.lock().unwrap().clear();
+    }
+
+    /// Push a new [`MorselEvent`]
+
     pub fn push(&self, event: MorselEvent) {
         self.events.lock().unwrap().push(event);
     }
 
+    /// Return a copy of the current list of [`MorselEvents`]
     pub fn events(&self) -> Vec<MorselEvent> {
         self.events.lock().unwrap().clone()
     }
@@ -151,10 +159,9 @@ impl MockMorselizer {
         }
     }
 
-    /// Set the observer
-    pub fn with_observer(mut self, observer: MorselObserver) -> Self {
-        self.observer = observer;
-        self
+    /// Return a reference to the observer
+    pub fn observer(&self) -> &MorselObserver {
+        &self.observer
     }
 
     /// Add the description of how a file will be planned