less abstraction

Author: alamb

datafusion/datasource/src/file_stream.rs

@@ -62,7 +62,7 @@ pub struct FileStream {
     /// Any morsels that have already been prepared for processing
     /// (TODO steal morsels from other streams)
     morsels: Vec<Box<dyn Morsel>>,
-    /// The current RecordBatches, if any
+    /// The current reader, if any
     reader: Option<BoxStream<'static, Result<RecordBatch>>>,
     /// The stream schema (file schema including partition columns and after
     /// projection).
@@ -133,7 +133,8 @@ impl FileStream {
             .expect("planner queue checked to be non-empty");
         self.file_stream_metrics.time_opening.start();
         self.state = FileStreamState::Active {
-            morsel_plan: MorselPlan::new(planner),
+            planner,
+            io_future: None,
         };
         Ok(true)
     }
@@ -155,7 +156,8 @@ impl FileStream {
             .expect("planner queue checked to be non-empty");
         self.file_stream_metrics.time_opening.start();
         self.state = FileStreamState::Active {
-            morsel_plan: MorselPlan::new(planner),
+            planner,
+            io_future: None,
         };
         Ok(())
     }
@@ -183,19 +185,16 @@ impl FileStream {
     /// stream is either pending or has produced/failed with a batch.
     fn poll_active_state(
         &mut self,
-        mut morsel_plan: MorselPlan,
+        planner: Arc<dyn MorselPlanner>,
+        mut io_future: Option<BoxFuture<'static, Result<()>>>,
         cx: &mut Context<'_>,
     ) -> StatePoll {
         // State transitions for the active planner:
         // 1. Poll any outstanding planner I/O before touching batch readers.
-        // 2. If a morsel is already ready, promote/drain it before asking the
-        //    planner for more work.
-        // 3. If there is no queued morsel, run the planner's CPU-only `plan`
-        //    step to produce more work or discover completion.
 
         let mut io_completed = false;
-        if let Some(mut io_future) = morsel_plan.take_io_future() {
-            match io_future.as_mut().poll(cx) {
+        if let Some(mut future) = io_future.take() {
+            match future.as_mut().poll(cx) {
                 Poll::Ready(Ok(())) => {
                     self.file_stream_metrics.files_opened.add(1);
                     self.file_stream_metrics.time_opening.stop();
@@ -217,25 +216,32 @@ impl FileStream {
                     }
                 }
                 Poll::Pending => {
-                    morsel_plan.set_io_future(io_future);
-                    self.state = FileStreamState::Active { morsel_plan };
+                    self.state = FileStreamState::Active {
+                        planner,
+                        io_future: Some(future),
+                    };
                     return StatePoll::Return(Poll::Pending);
                 }
             }
         }
 
+        // If there is no active reader, start one from the next available morsel
         if self.reader.is_none() {
+            // TODO: also try and drive any other planners to get more morsels /
+            // discover more IO to before launching into the next morsel.
             if let Some(morsel) = self.morsels.pop() {
                 self.reader = Some(morsel.into_stream());
                 self.file_stream_metrics.time_scanning_until_data.start();
                 self.file_stream_metrics.time_scanning_total.start();
             }
         }
 
+
+        // drive to the next record batch from the active reader, if any
         if let Some(reader) = self.reader.as_mut() {
             match reader.poll_next_unpin(cx) {
                 Poll::Pending => {
-                    self.state = FileStreamState::Active { morsel_plan };
+                    self.state = FileStreamState::Active { planner, io_future };
                     return StatePoll::Return(Poll::Pending);
                 }
                 Poll::Ready(Some(Ok(batch))) => {
@@ -258,7 +264,7 @@ impl FileStream {
                         None => batch,
                     };
                     self.file_stream_metrics.time_scanning_total.start();
-                    self.state = FileStreamState::Active { morsel_plan };
+                    self.state = FileStreamState::Active { planner, io_future };
                     return StatePoll::Return(Poll::Ready(Some(Ok(batch))));
                 }
                 Poll::Ready(Some(Err(err))) => {
@@ -293,14 +299,16 @@ impl FileStream {
             }
         }
 
-        if io_completed || !morsel_plan.has_io_future() {
-            let planner = morsel_plan.planner();
-            let current_plan = mem::replace(&mut morsel_plan, MorselPlan::new(planner));
+        if io_completed || io_future.is_none() {
+            let current_plan = MorselPlan::new(Arc::clone(&planner));
             match current_plan.plan() {
                 Ok(Some(mut new_plan)) => {
                     self.morsels.extend(new_plan.take_morsels());
+                    let next_planner = new_plan.planner();
+                    let next_io_future = new_plan.take_io_future();
                     self.state = FileStreamState::Active {
-                        morsel_plan: new_plan,
+                        planner: next_planner,
+                        io_future: next_io_future,
                     };
                     return StatePoll::Continue;
                 }
@@ -329,7 +337,7 @@ impl FileStream {
             }
         }
 
-        self.state = FileStreamState::Active { morsel_plan };
+        self.state = FileStreamState::Active { planner, io_future };
         StatePoll::Return(Poll::Pending)
     }
 
@@ -340,10 +348,12 @@ impl FileStream {
                     StatePoll::Continue => continue,
                     StatePoll::Return(result) => return result,
                 },
-                FileStreamState::Active { morsel_plan } => match self.poll_active_state(morsel_plan, cx) {
-                    StatePoll::Continue => continue,
-                    StatePoll::Return(result) => return result,
-                },
+                FileStreamState::Active { planner, io_future } => {
+                    match self.poll_active_state(planner, io_future, cx) {
+                        StatePoll::Continue => continue,
+                        StatePoll::Return(result) => return result,
+                    }
+                }
                 FileStreamState::Error | FileStreamState::Limit => {
                     return Poll::Ready(None);
                 }
@@ -422,12 +432,12 @@ pub enum NextOpen {
 pub enum FileStreamState {
     /// The idle state, no file is currently being read
     Idle,
-    /// Actively planning morsels
+    /// Actively managing a single planner and its in-flight I/O, if any.
     Active {
-        /// Plan to get more readers
-        ///
-        /// TODO make make the number of outstanding plans configurable to allow more parallelism
-        morsel_plan: MorselPlan
+        /// Planner that will produce the next morsels for this file.
+        planner: Arc<dyn MorselPlanner>,
+        /// In-flight I/O for the active planner, if it is currently blocked.
+        io_future: Option<BoxFuture<'static, Result<()>>>,
     },
     /// Encountered an error
     Error,

datafusion/datasource/src/morsel/mod.rs

@@ -74,9 +74,9 @@ pub trait Morselizer: Send + Sync + Debug {
 /// A Morsel Planner is responsible for creating morsels for a given scan.
 ///
 /// The MorselPlanner is the unit of I/O -- there is only ever a single IO
-/// outstanding for a MorselPlanner. DataFusion will potentially run multiple
-/// MorselPlanners in parallel which corresponds to multiple parallel I/O
-/// requests.
+/// outstanding for a specific  MorselPlanner. DataFusion will potentially run
+/// multiple MorselPlanners in parallel which corresponds to multiple parallel
+/// I/O requests.
 ///
 /// It is not a Rust [`Stream`] so that it can explicitly separate CPU bound
 /// work from IO work.