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,8 @@ 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.
+    pub(crate) shared_work_source: Arc<OnceLock<SharedWorkSource>>,
 }
 
 /// A builder for [`FileScanConfig`]'s.
@@ -551,6 +560,7 @@ impl FileScanConfigBuilder {
             expr_adapter_factory: expr_adapter,
             statistics,
             partitioned_by_file_group,
+            shared_work_source: Arc::new(OnceLock::new()),
         }
     }
 }

datafusion/datasource/src/file_scan_config.rs

@@ -19,31 +19,57 @@ 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.
+enum Reorderable {
+    /// This stream may reorder work, optionally using a shared file queue.
+    Yes {
+        shared_work_source: Option<SharedWorkSource>,
+    },
+    /// This stream may not reorder work.
+    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.
     pub fn new(config: &'a FileScanConfig) -> Self {
+        let reorderable = if config.preserve_order || config.partitioned_by_file_group {
+            Reorderable::No
+        } else {
+            Reorderable::Yes {
+                shared_work_source: Some(
+                    config
+                        .shared_work_source
+                        .get_or_init(SharedWorkSource::new)
+                        .clone(),
+                ),
+            }
+        };
+
         Self {
             config,
             partition: None,
             morselizer: None,
             metrics: None,
             on_error: OnError::Fail,
+            reorderable,
         }
     }
 
@@ -89,6 +115,7 @@ impl<'a> FileStreamBuilder<'a> {
             morselizer,
             metrics,
             on_error,
+            reorderable,
         } = self;
 
         let Some(partition) = partition else {
@@ -106,10 +133,23 @@ impl<'a> FileStreamBuilder<'a> {
                 "FileStreamBuilder invalid partition index: {partition}"
             );
         };
+        let files = file_group.into_inner();
+        let work_source = match reorderable {
+            Reorderable::Yes { shared_work_source } => {
+                if let Some(shared) = shared_work_source {
+                    shared.register_stream();
+                    shared.push_files(files);
+                    WorkSource::Shared(shared)
+                } else {
+                    WorkSource::Local(files.into())
+                }
+            }
+            Reorderable::No => WorkSource::Local(files.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,