Allow setting via environment variable

Author: alamb

datafusion/datasource/src/file_stream.rs

@@ -56,8 +56,12 @@ use futures::{FutureExt, Stream, StreamExt as _};
 /// However, setting this too high may lead to more memory buffering and
 /// resource contention if there are too many concurrent IOs.
 ///
-/// TODO make this a config option
-const TARGET_CONCURRENT_PLANNERS: usize = 2;
+/// Default target number of planners that may be active at once for a single
+/// `FileStream`.
+///
+/// This can be overridden temporarily for testing with the
+/// `DATAFUSION_FILESTREAM_TARGET_CONCURRENT_PLANNERS` environment variable.
+const DEFAULT_TARGET_CONCURRENT_PLANNERS: usize = 2;
 
 /// Keep at most this many morsels buffered before pausing additional planning.
 ///
@@ -72,6 +76,14 @@ fn max_buffered_morsels() -> usize {
         .unwrap_or(1)
 }
 
+fn target_concurrent_planners() -> usize {
+    std::env::var("DATAFUSION_FILESTREAM_TARGET_CONCURRENT_PLANNERS")
+        .ok()
+        .and_then(|value| value.parse::<usize>().ok())
+        .filter(|value| *value > 0)
+        .unwrap_or(DEFAULT_TARGET_CONCURRENT_PLANNERS)
+}
+
 /// A stream that iterates record batch by record batch, file over file.
 ///
 /// When running, a FileStream has some number of waiting planners (that are
@@ -128,6 +140,10 @@ pub struct FileStream {
     ///
     /// [`MorselPlan`]: crate::morsel::MorselPlan
     preserve_order: bool,
+    /// Target number of planners that may be active concurrently for this
+    /// stream. This is resolved once during construction so the scheduler
+    /// doesn't repeatedly consult the environment during polling.
+    target_concurrent_planners: usize,
     /// Is the stream complete?
     state: StreamState,
 }
@@ -181,6 +197,7 @@ impl FileStream {
             baseline_metrics: BaselineMetrics::new(metrics, partition),
             on_error: OnError::Fail,
             preserve_order: false,
+            target_concurrent_planners: target_concurrent_planners(),
             state: StreamState::Active,
         })
     }
@@ -208,7 +225,7 @@ impl FileStream {
     /// 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
+        self.waiting_planners.len() >= self.target_concurrent_planners
     }
 
     /// Run a planner on CPU until it either needs I/O or fully completes.
@@ -252,7 +269,7 @@ impl FileStream {
     fn start_next_files(&mut self) -> Result<()> {
         let max_buffered_morsels = max_buffered_morsels();
         while (self.waiting_planners.len() + self.ready_planners.len())
-            < TARGET_CONCURRENT_PLANNERS
+            < self.target_concurrent_planners
         {
             // In ordered mode, do not admit later files while there is any
             // earlier file work still buffered, waiting on I/O, or actively