Add Dynamic work scheduling in FileStream

Author: alamb

Cargo.lock

@@ -64,6 +64,7 @@ itertools = { workspace = true }
 liblzma = { workspace = true, optional = true }
 log = { workspace = true }
 object_store = { workspace = true }
+parking_lot = { workspace = true }
 rand = { workspace = true }
 tempfile = { workspace = true, optional = true }
 tokio = { workspace = true }

datafusion/datasource/Cargo.toml

@@ -22,7 +22,8 @@ use crate::file_groups::FileGroup;
 use crate::{
     PartitionedFile, display::FileGroupsDisplay, file::FileSource,
     file_compression_type::FileCompressionType, file_stream::FileStreamBuilder,
-    source::DataSource, statistics::MinMaxStatistics,
+    file_stream::work_source::SharedWorkSource, source::DataSource,
+    statistics::MinMaxStatistics,
 };
 use arrow::datatypes::FieldRef;
 use arrow::datatypes::{DataType, Schema, SchemaRef};
@@ -55,7 +56,13 @@ use datafusion_physical_plan::{
     metrics::ExecutionPlanMetricsSet,
 };
 use log::{debug, warn};
-use std::{any::Any, fmt::Debug, fmt::Formatter, fmt::Result as FmtResult, sync::Arc};
+use std::{
+    any::Any,
+    fmt::Debug,
+    fmt::Formatter,
+    fmt::Result as FmtResult,
+    sync::{Arc, OnceLock},
+};
 
 /// [`FileScanConfig`] represents scanning data from a group of files
 ///
@@ -209,6 +216,11 @@ pub struct FileScanConfig {
     /// If the number of file partitions > target_partitions, the file partitions will be grouped
     /// in a round-robin fashion such that number of file partitions = target_partitions.
     pub partitioned_by_file_group: bool,
+    /// Shared queue of unopened files for sibling streams in this scan.
+    ///
+    /// This is initialized once per `FileScanConfig` and reused by reorderable
+    /// `FileStream`s created from that config.
+    pub(crate) shared_work_source: Arc<OnceLock<SharedWorkSource>>,
 }
 
 /// A builder for [`FileScanConfig`]'s.
@@ -551,10 +563,34 @@ impl FileScanConfigBuilder {
             expr_adapter_factory: expr_adapter,
             statistics,
             partitioned_by_file_group,
+            shared_work_source: Arc::new(OnceLock::new()),
         }
     }
 }
 
+impl FileScanConfig {
+    /// Returns the shared unopened-file queue for reorderable streams in this scan.
+    ///
+    /// The queue is initialized once from all file groups so sibling streams
+    /// can begin stealing work immediately, even if they are built or polled
+    /// before every sibling `FileStream` has been constructed.
+    pub(crate) fn shared_work_source(&self) -> Option<SharedWorkSource> {
+        if self.preserve_order || self.partitioned_by_file_group {
+            return None;
+        }
+
+        Some(
+            self.shared_work_source
+                .get_or_init(|| {
+                    SharedWorkSource::new(
+                        self.file_groups.iter().flat_map(FileGroup::iter).cloned(),
+                    )
+                })
+                .clone(),
+        )
+    }
+}
+
 impl From<FileScanConfig> for FileScanConfigBuilder {
     fn from(config: FileScanConfig) -> Self {
         Self {

datafusion/datasource/src/file_scan_config.rs

@@ -19,31 +19,50 @@ use std::sync::Arc;
 
 use crate::file_scan_config::FileScanConfig;
 use crate::file_stream::scan_state::ScanState;
+use crate::file_stream::work_source::{SharedWorkSource, WorkSource};
 use crate::morsel::{FileOpenerMorselizer, Morselizer};
 use datafusion_common::{Result, internal_err};
 use datafusion_physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet};
 
 use super::metrics::FileStreamMetrics;
 use super::{FileOpener, FileStream, FileStreamState, OnError};
 
+/// Whether this stream may reorder work across sibling `FileStream`s.
+///
+/// This is derived entirely from [`FileScanConfig`]. Streams that must
+/// preserve file order or file-group partition boundaries are not reorderable.
+enum Reorderable {
+    /// This stream may reorder work using a shared queue of unopened files.
+    Yes(SharedWorkSource),
+    /// This stream must keep its own local file order.
+    No,
+}
+
 /// Builder for constructing a [`FileStream`].
 pub struct FileStreamBuilder<'a> {
     config: &'a FileScanConfig,
     partition: Option<usize>,
     morselizer: Option<Box<dyn Morselizer>>,
     metrics: Option<&'a ExecutionPlanMetricsSet>,
     on_error: OnError,
+    reorderable: Reorderable,
 }
 
 impl<'a> FileStreamBuilder<'a> {
-    /// Create a new builder.
+    /// Create a new builder for [`FileStream`].
     pub fn new(config: &'a FileScanConfig) -> Self {
+        let reorderable = match config.shared_work_source() {
+            Some(shared_work_source) => Reorderable::Yes(shared_work_source),
+            None => Reorderable::No,
+        };
+
         Self {
             config,
             partition: None,
             morselizer: None,
             metrics: None,
             on_error: OnError::Fail,
+            reorderable,
         }
     }
 
@@ -89,6 +108,7 @@ impl<'a> FileStreamBuilder<'a> {
             morselizer,
             metrics,
             on_error,
+            reorderable,
         } = self;
 
         let Some(partition) = partition else {
@@ -106,10 +126,14 @@ impl<'a> FileStreamBuilder<'a> {
                 "FileStreamBuilder invalid partition index: {partition}"
             );
         };
+        let work_source = match reorderable {
+            Reorderable::Yes(shared) => WorkSource::Shared(shared),
+            Reorderable::No => WorkSource::Local(file_group.into_inner().into()),
+        };
 
         let file_stream_metrics = FileStreamMetrics::new(metrics, partition);
         let scan_state = Box::new(ScanState::new(
-            file_group.into_inner(),
+            work_source,
             config.limit,
             morselizer,
             on_error,