test summary mode

Author: alamb

datafusion/datasource/src/file_stream.rs

@@ -870,6 +870,7 @@ mod tests {
         morselizer: MockMorselizer,
         file_names: Vec<String>,
         preserve_order: bool,
+        event_summaries: bool,
     }
 
     impl MorselTest {
@@ -878,6 +879,7 @@ mod tests {
                 morselizer: MockMorselizer::new(),
                 file_names: vec![],
                 preserve_order: false,
+                event_summaries: false,
             }
         }
 
@@ -893,6 +895,16 @@ mod tests {
             self
         }
 
+        /// Snapshot only the higher-level scheduler events.
+        ///
+        /// The full event trace is still useful for detailed tests, but for
+        /// more complex tests those lower level events obscure the important
+        /// events.
+        fn with_event_summaries(mut self) -> Self {
+            self.event_summaries = true;
+            self
+        }
+
         async fn run(self) -> Result<String> {
             // handle to shared observer
             let observer = self.morselizer.observer().clone();
@@ -936,7 +948,12 @@ mod tests {
             // maintaining long hand-written event assertions separately.
             let mut parts = vec!["----- Output Stream -----".to_string(), output];
             parts.push("----- File Stream Events -----".to_string());
-            parts.push(observer.format_events());
+            let events = if self.event_summaries {
+                observer.format_summary_events()
+            } else {
+                observer.format_events()
+            };
+            parts.push(events);
             Ok(parts.join("\n"))
         }
     }
@@ -1066,7 +1083,9 @@ mod tests {
             .return_none()
             .build();
 
-        let test = MorselTest::new().with_file("file1.parquet", parent_planner);
+        let test = MorselTest::new()
+            .with_file("file1.parquet", parent_planner)
+            .with_event_summaries();
 
         insta::assert_snapshot!(test.run().await.unwrap(), @r"
         ----- Output Stream -----
@@ -1076,19 +1095,11 @@ mod tests {
         ----- File Stream Events -----
         morselize_file: file1.parquet
         planner_created: PlannerId(0)
-        planner_called: PlannerId(0)
         morsel_produced: PlannerId(0), MorselId(10)
         planner_produced_child: PlannerId(0) -> PlannerId(1)
-        planner_called: PlannerId(0)
-        planner_called: PlannerId(1)
         morsel_produced: PlannerId(1), MorselId(11)
-        planner_called: PlannerId(1)
-        morsel_stream_started: MorselId(10)
         morsel_stream_batch_produced: MorselId(10), BatchId(42)
-        morsel_stream_finished: MorselId(10)
-        morsel_stream_started: MorselId(11)
         morsel_stream_batch_produced: MorselId(11), BatchId(43)
-        morsel_stream_finished: MorselId(11)
         ");
 
         Ok(())
@@ -1245,8 +1256,11 @@ mod tests {
             .return_none()
             .build();
 
-        let test = MorselTest::new().with_file("file1.parquet", parent_planner);
+        let test = MorselTest::new()
+            .with_file("file1.parquet", parent_planner)
+            .with_event_summaries();
 
+        // Note that the future for planner 1 must resolve before planner 2 begins
         insta::assert_snapshot!(test.run().await.unwrap(), @r"
         ----- Output Stream -----
         Batch: 41
@@ -1256,46 +1270,21 @@ mod tests {
         ----- File Stream Events -----
         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_produced_child: PlannerId(0) -> PlannerId(3)
-        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_resolved: PlannerId(1), IoFutureId(100)
-        io_future_polled: PlannerId(2), IoFutureId(101)
-        planner_called: PlannerId(3)
         io_future_created: PlannerId(3), IoFutureId(102)
-        io_future_polled: PlannerId(2), IoFutureId(101)
-        io_future_polled: PlannerId(3), IoFutureId(102)
-        io_future_polled: PlannerId(2), IoFutureId(101)
         io_future_resolved: PlannerId(2), IoFutureId(101)
-        io_future_polled: PlannerId(3), IoFutureId(102)
         io_future_resolved: PlannerId(3), IoFutureId(102)
-        planner_called: PlannerId(1)
         morsel_produced: PlannerId(1), MorselId(11)
-        planner_called: PlannerId(1)
-        planner_called: PlannerId(2)
         morsel_produced: PlannerId(2), MorselId(12)
-        planner_called: PlannerId(2)
-        planner_called: PlannerId(3)
         morsel_produced: PlannerId(3), MorselId(13)
-        planner_called: PlannerId(3)
-        morsel_stream_started: MorselId(11)
         morsel_stream_batch_produced: MorselId(11), BatchId(41)
-        morsel_stream_finished: MorselId(11)
-        morsel_stream_started: MorselId(12)
         morsel_stream_batch_produced: MorselId(12), BatchId(42)
-        morsel_stream_finished: MorselId(12)
-        morsel_stream_started: MorselId(13)
         morsel_stream_batch_produced: MorselId(13), BatchId(43)
-        morsel_stream_finished: MorselId(13)
         ");
 
         Ok(())

datafusion/datasource/src/morsel/test_utils.rs

@@ -138,6 +138,38 @@ impl MorselObserver {
             .collect::<Vec<_>>()
             .join("\n")
     }
+
+    /// Format only the high-level scheduler events.
+    ///
+    /// This is useful for tests where the exact `plan()` / `poll()` interleave
+    /// is not the main point of the assertion. Summary snapshots are easier to
+    /// review when validating broader behavior such as ordering or outstanding
+    /// I/O limits.
+    pub fn format_summary_events(&self) -> String {
+        self.events()
+            .into_iter()
+            .filter(MorselEvent::is_summary_event)
+            .map(|event| event.to_string())
+            .collect::<Vec<_>>()
+            .join("\n")
+    }
+}
+
+impl MorselEvent {
+    /// Return true for the higher-level events that are most useful in compact
+    /// scheduler snapshots.
+    pub fn is_summary_event(&self) -> bool {
+        matches!(
+            self,
+            MorselEvent::MorselizeFile { .. }
+                | MorselEvent::PlannerCreated { .. }
+                | MorselEvent::PlannerProducedChild { .. }
+                | MorselEvent::IoFutureCreated { .. }
+                | MorselEvent::IoFutureResolved { .. }
+                | MorselEvent::MorselProduced { .. }
+                | MorselEvent::MorselStreamBatchProduced { .. }
+        )
+    }
 }
 
 /// Test [`Morselizer`] that maps file paths to a fixed set of planner specs.