datasource: opt-in open-ahead so FileStream overlaps the next file's open with scanning

datafusion/common/src/config.rs

@@ -1039,6 +1039,16 @@ config_namespace! {
         /// runtime reassignment occurs.
         pub enable_file_stream_work_stealing: bool, default = true
 
+        /// When `true`, each file-scan partition starts opening its next file
+        /// (footer, page index and bloom filter I/O) while the current file is
+        /// still being scanned, keeping at most one file in flight ahead of
+        /// the active reader. This overlaps per-file open latency with data
+        /// reads, which matters for scans over many files on high-latency
+        /// object stores. The cost is holding one extra file's metadata per
+        /// partition. When `false` (the default) a partition opens the next
+        /// file only after the current one is fully consumed.
+        pub file_stream_open_ahead: bool, default = false
+
         /// Aggregation ratio (number of distinct groups / number of input rows)
         /// threshold for skipping partial aggregation. If the value is greater
         /// then partial aggregation will skip aggregation for further input

datafusion/datasource/src/file_scan_config/mod.rs

@@ -720,11 +720,18 @@ impl DataSource for FileScanConfig {
             .and_then(|state| state.downcast_ref::<SharedWorkSource>())
             .cloned();
 
+        let open_ahead = context
+            .session_config()
+            .options()
+            .execution
+            .file_stream_open_ahead;
+
         let stream = FileStreamBuilder::new(self)
             .with_partition(partition)
             .with_shared_work_source(shared_work_source)
             .with_morselizer(morselizer)
             .with_metrics(source.metrics())
+            .with_open_ahead(open_ahead)
             .build()?;
         Ok(Box::pin(cooperative(stream)))
     }

datafusion/datasource/src/file_stream/builder.rs

@@ -35,6 +35,7 @@ pub struct FileStreamBuilder<'a> {
     metrics: Option<&'a ExecutionPlanMetricsSet>,
     on_error: OnError,
     shared_work_source: Option<SharedWorkSource>,
+    open_ahead: bool,
 }
 
 impl<'a> FileStreamBuilder<'a> {
@@ -47,6 +48,7 @@ impl<'a> FileStreamBuilder<'a> {
             metrics: None,
             on_error: OnError::Fail,
             shared_work_source: None,
+            open_ahead: false,
         }
     }
 
@@ -93,6 +95,14 @@ impl<'a> FileStreamBuilder<'a> {
         self
     }
 
+    /// Configure whether the stream opens its next file while the current
+    /// file is still being scanned (see
+    /// `datafusion.execution.file_stream_open_ahead`).
+    pub fn with_open_ahead(mut self, open_ahead: bool) -> Self {
+        self.open_ahead = open_ahead;
+        self
+    }
+
     /// Build the configured [`FileStream`].
     pub fn build(self) -> Result<FileStream> {
         let Self {
@@ -102,6 +112,7 @@ impl<'a> FileStreamBuilder<'a> {
             metrics,
             on_error,
             shared_work_source,
+            open_ahead,
         } = self;
 
         let Some(partition) = partition else {
@@ -130,6 +141,7 @@ impl<'a> FileStreamBuilder<'a> {
             config.limit,
             morselizer,
             on_error,
+            open_ahead,
             file_stream_metrics,
         ));
 

datafusion/datasource/src/file_stream/metrics.rs

@@ -33,6 +33,13 @@ impl StartableTime {
         self.start = Some(Instant::now());
     }
 
+    /// Start the timer unless it is already running.
+    pub fn start_if_stopped(&mut self) {
+        if self.start.is_none() {
+            self.start = Some(Instant::now());
+        }
+    }
+
     pub fn stop(&mut self) {
         if let Some(start) = self.start.take() {
             self.metrics.add_elapsed(start);

datafusion/datasource/src/file_stream/mod.rs

@@ -763,6 +763,77 @@ mod tests {
         Ok(())
     }
 
+    /// Verifies that with `open_ahead` the next file is morselized and its
+    /// I/O is issued while the previous file's morsel is still producing
+    /// batches, so per-file open latency overlaps with reading.
+    #[tokio::test]
+    async fn morsel_open_ahead_overlaps_next_file_io() -> Result<()> {
+        let test = FileStreamMorselTest::new()
+            .with_open_ahead(true)
+            .with_file(
+                MockPlanner::builder("file1.parquet")
+                    .add_plan(
+                        PendingPlannerBuilder::new(IoFutureId(1))
+                            .with_polls_to_resolve(PollsToResolve(1)),
+                    )
+                    .add_plan(
+                        MockPlanBuilder::new()
+                            .with_morsel_batches(MorselId(10), vec![42, 43]),
+                    )
+                    .return_none(),
+            )
+            .with_file(
+                MockPlanner::builder("file2.parquet")
+                    .add_plan(
+                        PendingPlannerBuilder::new(IoFutureId(2))
+                            .with_polls_to_resolve(PollsToResolve(1)),
+                    )
+                    .add_plan(
+                        MockPlanBuilder::new()
+                            .with_morsel_batches(MorselId(20), vec![44]),
+                    )
+                    .return_none(),
+            );
+
+        // file2 is morselized and its I/O issued while morsel 10 is still
+        // producing batches; the I/O resolves before morsel 10 finishes.
+        insta::assert_snapshot!(test.run().await.unwrap(), @r"
+        ----- Output Stream -----
+        Batch: 42
+        Batch: 43
+        Batch: 44
+        Done
+        ----- File Stream Events -----
+        morselize_file: file1.parquet
+        planner_created: file1.parquet
+        planner_called: file1.parquet
+        io_future_created: file1.parquet, IoFutureId(1)
+        io_future_polled: file1.parquet, IoFutureId(1)
+        io_future_polled: file1.parquet, IoFutureId(1)
+        io_future_resolved: file1.parquet, IoFutureId(1)
+        planner_called: file1.parquet
+        morsel_produced: file1.parquet, MorselId(10)
+        morsel_stream_started: MorselId(10)
+        morselize_file: file2.parquet
+        planner_created: file2.parquet
+        planner_called: file2.parquet
+        io_future_created: file2.parquet, IoFutureId(2)
+        morsel_stream_batch_produced: MorselId(10), BatchId(42)
+        io_future_polled: file2.parquet, IoFutureId(2)
+        morsel_stream_batch_produced: MorselId(10), BatchId(43)
+        io_future_polled: file2.parquet, IoFutureId(2)
+        io_future_resolved: file2.parquet, IoFutureId(2)
+        planner_called: file2.parquet
+        morsel_produced: file2.parquet, MorselId(20)
+        morsel_stream_finished: MorselId(10)
+        morsel_stream_started: MorselId(20)
+        morsel_stream_batch_produced: MorselId(20), BatchId(44)
+        morsel_stream_finished: MorselId(20)
+        ");
+
+        Ok(())
+    }
+
     /// Verifies that a planner can traverse two sequential I/O phases before
     /// producing one batch, similar to Parquet.
     #[tokio::test]
@@ -1368,6 +1439,7 @@ mod tests {
         preserve_order: bool,
         declared_output_partitioning: bool,
         enable_file_stream_work_stealing: bool,
+        open_ahead: bool,
         file_stream_events: bool,
         build_streams_on_first_read: bool,
         reads: Vec<PartitionId>,
@@ -1383,6 +1455,7 @@ mod tests {
                 preserve_order: false,
                 declared_output_partitioning: false,
                 enable_file_stream_work_stealing: true,
+                open_ahead: false,
                 file_stream_events: true,
                 build_streams_on_first_read: false,
                 reads: vec![],
@@ -1436,6 +1509,14 @@ mod tests {
             self
         }
 
+        /// Sets `datafusion.execution.file_stream_open_ahead`: when enabled,
+        /// each stream starts planning its next file while the active reader
+        /// is still producing batches.
+        fn with_open_ahead(mut self, open_ahead: bool) -> Self {
+            self.open_ahead = open_ahead;
+            self
+        }
+
         /// Controls whether scheduler events are included in the snapshot.
         ///
         /// When disabled, `run()` still includes the event section header but
@@ -1527,6 +1608,7 @@ mod tests {
                         .with_shared_work_source(shared_work_source.clone())
                         .with_morselizer(Box::new(self.morselizer.clone()))
                         .with_metrics(&metrics_set)
+                        .with_open_ahead(self.open_ahead)
                         .build()?;
                     partitions[partition].set_stream(stream);
                 }
@@ -1554,6 +1636,7 @@ mod tests {
                         .with_shared_work_source(shared_work_source.clone())
                         .with_morselizer(Box::new(self.morselizer.clone()))
                         .with_metrics(&metrics_set)
+                        .with_open_ahead(self.open_ahead)
                         .build()?;
                     partition_state.set_stream(stream);
                 }

datafusion/datasource/src/file_stream/scan_state.rs

@@ -81,6 +81,9 @@ pub(super) struct ScanState {
     /// Once the I/O completes, yields the next planner and is pushed back
     /// onto `ready_planners`.
     pending_planner: Option<PendingMorselPlanner>,
+    /// Whether to start opening the next file while the active reader is
+    /// still producing batches. See [`Self::advance_open_ahead`].
+    open_ahead: bool,
     /// Metrics for the active scan queues.
     metrics: FileStreamMetrics,
 }
@@ -91,6 +94,7 @@ impl ScanState {
         remain: Option<usize>,
         morselizer: Box<dyn Morselizer>,
         on_error: OnError,
+        open_ahead: bool,
         metrics: FileStreamMetrics,
     ) -> Self {
         Self {
@@ -102,6 +106,7 @@ impl ScanState {
             ready_morsels: Default::default(),
             reader: None,
             pending_planner: None,
+            open_ahead,
             metrics,
         }
     }
@@ -115,16 +120,20 @@ impl ScanState {
     ///
     /// Work is attempted in this order:
     /// 1. resolve any pending planner I/O
-    /// 2. poll the active reader
-    /// 3. turn a ready morsel into the active reader
-    /// 4. run CPU planning on a ready planner
-    /// 5. morselize the next unopened file
+    /// 2. with `open_ahead`, start planning the next file behind the active
+    ///    reader (see [`Self::advance_open_ahead`])
+    /// 3. poll the active reader
+    /// 4. turn a ready morsel into the active reader
+    /// 5. run CPU planning on a ready planner
+    /// 6. morselize the next unopened file
     ///
     /// The return [`ScanAndReturn`] tells `poll_inner` how to update the
     /// outer `FileStreamState`.
     pub(super) fn poll_scan(&mut self, cx: &mut Context<'_>) -> ScanAndReturn {
-        let _processing_timer: ScopedTimerGuard<'_> =
-            self.metrics.time_processing.timer();
+        // Guard a clone of the shared timer so the borrow does not pin
+        // `self.metrics` for the whole poll (the look-ahead needs `&mut self`).
+        let time_processing = self.metrics.time_processing.clone();
+        let _processing_timer: ScopedTimerGuard<'_> = time_processing.timer();
 
         // Try and resolve outstanding IO first. If it is still pending, check
         // the current reader or ready morsels before yielding. New planning
@@ -154,6 +163,16 @@ impl ScanState {
             }
         }
 
+        // With open-ahead enabled, use the time the active reader is busy to
+        // get the next file's planning (and its single outstanding I/O)
+        // under way, so per-file open latency overlaps with data reads.
+        if self.open_ahead
+            && self.reader.is_some()
+            && let Some(ret) = self.advance_open_ahead()
+        {
+            return ret;
+        }
+
         // Next try and get the next batch from the active reader, if any.
         if let Some(reader) = self.reader.as_mut() {
             match reader.poll_next_unpin(cx) {
@@ -225,6 +244,12 @@ impl ScanState {
         // still outstanding because they may need additional IO and ScanState
         // currently only permits a single outstanding IO
         if self.pending_planner.is_some() {
+            if self.open_ahead {
+                // The next file was claimed while a reader was active without
+                // starting the opening timer; only the wait that is actually
+                // exposed (no reader to drain) counts as opening time.
+                self.metrics.time_opening.start_if_stopped();
+            }
             return ScanAndReturn::Return(Poll::Pending);
         }
 
@@ -291,6 +316,86 @@ impl ScanState {
     }
 }
 
+impl ScanState {
+    /// Drives planning of the *next* file while a reader is active, keeping at
+    /// most one file in flight ahead of it.
+    ///
+    /// Claims the next unopened file once nothing else is queued and runs its
+    /// CPU planning until it either yields a ready morsel or blocks on I/O.
+    /// The single-outstanding-I/O rule still holds: the pending planner is
+    /// polled by [`Self::poll_scan`] before the reader on every iteration, so
+    /// the file's open latency overlaps with the reader's data reads instead
+    /// of following them.
+    ///
+    /// The opening timer is deliberately not started here: time spent opening
+    /// a file behind an active reader is not exposed to the query. It is
+    /// started by `poll_scan` only if the stream later has to wait for that
+    /// open with no reader left to drain.
+    ///
+    /// Returns `Some` when the outer loop must return (an error surfaced by
+    /// the look-ahead), `None` to continue with the active reader.
+    fn advance_open_ahead(&mut self) -> Option<ScanAndReturn> {
+        loop {
+            // The next file is already either fully open or waiting on I/O.
+            if self.pending_planner.is_some() || !self.ready_morsels.is_empty() {
+                return None;
+            }
+
+            if let Some(planner) = self.ready_planners.pop_front() {
+                match planner.plan() {
+                    Ok(Some(mut plan)) => {
+                        self.ready_morsels.extend(plan.take_morsels());
+                        self.ready_planners.extend(plan.take_ready_planners());
+                        if let Some(pending_planner) = plan.take_pending_planner() {
+                            if self.pending_planner.is_some() {
+                                return Some(ScanAndReturn::Error(
+                                    internal_datafusion_err!(
+                                        "Conflicting pending planner state in FileStream ScanState"
+                                    ),
+                                ));
+                            }
+                            self.pending_planner = Some(pending_planner);
+                        }
+                    }
+                    // Nothing to read from this file (e.g. pruned late).
+                    Ok(None) => {
+                        self.metrics.files_processed.add(1);
+                    }
+                    Err(err) => {
+                        self.metrics.file_open_errors.add(1);
+                        match self.on_error {
+                            OnError::Skip => {
+                                self.metrics.files_processed.add(1);
+                            }
+                            OnError::Fail => return Some(ScanAndReturn::Error(err)),
+                        }
+                    }
+                }
+                continue;
+            }
+
+            // Nothing queued for the next file: claim it (`None` once the
+            // work source is drained: keep serving the active reader).
+            let part_file = self.work_source.pop_front()?;
+            match self.morselizer.plan_file(part_file) {
+                Ok(planner) => {
+                    self.metrics.files_opened.add(1);
+                    self.ready_planners.push_back(planner);
+                }
+                Err(err) => {
+                    self.metrics.file_open_errors.add(1);
+                    match self.on_error {
+                        OnError::Skip => {
+                            self.metrics.files_processed.add(1);
+                        }
+                        OnError::Fail => return Some(ScanAndReturn::Error(err)),
+                    }
+                }
+            }
+        }
+    }
+}
+
 /// What should be done on the next iteration of [`ScanState::poll_scan`]?
 pub(super) enum ScanAndReturn {
     /// Poll again.