diff --git a/datafusion/datasource/src/file_scan_config/mod.rs b/datafusion/datasource/src/file_scan_config/mod.rs index b5915d6ebd11c..52d19828935af 100644 --- a/datafusion/datasource/src/file_scan_config/mod.rs +++ b/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> { if self.preserve_order || self.partitioned_by_file_group { return None; } - Some(Arc::new(SharedWorkSource::from_config(self)) as Arc) + Some(Arc::new(SharedWorkSource::empty()) as Arc) } } diff --git a/datafusion/datasource/src/file_stream/builder.rs b/datafusion/datasource/src/file_stream/builder.rs index 7034e902550a9..bfb12a49077ee 100644 --- a/datafusion/datasource/src/file_stream/builder.rs +++ b/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); diff --git a/datafusion/datasource/src/file_stream/mod.rs b/datafusion/datasource/src/file_stream/mod.rs index e277690cff810..ffd9905e9ede4 100644 --- a/datafusion/datasource/src/file_stream/mod.rs +++ b/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::().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::>(); - 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::().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()) diff --git a/datafusion/datasource/src/file_stream/work_source.rs b/datafusion/datasource/src/file_stream/work_source.rs index c00048453b304..4b00058ff2179 100644 --- a/datafusion/datasource/src/file_stream/work_source.rs +++ b/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,14 +54,18 @@ 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, @@ -70,40 +73,47 @@ pub(crate) struct SharedWorkSource { #[derive(Debug, Default)] pub(super) struct SharedWorkSourceInner { - files: Mutex>, + state: Mutex, +} + +#[derive(Debug, Default)] +struct SharedWorkSourceState { + files: VecDeque, + registered_partitions: HashSet, } impl SharedWorkSource { - /// Create a shared work source containing the provided unopened files. - pub(crate) fn new(files: impl IntoIterator) -> 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 { - self.inner.files.lock().pop_front() + self.inner.state.lock().files.pop_front() } }