limit concurrent planners

Author: alamb

datafusion/datasource/src/file_stream.rs

@@ -200,6 +200,17 @@ impl FileStream {
         self
     }
 
+    /// Return true if this stream already has the target number of planner I/O
+    /// futures outstanding.
+    ///
+    /// Child planners returned from a single parent planner should respect the
+    /// same outstanding-I/O cap as top-level file planners. Otherwise a single
+    /// parent plan can fan out into arbitrarily many waiting I/O futures and
+    /// bypass `TARGET_CONCURRENT_PLANNERS`.
+    fn planner_io_at_capacity(&self) -> bool {
+        self.waiting_planners.len() >= TARGET_CONCURRENT_PLANNERS
+    }
+
     /// Run a planner on CPU until it either needs I/O or fully completes.
     ///
     /// Any morsels produced along the way are appended to `self.morsels`. If
@@ -349,6 +360,13 @@ impl FileStream {
                 {
                     break;
                 }
+                // Apply the same outstanding-I/O cap to child planners that we
+                // apply when admitting new files. Once enough planners are
+                // waiting on I/O, leave any additional ready planners queued
+                // until one of the existing I/O futures completes.
+                if self.planner_io_at_capacity() {
+                    break;
+                }
                 let Some(planner) = self.ready_planners.pop_front() else {
                     break;
                 };
@@ -1175,6 +1193,97 @@ mod tests {
         Ok(())
     }
 
+    /// Verifies that child planners still respect the global outstanding-I/O
+    /// cap. Even if a parent planner returns three ready children, only two of
+    /// them should be allowed to create waiting I/O futures at once.
+    #[tokio::test]
+    async fn morsel_framework_child_planner_io_respects_global_cap() -> Result<()> {
+        let planner_1 = MockPlanner::builder()
+            .with_id(PlannerId(1))
+            .return_plan(ReturnPlanBuilder::new().with_io(IoFutureId(100), 1))
+            .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), 3))
+            .return_morsel(MorselId(12), 42)
+            .return_none()
+            .build();
+        let planner_3 = MockPlanner::builder()
+            .with_id(PlannerId(3))
+            .return_plan(ReturnPlanBuilder::new().with_io(IoFutureId(102), 1))
+            .return_morsel(MorselId(13), 43)
+            .return_none()
+            .build();
+
+        let parent_planner = MockPlanner::builder()
+            .with_id(PlannerId(0))
+            .return_plan(
+                ReturnPlanBuilder::new()
+                    .with_planner(planner_1)
+                    .with_planner(planner_2)
+                    .with_planner(planner_3),
+            )
+            .return_none()
+            .build();
+
+        let test = MorselTest::new().with_file("file1.parquet", parent_planner);
+
+        insta::assert_snapshot!(test.run().await.unwrap(), @r"
+        ----- Output Stream -----
+        Batch: 41
+        Batch: 42
+        Batch: 43
+        Done
+        ----- 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(())
+    }
+
     /// Verifies that `FileStream` overlaps planner I/O across multiple files
     /// rather than waiting for the first file to finish before starting the
     /// second.