fix(datasource): scope file scan work stealing to opened partitions

datafusion/datasource/src/file_scan_config/mod.rs

@@ -1101,15 +1101,19 @@ impl DataSource for FileScanConfig {
     /// Create any shared state that should be passed between sibling streams
     /// during one execution.
     ///
+    /// The shared source starts empty. Each opened partition registers its own
+    /// file group, so sibling streams can only steal work from partitions that
+    /// are actually executing in this process.
+    ///
     /// This returns `None` when sibling streams must not share work, such as
     /// when file order must be preserved or the file groups define the output
-    /// partitioning needed for the rest of the plan
+    /// partitioning needed for the rest of the plan.
     fn create_sibling_state(&self) -> Option<Arc<dyn Any + Send + Sync>> {
         if self.preserve_order || self.partitioned_by_file_group {
             return None;
         }
 
-        Some(Arc::new(SharedWorkSource::from_config(self)) as Arc<dyn Any + Send + Sync>)
+        Some(Arc::new(SharedWorkSource::empty()) as Arc<dyn Any + Send + Sync>)
     }
 }
 

datafusion/datasource/src/file_stream/builder.rs

@@ -114,14 +114,17 @@ impl<'a> FileStreamBuilder<'a> {
             return internal_err!("FileStreamBuilder missing required metrics");
         };
         let projected_schema = config.projected_schema()?;
-        let Some(file_group) = config.file_groups.get(partition).cloned() else {
+        let Some(file_group) = config.file_groups.get(partition) else {
             return internal_err!(
                 "FileStreamBuilder invalid partition index: {partition}"
             );
         };
         let work_source = match shared_work_source {
-            Some(shared) => WorkSource::Shared(shared),
-            None => WorkSource::Local(file_group.into_inner().into()),
+            Some(shared) => {
+                shared.register_partition(partition, config);
+                WorkSource::Shared(shared)
+            }
+            None => WorkSource::Local(file_group.clone().into_inner().into()),
         };
 
         let file_stream_metrics = FileStreamMetrics::new(metrics, partition);

datafusion/datasource/src/file_stream/mod.rs

@@ -1077,6 +1077,31 @@ mod tests {
         Ok(())
     }
 
+    /// Verifies that a process that opens only one partition only registers
+    /// that partition's files in the shared queue.
+    #[tokio::test]
+    async fn morsel_only_opened_partition_is_registered() -> Result<()> {
+        let test = two_partition_morsel_test();
+        let config = test.test_config();
+        let shared_work_source = config
+            .create_sibling_state()
+            .and_then(|state| state.as_ref().downcast_ref::<SharedWorkSource>().cloned())
+            .expect("shared work source");
+        let metrics = ExecutionPlanMetricsSet::new();
+
+        let stream = FileStreamBuilder::new(&config)
+            .with_partition(1)
+            .with_shared_work_source(Some(shared_work_source))
+            .with_morselizer(Box::new(test.morselizer))
+            .with_metrics(&metrics)
+            .build()?;
+
+        insta::assert_snapshot!(drain_stream_output(stream).await?, @"Batch: 201");
+        assert_eq!(metric_count(&metrics, "files_opened"), 1);
+
+        Ok(())
+    }
+
     /// Verifies that a stream that must preserve order keeps its files local
     /// and therefore cannot steal from a sibling shared queue.
     #[tokio::test]
@@ -1166,11 +1191,11 @@ mod tests {
         Ok(())
     }
 
-    /// Ensures that if a sibling is built and polled
-    /// before another sibling has been built and contributed its files to the
-    /// shared queue, the first sibling does not finish prematurely.
+    /// Ensures that if an empty sibling is built and polled before another
+    /// partition registers its files, the later partition still reads its own
+    /// files.
     #[tokio::test]
-    async fn morsel_empty_sibling_can_finish_before_shared_work_exists() -> Result<()> {
+    async fn morsel_empty_sibling_before_late_registration() -> Result<()> {
         let test = FileStreamMorselTest::new()
             .with_file_in_partition(
                 PartitionId(0),
@@ -1185,21 +1210,19 @@ mod tests {
                     .return_none(),
             )
             // Build streams lazily so partition 1 can poll the shared queue
-            // before partition 0 has contributed its files. Once partition 0
-            // is built, a later poll of partition 1 can still steal one of
-            // them from the shared queue.
+            // before partition 0 has registered its files.
             .with_build_streams_on_first_read(true)
             .with_reads(vec![PartitionId(1), PartitionId(0), PartitionId(1)])
             .with_file_stream_events(false);
 
-        // Partition 1 polls too early once, then later steals one file after
-        // partition 0 has populated the shared queue.
+        // Partition 1 polls too early and finishes. Partition 0 later
+        // registers and reads its own files; no work is lost.
         insta::assert_snapshot!(test.run().await.unwrap(), @r"
         ----- Partition 0 -----
+        Batch: 101
         Batch: 102
         Done
         ----- Partition 1 -----
-        Batch: 101
         Done
         ----- File Stream Events -----
         (omitted due to with_file_stream_events(false))
@@ -1222,18 +1245,18 @@ mod tests {
         let limited_metrics = ExecutionPlanMetricsSet::new();
         let unlimited_metrics = ExecutionPlanMetricsSet::new();
 
-        let limited_stream = FileStreamBuilder::new(&limited_config)
-            .with_partition(1)
+        let unlimited_stream = FileStreamBuilder::new(&unlimited_config)
+            .with_partition(0)
             .with_shared_work_source(Some(shared_work_source.clone()))
             .with_morselizer(Box::new(test.morselizer.clone()))
-            .with_metrics(&limited_metrics)
+            .with_metrics(&unlimited_metrics)
             .build()?;
 
-        let unlimited_stream = FileStreamBuilder::new(&unlimited_config)
-            .with_partition(0)
+        let limited_stream = FileStreamBuilder::new(&limited_config)
+            .with_partition(1)
             .with_shared_work_source(Some(shared_work_source))
             .with_morselizer(Box::new(test.morselizer))
-            .with_metrics(&unlimited_metrics)
+            .with_metrics(&limited_metrics)
             .build()?;
 
         let limited_output = drain_stream_output(limited_stream).await?;
@@ -1406,7 +1429,7 @@ mod tests {
         ///
         /// The default builds all streams before polling begins, which matches
         /// normal execution. Tests may enable lazy creation to model races
-        /// where one sibling polls before another has contributed its files to
+        /// where one sibling polls before another has registered its files in
         /// the shared queue.
         fn with_build_streams_on_first_read(
             mut self,
@@ -1449,30 +1472,18 @@ mod tests {
                 .map(|_| PartitionState::new())
                 .collect::<Vec<_>>();
 
-            let mut build_order = Vec::new();
-            for partition in self.reads.iter().map(|partition| partition.0) {
-                if !build_order.contains(&partition) {
-                    build_order.push(partition);
-                }
-            }
-            for partition in 0..partition_count {
-                if !build_order.contains(&partition) {
-                    build_order.push(partition);
-                }
-            }
-
             let config = self.test_config();
             // `DataSourceExec::execute` creates one execution-local shared
             // state object via `create_sibling_state()` and then passes it
-            // to `open_with_sibling_state(...)`. These tests build
-            // `FileStream`s directly, bypassing `DataSourceExec`, so they must
-            // perform the same setup explicitly when exercising sibling-stream
-            // work stealing.
+            // to `open_with_args(...)`. These tests build `FileStream`s
+            // directly, bypassing `DataSourceExec`, so they must perform the
+            // same setup explicitly when exercising sibling-stream work
+            // stealing.
             let shared_work_source = config.create_sibling_state().and_then(|state| {
                 state.as_ref().downcast_ref::<SharedWorkSource>().cloned()
             });
             if !self.build_streams_on_first_read {
-                for partition in build_order {
+                for partition in 0..partition_count {
                     let stream = FileStreamBuilder::new(&config)
                         .with_partition(partition)
                         .with_shared_work_source(shared_work_source.clone())

datafusion/datasource/src/file_stream/work_source.rs

@@ -15,11 +15,10 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::collections::VecDeque;
+use std::collections::{HashSet, VecDeque};
 use std::sync::Arc;
 
 use crate::PartitionedFile;
-use crate::file_groups::FileGroup;
 use crate::file_scan_config::FileScanConfig;
 use parking_lot::Mutex;
 
@@ -55,55 +54,66 @@ impl WorkSource {
     }
 }
 
-/// Shared source of work for sibling `FileStream`s
+/// Shared source of work for sibling `FileStream`s.
 ///
-/// The queue is created once per execution and shared by all reorderable
-/// sibling streams for that execution. Whichever stream becomes idle first may
-/// take the next unopened file from the front of the queue.
+/// The source is created once per execution and shared by all reorderable
+/// sibling streams for that execution. It starts empty: each stream registers
+/// its own partition's files when it is opened. Whichever stream becomes idle
+/// first may then take the next unopened file from any partition that has been
+/// registered in this process.
 ///
-/// It uses a [`Mutex`] internally to provide thread-safe access
-/// to the shared file queue.
+/// Registering files lazily keeps distributed execution correct. If a process
+/// opens only partition `p`, only that partition's files are available to the
+/// shared queue; if it opens all partitions, all files are available for local
+/// work stealing.
 #[derive(Debug, Clone)]
 pub(crate) struct SharedWorkSource {
     inner: Arc<SharedWorkSourceInner>,
 }
 
 #[derive(Debug, Default)]
 pub(super) struct SharedWorkSourceInner {
-    files: Mutex<VecDeque<PartitionedFile>>,
+    state: Mutex<SharedWorkSourceState>,
+}
+
+#[derive(Debug, Default)]
+struct SharedWorkSourceState {
+    files: VecDeque<PartitionedFile>,
+    registered_partitions: HashSet<usize>,
 }
 
 impl SharedWorkSource {
-    /// Create a shared work source containing the provided unopened files.
-    pub(crate) fn new(files: impl IntoIterator<Item = PartitionedFile>) -> Self {
-        let files = files.into_iter().collect();
+    /// Create an empty shared work source.
+    pub(crate) fn empty() -> Self {
         Self {
-            inner: Arc::new(SharedWorkSourceInner {
-                files: Mutex::new(files),
-            }),
+            inner: Arc::new(SharedWorkSourceInner::default()),
         }
     }
 
-    /// Create a shared work source for the unopened files in `config`.
+    /// Register the files for `partition` in this process-local work source.
     ///
-    /// Files are reordered by the file source (e.g. by statistics for TopK)
-    /// before being placed in the shared queue, so the most promising files
-    /// are processed first across all partitions.
-    pub(crate) fn from_config(config: &FileScanConfig) -> Self {
-        let files: Vec<_> = config
-            .file_groups
-            .iter()
-            .flat_map(FileGroup::iter)
-            .cloned()
-            .collect();
-        let files = config.file_source.reorder_files(files);
-        Self::new(files)
+    /// A partition is registered at most once. Files are reordered by the file
+    /// source (e.g. by statistics for TopK) together with any already queued
+    /// files before becoming available for sibling streams to steal.
+    pub(crate) fn register_partition(&self, partition: usize, config: &FileScanConfig) {
+        let Some(file_group) = config.file_groups.get(partition) else {
+            return;
+        };
+
+        let mut state = self.inner.state.lock();
+        if !state.registered_partitions.insert(partition) || file_group.is_empty() {
+            return;
+        }
+
+        state.files.extend(file_group.iter().cloned());
+        let files = state.files.drain(..).collect();
+        state.files = config.file_source.reorder_files(files).into();
     }
 
     /// Pop the next file from the shared work queue.
     ///
-    /// Returns `None` if the queue is empty
+    /// Returns `None` if the queue is empty.
     fn pop_front(&self) -> Option<PartitionedFile> {
-        self.inner.files.lock().pop_front()
+        self.inner.state.lock().files.pop_front()
     }
 }