From 1151bb4f2ef5271c9241e94a9a13e27a59dfa7a8 Mon Sep 17 00:00:00 2001 From: Phillip LeBlanc <879445+phillipleblanc@users.noreply.github.com> Date: Fri, 26 Jun 2026 13:51:21 +0900 Subject: [PATCH 1/4] feat(datasource): add enable_file_scan_work_stealing option to gate cross-partition scan work-stealing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DataSourceExec shares one SharedWorkSource queue across all of a scan's partition streams via a per-instance OnceLock, so sibling streams in the same process steal work and each file is opened exactly once. That is correct only when all output partitions execute in the same process. Distributed engines (e.g. Ballista) run each partition in a separate process, where every process rebuilds the full queue over all file_groups and reads the entire input — N-fold scan amplification and, with a Hash-shuffled partial aggregate above the scan, N-fold inflated aggregate results. Add `execution.enable_file_scan_work_stealing` (default true, so single-node behavior is unchanged) and gate `create_sibling_state` in `DataSourceExec::execute` on it. When false, each partition reads only its own file group (`WorkSource::Local`) — the correct behavior under distributed execution. Claude-Session: https://claude.ai/code/session_01G8drJk8Fx5ynRQZ4WkuVju --- datafusion/common/src/config.rs | 16 ++++++++++++++++ datafusion/datasource/src/source.rs | 21 +++++++++++++++++---- 2 files changed, 33 insertions(+), 4 deletions(-) diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 50a8f71154889..fc309322ea46f 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -636,6 +636,22 @@ config_namespace! { /// Currently experimental pub split_file_groups_by_statistics: bool, default = false + /// Whether sibling partition streams of a single file scan may share a + /// common queue of unopened files and steal work from one another + /// (dynamic work scheduling). + /// + /// This is safe and beneficial when all output partitions of a scan are + /// executed by sibling streams within the same process, since the shared + /// queue guarantees every file is opened exactly once across the siblings. + /// + /// It MUST be disabled when partitions are executed in isolation from one + /// another — for example under a distributed engine (Ballista) where each + /// partition runs in a separate process. In that case each process would + /// build its own shared queue over *all* files and read the entire input, + /// causing N-fold scan amplification and N-fold inflated aggregate results. + /// When disabled, each partition reads only its own file group. + pub enable_file_scan_work_stealing: bool, default = true + /// Should DataFusion keep the columns used for partition_by in the output RecordBatches pub keep_partition_by_columns: bool, default = false diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs index af4bc09504937..9b434dcc66bea 100644 --- a/datafusion/datasource/src/source.rs +++ b/datafusion/datasource/src/source.rs @@ -389,10 +389,23 @@ impl ExecutionPlan for DataSourceExec { partition: usize, context: Arc, ) -> Result { - let shared_state = self - .execution_state - .get_or_init(|| self.data_source.create_sibling_state()) - .clone(); + // Sibling work-stealing shares one queue of unopened files across all + // partition streams of this scan. That is only correct when those streams + // run in the same process. Under a distributed engine each partition runs + // in isolation, so sharing must be disabled (each partition then reads + // only its own file group) to avoid reading the whole input per partition. + let shared_state = if context + .session_config() + .options() + .execution + .enable_file_scan_work_stealing + { + self.execution_state + .get_or_init(|| self.data_source.create_sibling_state()) + .clone() + } else { + None + }; let args = OpenArgs::new(partition, Arc::clone(&context)) .with_shared_state(shared_state); let stream = self.data_source.open_with_args(args)?; From 3e97c01e2856d7397c4600654472a40d6bc4fe4d Mon Sep 17 00:00:00 2001 From: Phillip LeBlanc <879445+phillipleblanc@users.noreply.github.com> Date: Sat, 27 Jun 2026 07:34:05 +0900 Subject: [PATCH 2/4] address review: clarify work-stealing gate scope + add regression test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Copilot review follow-ups on the `enable_file_scan_work_stealing` gate in `DataSourceExec::execute`: - Document that `create_sibling_state` is the work-stealing hook and is currently only implemented by file-scan sources (`FileScanConfig`); every other `DataSource` inherits the default `None`, so the gate is effectively file-scan scoped and the option name is accurate. Note that a future `DataSource` introducing its own sibling state would also be governed by the flag and should revisit the intended behavior. - Add a regression test that drives `DataSourceExec::execute` with a probe `DataSource` and asserts the gate flips sibling state: shared when `enable_file_scan_work_stealing=true`, absent when `false` (so each partition opens in isolation and reads only its own file group — no cross-partition work-stealing under distributed execution). Claude-Session: https://claude.ai/code/session_01G8drJk8Fx5ynRQZ4WkuVju --- datafusion/datasource/src/source.rs | 131 ++++++++++++++++++++++++++++ 1 file changed, 131 insertions(+) diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs index 9b434dcc66bea..21ddf0f74be56 100644 --- a/datafusion/datasource/src/source.rs +++ b/datafusion/datasource/src/source.rs @@ -394,6 +394,14 @@ impl ExecutionPlan for DataSourceExec { // run in the same process. Under a distributed engine each partition runs // in isolation, so sharing must be disabled (each partition then reads // only its own file group) to avoid reading the whole input per partition. + // + // `create_sibling_state` is the work-stealing hook and is currently only + // implemented by file-scan sources (`FileScanConfig`); every other + // `DataSource` inherits the default `None`, so this gate is effectively + // file-scan scoped and the `enable_file_scan_work_stealing` name is + // accurate. A future `DataSource` that introduces its own sibling state + // would also be governed by this flag and should revisit whether that is + // the intended behavior. let shared_state = if context .session_config() .options() @@ -630,3 +638,126 @@ where Self::new(Arc::new(source)) } } + +#[cfg(test)] +mod work_stealing_gate_tests { + use super::*; + use std::any::Any; + use std::sync::atomic::{AtomicBool, Ordering}; + + use arrow::datatypes::{Schema, SchemaRef}; + use datafusion_execution::config::SessionConfig; + use datafusion_physical_plan::stream::EmptyRecordBatchStream; + + /// A minimal [`DataSource`] that always offers sibling (work-stealing) state and + /// records whether that state was actually handed to it through + /// [`DataSource::open_with_args`]. This lets the test observe how + /// [`DataSourceExec::execute`] gates work-stealing on the + /// `enable_file_scan_work_stealing` option without depending on a real file + /// scan. + #[derive(Debug)] + struct GateProbeSource { + schema: SchemaRef, + saw_sibling_state: Arc, + } + + impl DataSource for GateProbeSource { + fn open( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + Ok(Box::pin(EmptyRecordBatchStream::new(Arc::clone( + &self.schema, + )))) + } + + fn open_with_args(&self, args: OpenArgs) -> Result { + self.saw_sibling_state + .store(args.sibling_state.is_some(), Ordering::SeqCst); + self.open(args.partition, args.context) + } + + /// Always offer sibling state so the gate has an observable effect. + fn create_sibling_state(&self) -> Option> { + Some(Arc::new(()) as Arc) + } + + fn fmt_as(&self, _t: DisplayFormatType, _f: &mut Formatter) -> fmt::Result { + Ok(()) + } + + fn output_partitioning(&self) -> Partitioning { + Partitioning::UnknownPartitioning(1) + } + + fn eq_properties(&self) -> EquivalenceProperties { + EquivalenceProperties::new(Arc::clone(&self.schema)) + } + + fn partition_statistics( + &self, + _partition: Option, + ) -> Result> { + Ok(Arc::new(Statistics::new_unknown(self.schema.as_ref()))) + } + + fn with_fetch(&self, _limit: Option) -> Option> { + None + } + + fn fetch(&self) -> Option { + None + } + + fn try_swapping_with_projection( + &self, + _projection: &ProjectionExprs, + ) -> Result>> { + Ok(None) + } + } + + /// Execute partition 0 of a `DataSourceExec` wrapping a [`GateProbeSource`] with + /// the given `enable_file_scan_work_stealing` setting; returns whether the + /// source was handed sibling (work-stealing) state. + fn executed_with_sibling_state(enable_work_stealing: bool) -> Result { + let schema: SchemaRef = Arc::new(Schema::empty()); + let saw_sibling_state = Arc::new(AtomicBool::new(false)); + let source = Arc::new(GateProbeSource { + schema, + saw_sibling_state: Arc::clone(&saw_sibling_state), + }); + let exec = DataSourceExec::new(source); + + let mut config = SessionConfig::new(); + config + .options_mut() + .execution + .enable_file_scan_work_stealing = enable_work_stealing; + let ctx = Arc::new(TaskContext::default().with_session_config(config)); + + // execute() calls `open_with_args` synchronously before returning the + // stream, so the probe flag is set by the time this returns. + let _stream = exec.execute(0, ctx)?; + Ok(saw_sibling_state.load(Ordering::SeqCst)) + } + + #[test] + fn work_stealing_gate_controls_sibling_state() -> Result<()> { + // Single-process default: work-stealing on, sibling state is shared so + // partitions cooperatively drain one file queue. + assert!( + executed_with_sibling_state(true)?, + "enable_file_scan_work_stealing=true should pass shared sibling state to the source", + ); + // Distributed setting: work-stealing off, the source opens each partition in + // isolation (no shared file queue) so it reads only its own file group and + // does not re-read files belonging to other partitions. + assert!( + !executed_with_sibling_state(false)?, + "enable_file_scan_work_stealing=false must NOT pass sibling state (no cross-partition work-stealing)", + ); + Ok(()) + } +} From 12ecaf0bfffc65d7eca9a7e2db5650eaa8dd29e2 Mon Sep 17 00:00:00 2001 From: Phillip LeBlanc <879445+phillipleblanc@users.noreply.github.com> Date: Mon, 29 Jun 2026 08:32:27 +0900 Subject: [PATCH 3/4] style(datasource): match surrounding comment style for work-stealing gate Claude-Session: https://claude.ai/code/session_01G8drJk8Fx5ynRQZ4WkuVju --- datafusion/common/src/config.rs | 18 ++++---------- datafusion/datasource/src/source.rs | 37 ++++++++--------------------- 2 files changed, 14 insertions(+), 41 deletions(-) diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index fc309322ea46f..e03b0e6353658 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -636,20 +636,10 @@ config_namespace! { /// Currently experimental pub split_file_groups_by_statistics: bool, default = false - /// Whether sibling partition streams of a single file scan may share a - /// common queue of unopened files and steal work from one another - /// (dynamic work scheduling). - /// - /// This is safe and beneficial when all output partitions of a scan are - /// executed by sibling streams within the same process, since the shared - /// queue guarantees every file is opened exactly once across the siblings. - /// - /// It MUST be disabled when partitions are executed in isolation from one - /// another — for example under a distributed engine (Ballista) where each - /// partition runs in a separate process. In that case each process would - /// build its own shared queue over *all* files and read the entire input, - /// causing N-fold scan amplification and N-fold inflated aggregate results. - /// When disabled, each partition reads only its own file group. + /// Whether sibling partition streams of a file scan share a queue of + /// unopened files and steal work from one another. Must be disabled when + /// partitions run in separate processes (e.g. distributed execution), + /// where each would otherwise scan the whole input, not just its group. pub enable_file_scan_work_stealing: bool, default = true /// Should DataFusion keep the columns used for partition_by in the output RecordBatches diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs index 21ddf0f74be56..51ca582a59c3f 100644 --- a/datafusion/datasource/src/source.rs +++ b/datafusion/datasource/src/source.rs @@ -389,19 +389,10 @@ impl ExecutionPlan for DataSourceExec { partition: usize, context: Arc, ) -> Result { - // Sibling work-stealing shares one queue of unopened files across all - // partition streams of this scan. That is only correct when those streams - // run in the same process. Under a distributed engine each partition runs - // in isolation, so sharing must be disabled (each partition then reads - // only its own file group) to avoid reading the whole input per partition. - // - // `create_sibling_state` is the work-stealing hook and is currently only - // implemented by file-scan sources (`FileScanConfig`); every other - // `DataSource` inherits the default `None`, so this gate is effectively - // file-scan scoped and the `enable_file_scan_work_stealing` name is - // accurate. A future `DataSource` that introduces its own sibling state - // would also be governed by this flag and should revisit whether that is - // the intended behavior. + // Sibling work-stealing shares one file queue across a scan's partition + // streams, which is only correct when they run in the same process. + // Disable it for isolated partitions (e.g. distributed execution) so each + // reads only its own file group. let shared_state = if context .session_config() .options() @@ -649,12 +640,8 @@ mod work_stealing_gate_tests { use datafusion_execution::config::SessionConfig; use datafusion_physical_plan::stream::EmptyRecordBatchStream; - /// A minimal [`DataSource`] that always offers sibling (work-stealing) state and - /// records whether that state was actually handed to it through - /// [`DataSource::open_with_args`]. This lets the test observe how - /// [`DataSourceExec::execute`] gates work-stealing on the - /// `enable_file_scan_work_stealing` option without depending on a real file - /// scan. + /// A [`DataSource`] that always offers sibling state and records whether + /// [`DataSourceExec::execute`] actually handed it that state. #[derive(Debug)] struct GateProbeSource { schema: SchemaRef, @@ -718,9 +705,8 @@ mod work_stealing_gate_tests { } } - /// Execute partition 0 of a `DataSourceExec` wrapping a [`GateProbeSource`] with - /// the given `enable_file_scan_work_stealing` setting; returns whether the - /// source was handed sibling (work-stealing) state. + /// Execute partition 0 with the given `enable_file_scan_work_stealing` value + /// and return whether the source was handed sibling state. fn executed_with_sibling_state(enable_work_stealing: bool) -> Result { let schema: SchemaRef = Arc::new(Schema::empty()); let saw_sibling_state = Arc::new(AtomicBool::new(false)); @@ -745,15 +731,12 @@ mod work_stealing_gate_tests { #[test] fn work_stealing_gate_controls_sibling_state() -> Result<()> { - // Single-process default: work-stealing on, sibling state is shared so - // partitions cooperatively drain one file queue. + // Default (single-process): sibling state is shared for work-stealing. assert!( executed_with_sibling_state(true)?, "enable_file_scan_work_stealing=true should pass shared sibling state to the source", ); - // Distributed setting: work-stealing off, the source opens each partition in - // isolation (no shared file queue) so it reads only its own file group and - // does not re-read files belonging to other partitions. + // Disabled (distributed): no sibling state, so no cross-partition stealing. assert!( !executed_with_sibling_state(false)?, "enable_file_scan_work_stealing=false must NOT pass sibling state (no cross-partition work-stealing)", From 8a59727939bc4fa3fe94800957cf7c6f54900c35 Mon Sep 17 00:00:00 2001 From: Phillip LeBlanc <879445+phillipleblanc@users.noreply.github.com> Date: Mon, 29 Jun 2026 09:25:56 +0900 Subject: [PATCH 4/4] fix(datasource): scope file scan work stealing to opened partitions Create SharedWorkSource empty and register each partition's file group only when that partition is opened. This preserves local sibling work stealing while preventing distributed workers that execute a single partition from scanning every file group. Remove the execution.enable_file_scan_work_stealing config gate and add/adjust morsel tests for single-partition registration and late partition registration. --- datafusion/common/src/config.rs | 6 - .../datasource/src/file_scan_config/mod.rs | 8 +- .../datasource/src/file_stream/builder.rs | 9 +- datafusion/datasource/src/file_stream/mod.rs | 79 +++++----- .../datasource/src/file_stream/work_source.rs | 70 +++++---- datafusion/datasource/src/source.rs | 135 +----------------- 6 files changed, 101 insertions(+), 206 deletions(-) diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index e03b0e6353658..50a8f71154889 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -636,12 +636,6 @@ config_namespace! { /// Currently experimental pub split_file_groups_by_statistics: bool, default = false - /// Whether sibling partition streams of a file scan share a queue of - /// unopened files and steal work from one another. Must be disabled when - /// partitions run in separate processes (e.g. distributed execution), - /// where each would otherwise scan the whole input, not just its group. - pub enable_file_scan_work_stealing: bool, default = true - /// Should DataFusion keep the columns used for partition_by in the output RecordBatches pub keep_partition_by_columns: bool, default = false 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() } } diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs index 51ca582a59c3f..af4bc09504937 100644 --- a/datafusion/datasource/src/source.rs +++ b/datafusion/datasource/src/source.rs @@ -389,22 +389,10 @@ impl ExecutionPlan for DataSourceExec { partition: usize, context: Arc, ) -> Result { - // Sibling work-stealing shares one file queue across a scan's partition - // streams, which is only correct when they run in the same process. - // Disable it for isolated partitions (e.g. distributed execution) so each - // reads only its own file group. - let shared_state = if context - .session_config() - .options() - .execution - .enable_file_scan_work_stealing - { - self.execution_state - .get_or_init(|| self.data_source.create_sibling_state()) - .clone() - } else { - None - }; + let shared_state = self + .execution_state + .get_or_init(|| self.data_source.create_sibling_state()) + .clone(); let args = OpenArgs::new(partition, Arc::clone(&context)) .with_shared_state(shared_state); let stream = self.data_source.open_with_args(args)?; @@ -629,118 +617,3 @@ where Self::new(Arc::new(source)) } } - -#[cfg(test)] -mod work_stealing_gate_tests { - use super::*; - use std::any::Any; - use std::sync::atomic::{AtomicBool, Ordering}; - - use arrow::datatypes::{Schema, SchemaRef}; - use datafusion_execution::config::SessionConfig; - use datafusion_physical_plan::stream::EmptyRecordBatchStream; - - /// A [`DataSource`] that always offers sibling state and records whether - /// [`DataSourceExec::execute`] actually handed it that state. - #[derive(Debug)] - struct GateProbeSource { - schema: SchemaRef, - saw_sibling_state: Arc, - } - - impl DataSource for GateProbeSource { - fn open( - &self, - _partition: usize, - _context: Arc, - ) -> Result { - Ok(Box::pin(EmptyRecordBatchStream::new(Arc::clone( - &self.schema, - )))) - } - - fn open_with_args(&self, args: OpenArgs) -> Result { - self.saw_sibling_state - .store(args.sibling_state.is_some(), Ordering::SeqCst); - self.open(args.partition, args.context) - } - - /// Always offer sibling state so the gate has an observable effect. - fn create_sibling_state(&self) -> Option> { - Some(Arc::new(()) as Arc) - } - - fn fmt_as(&self, _t: DisplayFormatType, _f: &mut Formatter) -> fmt::Result { - Ok(()) - } - - fn output_partitioning(&self) -> Partitioning { - Partitioning::UnknownPartitioning(1) - } - - fn eq_properties(&self) -> EquivalenceProperties { - EquivalenceProperties::new(Arc::clone(&self.schema)) - } - - fn partition_statistics( - &self, - _partition: Option, - ) -> Result> { - Ok(Arc::new(Statistics::new_unknown(self.schema.as_ref()))) - } - - fn with_fetch(&self, _limit: Option) -> Option> { - None - } - - fn fetch(&self) -> Option { - None - } - - fn try_swapping_with_projection( - &self, - _projection: &ProjectionExprs, - ) -> Result>> { - Ok(None) - } - } - - /// Execute partition 0 with the given `enable_file_scan_work_stealing` value - /// and return whether the source was handed sibling state. - fn executed_with_sibling_state(enable_work_stealing: bool) -> Result { - let schema: SchemaRef = Arc::new(Schema::empty()); - let saw_sibling_state = Arc::new(AtomicBool::new(false)); - let source = Arc::new(GateProbeSource { - schema, - saw_sibling_state: Arc::clone(&saw_sibling_state), - }); - let exec = DataSourceExec::new(source); - - let mut config = SessionConfig::new(); - config - .options_mut() - .execution - .enable_file_scan_work_stealing = enable_work_stealing; - let ctx = Arc::new(TaskContext::default().with_session_config(config)); - - // execute() calls `open_with_args` synchronously before returning the - // stream, so the probe flag is set by the time this returns. - let _stream = exec.execute(0, ctx)?; - Ok(saw_sibling_state.load(Ordering::SeqCst)) - } - - #[test] - fn work_stealing_gate_controls_sibling_state() -> Result<()> { - // Default (single-process): sibling state is shared for work-stealing. - assert!( - executed_with_sibling_state(true)?, - "enable_file_scan_work_stealing=true should pass shared sibling state to the source", - ); - // Disabled (distributed): no sibling state, so no cross-partition stealing. - assert!( - !executed_with_sibling_state(false)?, - "enable_file_scan_work_stealing=false must NOT pass sibling state (no cross-partition work-stealing)", - ); - Ok(()) - } -}