comments and test with IO

Author: alamb

datafusion/datasource/src/file_stream.rs

@@ -613,8 +613,8 @@ impl FileStreamMetrics {
 mod tests {
     use crate::file_scan_config::{FileScanConfig, FileScanConfigBuilder};
     use crate::morsel::test_utils::{
-        BatchId, MockMorselizer, MockPlanner, MorselEventsBuilder, MorselId,
-        MorselObserver, PlannerId,
+        BatchId, IoFutureId, MockMorselizer, MockPlanner, MorselEventsBuilder, MorselId,
+        MorselObserver, PlannerId, ReturnPlanBuilder,
     };
     use crate::tests::make_partition;
     use crate::{PartitionedFile, TableSchema};
@@ -815,6 +815,49 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn morsel_framework_single_planner_single_io_then_morsel() -> Result<()> {
+        let observer = MorselObserver::new();
+        let morselizer = MockMorselizer::new()
+            .with_file(
+                "file1.parquet",
+                MockPlanner::builder()
+                    .with_id(PlannerId(0))
+                    .return_plan(ReturnPlanBuilder::new().with_io(IoFutureId(100), 1))
+                    .return_morsel(MorselId(10), 42)
+                    .return_none()
+                    .build(),
+            )
+            .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(),
+            MorselEventsBuilder::new()
+                .morselize_file("file1.parquet")
+                .planner_created(PlannerId(0))
+                .planner_called(PlannerId(0))
+                .io_future_created(PlannerId(0), IoFutureId(100))
+                .io_future_polled(PlannerId(0), IoFutureId(100))
+                .io_future_polled(PlannerId(0), IoFutureId(100))
+                .io_future_resolved(PlannerId(0), IoFutureId(100))
+                .planner_called(PlannerId(0))
+                .morsel_produced(PlannerId(0), MorselId(10))
+                .planner_called(PlannerId(0))
+                .morsel_stream_started(MorselId(10))
+                .morsel_stream_batch_produced(MorselId(10), BatchId(42))
+                .morsel_stream_finished(MorselId(10))
+                .build()
+        );
+
+        Ok(())
+    }
+
     fn test_config(file_names: Vec<&str>) -> FileScanConfig {
         let file_group = file_names
             .into_iter()

datafusion/datasource/src/morsel/test_utils.rs

@@ -370,33 +370,53 @@ impl MockPlannerBuilder {
 /// Builder for `PlannerStep::ReturnPlan`.
 #[derive(Debug, Default)]
 pub struct ReturnPlanBuilder {
+    /// Morsels that should be returned immediately by this planner step.
     morsels: Vec<MockMorselSpec>,
+    /// Child planners that should be returned immediately by this planner step.
     planners: Vec<MockPlanner>,
+    /// Identifier for the mock I/O future returned by this step, if any.
     io_future_id: Option<IoFutureId>,
+    /// Number of `Poll::Pending` results the mock I/O future should yield
+    /// before resolving successfully.
+    ///
+    /// This is a deterministic test-only knob. It does not model elapsed time
+    /// or bytes read; it only controls how many scheduler polls are required
+    /// before the mock I/O future becomes ready.
     io_polls: usize,
 }
 
 impl ReturnPlanBuilder {
+    /// Create an empty return-plan builder.
     pub fn new() -> Self {
         Self::default()
     }
 
+    /// Add a morsel that should be returned immediately by this planner step.
     pub fn with_morsel(mut self, morsel: MockMorselSpec) -> Self {
         self.morsels.push(morsel);
         self
     }
 
+    /// Add a child planner that should be returned immediately by this step.
     pub fn with_planner(mut self, planner: MockPlanner) -> Self {
         self.planners.push(planner);
         self
     }
 
+    /// Return a mock I/O future from this step.
+    ///
+    /// `io_future_id` is recorded in the emitted `MorselEvent`s so tests can
+    /// distinguish multiple I/O phases from the same planner.
+    ///
+    /// `io_polls` controls how many times that future returns `Poll::Pending`
+    /// before it resolves with `Poll::Ready(Ok(()))`.
     pub fn with_io(mut self, io_future_id: IoFutureId, io_polls: usize) -> Self {
         self.io_future_id = Some(io_future_id);
         self.io_polls = io_polls;
         self
     }
 
+    /// Build the corresponding [`PlannerStep::ReturnPlan`]
     pub fn build(self) -> PlannerStep {
         PlannerStep::ReturnPlan {
             morsels: self.morsels,