From 17ad806a5081632b4952d737bda23ade005215dd Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 2 Jul 2026 07:57:01 -0600 Subject: [PATCH 1/4] test: reproduce isolated-per-task FileStream work-stealing scan bug Add an ignored regression test for #23293. FileStream sibling work-stealing seeds one shared work queue from every file group and relies on all output partitions being polled concurrently in one process. An executor that runs each partition as an isolated task polls only one partition, which then drains the whole queue and reads files belonging to other partitions, inflating scan output by the partition count. The test builds and drives only partition 0 and asserts it reads solely its own file. It fails on main by design, so it is #[ignore]d with its assertion intact as a caught regression to triage. (cherry picked from commit 1b2760bfb396a7cfa9d963608e8b72504c700420) --- datafusion/datasource/src/file_stream/mod.rs | 71 ++++++++++++++++++++ 1 file changed, 71 insertions(+) diff --git a/datafusion/datasource/src/file_stream/mod.rs b/datafusion/datasource/src/file_stream/mod.rs index e277690cff810..fd6fc4a3c912e 100644 --- a/datafusion/datasource/src/file_stream/mod.rs +++ b/datafusion/datasource/src/file_stream/mod.rs @@ -1131,6 +1131,77 @@ mod tests { Ok(()) } + /// Reproduces . + /// + /// Executors that run each output partition as an isolated task in a + /// separate process (Ballista, datafusion-distributed) build one plan + /// instance and call `execute(partition)` for a single partition. The + /// sibling partitions are never polled in that process. A `DataSourceExec` + /// seeds one shared work queue from every file group, so the single polled + /// partition drains the whole queue and reads files that belong to other + /// partitions. Each isolated task does the same, inflating the scan output + /// by the partition count. + /// + /// This test models that setup: it builds and drives only partition 0. + /// Partition 0 should read only its own file (`file1.parquet`), but today + /// it also reads partition 1's file, so it fails with + /// `"Batch: 101\nBatch: 201"`. + /// + /// Ignored because it fails on `main` by design: it is a caught regression + /// kept with its assertion intact. Run it with + /// `cargo test -p datafusion-datasource -- --ignored + /// regression_isolated_partition_reads_only_its_own_files`. + #[tokio::test] + #[ignore = "reproduces #23293: shared FileStream work queue is drained by a single isolated task"] + async fn regression_isolated_partition_reads_only_its_own_files() -> Result<()> { + let test = FileStreamMorselTest::new() + .with_file_in_partition( + PartitionId(0), + MockPlanner::builder("file1.parquet") + .add_plan(MockPlanBuilder::new().with_morsel(MorselId(10), 101)) + .return_none(), + ) + .with_file_in_partition( + PartitionId(1), + MockPlanner::builder("file2.parquet") + .add_plan(MockPlanBuilder::new().with_morsel(MorselId(11), 201)) + .return_none(), + ); + + let config = test.test_config(); + let metrics_set = ExecutionPlanMetricsSet::new(); + + // A `DataSourceExec` builds exactly one shared work source per + // execution, seeded from every file group. An isolated task builds that + // state and then executes only its own partition. + let shared_work_source = config + .create_sibling_state() + .and_then(|state| state.as_ref().downcast_ref::().cloned()); + assert!( + shared_work_source.is_some(), + "a plain repartitioned scan should use a shared work source" + ); + + // Build and drive ONLY partition 0. Partition 1's stream is never + // created, exactly as in an isolated per-task executor. + let partition0 = FileStreamBuilder::new(&config) + .with_partition(0) + .with_shared_work_source(shared_work_source) + .with_morselizer(Box::new(test.morselizer.clone())) + .with_metrics(&metrics_set) + .build()?; + + let output = drain_stream_output(partition0).await?; + + assert_eq!( + output, "Batch: 101", + "partition 0 read files belonging to a sibling partition: a single \ + isolated task drained the shared work queue (issue #23293)" + ); + + Ok(()) + } + /// Verifies that an empty sibling can immediately steal shared files when /// it is polled before the stream that originally owned them. #[tokio::test] From a4678405c8580176d3da9adfb61c65f66a2054a8 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 2 Jul 2026 08:31:48 -0600 Subject: [PATCH 2/4] feat: add datafusion.execution.enable_file_stream_work_stealing config FileStream sibling work-stealing (WorkSource::Shared) seeds one shared work queue from every file group and lets whichever output partition goes idle first steal the next file or byte-range morsel. This assumes all output partitions of a scan are polled concurrently in one process. Executors that run each output partition as an isolated task in a separate process (Ballista, datafusion-distributed) never poll the sibling partitions, so the single polled partition drains the whole queue and reads files belonging to other partitions, inflating the scan output by the partition count. That is a correctness bug for those executors, and the existing escape hatches (preserve_order, partitioned_by_file_group) are plan-level flags, not a session config they can set centrally. Add datafusion.execution.enable_file_stream_work_stealing (default true), checked in FileScanConfig::create_sibling_state: when false it returns None so each partition falls back to reading only its own file group. This mirrors the enable_dynamic_filter_pushdown precedent and round-trips through datafusion-proto as a config value. Thread ConfigOptions into DataSource::create_sibling_state so the flag is read from the session config at execute time. Turn the #[ignore]'d reproduction test into a passing regression test that drives only partition 0 and checks both behaviors: with the default (stealing on) partition 0 also reads partition 1's file, and with the flag off it reads only its own. Closes #23293. (cherry picked from commit 482b04ba5589f39a19b2c684306ae02be6f72e1f) --- datafusion/common/src/config.rs | 15 ++ .../datasource/src/file_scan_config/mod.rs | 16 ++- datafusion/datasource/src/file_stream/mod.rs | 133 +++++++++++------- datafusion/datasource/src/source.rs | 10 +- .../test_files/information_schema.slt | 2 + docs/source/user-guide/configs.md | 1 + 6 files changed, 117 insertions(+), 60 deletions(-) diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 9096e5c366dac..3b2cde43f695b 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -639,6 +639,21 @@ config_namespace! { /// Should DataFusion keep the columns used for partition_by in the output RecordBatches pub keep_partition_by_columns: bool, default = false + /// When `true` (the default), sibling partition streams of a single file + /// scan share a work queue: whichever output partition goes idle first + /// steals the next unopened file (or byte-range morsel) from the shared + /// queue. This balances work when all output partitions are polled + /// concurrently in one process. + /// + /// Executors that run each output partition as an isolated task in a + /// separate process (for example Ballista and datafusion-distributed) + /// never poll the sibling partitions, so the single polled partition + /// drains the whole queue and reads files belonging to other partitions, + /// inflating the scan output by the partition count. Such executors + /// should set this to `false`, which makes each partition read only its + /// own file group. + pub enable_file_stream_work_stealing: bool, default = true + /// Aggregation ratio (number of distinct groups / number of input rows) /// threshold for skipping partial aggregation. If the value is greater /// then partial aggregation will skip aggregation for further input diff --git a/datafusion/datasource/src/file_scan_config/mod.rs b/datafusion/datasource/src/file_scan_config/mod.rs index 4bf86e17d387d..700488e3ae707 100644 --- a/datafusion/datasource/src/file_scan_config/mod.rs +++ b/datafusion/datasource/src/file_scan_config/mod.rs @@ -1038,10 +1038,18 @@ impl DataSource for FileScanConfig { /// during one execution. /// /// 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 - fn create_sibling_state(&self) -> Option> { - if self.preserve_order || self.partitioned_by_file_group { + /// when file order must be preserved, the file groups define the output + /// partitioning needed for the rest of the plan, or work stealing is + /// disabled via + /// `datafusion.execution.enable_file_stream_work_stealing`. + fn create_sibling_state( + &self, + config: &ConfigOptions, + ) -> Option> { + if self.preserve_order + || self.partitioned_by_file_group + || !config.execution.enable_file_stream_work_stealing + { return None; } diff --git a/datafusion/datasource/src/file_stream/mod.rs b/datafusion/datasource/src/file_stream/mod.rs index fd6fc4a3c912e..c5c18e37bc6a6 100644 --- a/datafusion/datasource/src/file_stream/mod.rs +++ b/datafusion/datasource/src/file_stream/mod.rs @@ -182,6 +182,7 @@ mod tests { use arrow::array::{AsArray, RecordBatch}; use arrow::datatypes::{DataType, Field, Int32Type, Schema}; use datafusion_common::DataFusionError; + use datafusion_common::config::ConfigOptions; use datafusion_common::error::Result; use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; @@ -1131,7 +1132,7 @@ mod tests { Ok(()) } - /// Reproduces . + /// Covers . /// /// Executors that run each output partition as an isolated task in a /// separate process (Ballista, datafusion-distributed) build one plan @@ -1139,64 +1140,86 @@ mod tests { /// sibling partitions are never polled in that process. A `DataSourceExec` /// seeds one shared work queue from every file group, so the single polled /// partition drains the whole queue and reads files that belong to other - /// partitions. Each isolated task does the same, inflating the scan output - /// by the partition count. + /// partitions, inflating the scan output by the partition count. /// - /// This test models that setup: it builds and drives only partition 0. - /// Partition 0 should read only its own file (`file1.parquet`), but today - /// it also reads partition 1's file, so it fails with - /// `"Batch: 101\nBatch: 201"`. - /// - /// Ignored because it fails on `main` by design: it is a caught regression - /// kept with its assertion intact. Run it with - /// `cargo test -p datafusion-datasource -- --ignored - /// regression_isolated_partition_reads_only_its_own_files`. + /// `datafusion.execution.enable_file_stream_work_stealing` turns the shared + /// queue off: `create_sibling_state` then returns `None` and each partition + /// reads only its own file group. This test drives only partition 0 (as an + /// isolated task would) and checks both behaviors: with the default + /// (stealing on) partition 0 also reads partition 1's file, and with the + /// flag off it reads only its own. #[tokio::test] - #[ignore = "reproduces #23293: shared FileStream work queue is drained by a single isolated task"] - async fn regression_isolated_partition_reads_only_its_own_files() -> Result<()> { - let test = FileStreamMorselTest::new() - .with_file_in_partition( - PartitionId(0), - MockPlanner::builder("file1.parquet") - .add_plan(MockPlanBuilder::new().with_morsel(MorselId(10), 101)) - .return_none(), - ) - .with_file_in_partition( - PartitionId(1), - MockPlanner::builder("file2.parquet") - .add_plan(MockPlanBuilder::new().with_morsel(MorselId(11), 201)) - .return_none(), - ); + async fn isolated_partition_respects_work_stealing_config() -> Result<()> { + // A two-file scan: partition 0 owns file1 (batch 101), partition 1 owns + // file2 (batch 201). Rebuilt per phase because the mock planners are + // consumed as files are opened. + let make_test = || { + FileStreamMorselTest::new() + .with_file_in_partition( + PartitionId(0), + MockPlanner::builder("file1.parquet") + .add_plan(MockPlanBuilder::new().with_morsel(MorselId(10), 101)) + .return_none(), + ) + .with_file_in_partition( + PartitionId(1), + MockPlanner::builder("file2.parquet") + .add_plan(MockPlanBuilder::new().with_morsel(MorselId(11), 201)) + .return_none(), + ) + }; + // Drive only partition 0, exactly as an isolated per-task executor does. + // Partition 1's stream is never created. + async fn drive_partition0( + test: &FileStreamMorselTest, + config: &FileScanConfig, + options: &ConfigOptions, + ) -> Result<(bool, String)> { + let shared_work_source = + config.create_sibling_state(options).and_then(|state| { + state.as_ref().downcast_ref::().cloned() + }); + let shared = shared_work_source.is_some(); + let metrics_set = ExecutionPlanMetricsSet::new(); + let partition0 = FileStreamBuilder::new(config) + .with_partition(0) + .with_shared_work_source(shared_work_source) + .with_morselizer(Box::new(test.morselizer.clone())) + .with_metrics(&metrics_set) + .build()?; + Ok((shared, drain_stream_output(partition0).await?)) + } + + // With work stealing enabled (the default), the single polled partition + // drains the shared queue and reads a sibling's file (issue #23293). + let mut options = ConfigOptions::default(); + assert!(options.execution.enable_file_stream_work_stealing); + let test = make_test(); let config = test.test_config(); - let metrics_set = ExecutionPlanMetricsSet::new(); - - // A `DataSourceExec` builds exactly one shared work source per - // execution, seeded from every file group. An isolated task builds that - // state and then executes only its own partition. - let shared_work_source = config - .create_sibling_state() - .and_then(|state| state.as_ref().downcast_ref::().cloned()); + let (shared, output) = drive_partition0(&test, &config, &options).await?; assert!( - shared_work_source.is_some(), - "a plain repartitioned scan should use a shared work source" + shared, + "the default config shares one work queue across siblings" + ); + assert_eq!( + output, "Batch: 101\nBatch: 201", + "with work stealing on, an isolated partition drains the shared queue" ); - // Build and drive ONLY partition 0. Partition 1's stream is never - // created, exactly as in an isolated per-task executor. - let partition0 = FileStreamBuilder::new(&config) - .with_partition(0) - .with_shared_work_source(shared_work_source) - .with_morselizer(Box::new(test.morselizer.clone())) - .with_metrics(&metrics_set) - .build()?; - - let output = drain_stream_output(partition0).await?; - + // Disabling work stealing drops the shared queue, so partition 0 reads + // only its own file group. + options.execution.enable_file_stream_work_stealing = false; + let test = make_test(); + let config = test.test_config(); + let (shared, output) = drive_partition0(&test, &config, &options).await?; + assert!( + !shared, + "disabling work stealing drops the shared work source" + ); assert_eq!( output, "Batch: 101", - "partition 0 read files belonging to a sibling partition: a single \ - isolated task drained the shared work queue (issue #23293)" + "with work stealing off, partition 0 reads only its own file group" ); Ok(()) @@ -1287,7 +1310,7 @@ mod tests { let unlimited_config = test.test_config(); let limited_config = test.clone().with_limit(1).test_config(); let shared_work_source = limited_config - .create_sibling_state() + .create_sibling_state(&ConfigOptions::default()) .and_then(|state| state.as_ref().downcast_ref::().cloned()) .expect("shared work source"); let limited_metrics = ExecutionPlanMetricsSet::new(); @@ -1539,9 +1562,11 @@ mod tests { // `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() - }); + let shared_work_source = config + .create_sibling_state(&ConfigOptions::default()) + .and_then(|state| { + state.as_ref().downcast_ref::().cloned() + }); if !self.build_streams_on_first_read { for partition in build_order { let stream = FileStreamBuilder::new(&config) diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs index af4bc09504937..d68271a05947d 100644 --- a/datafusion/datasource/src/source.rs +++ b/datafusion/datasource/src/source.rs @@ -245,7 +245,10 @@ pub trait DataSource: Any + Send + Sync + Debug { /// /// Returns `None` (the default) if this data source has /// no sibling-shared execution state. - fn create_sibling_state(&self) -> Option> { + fn create_sibling_state( + &self, + _config: &ConfigOptions, + ) -> Option> { None } @@ -391,7 +394,10 @@ impl ExecutionPlan for DataSourceExec { ) -> Result { let shared_state = self .execution_state - .get_or_init(|| self.data_source.create_sibling_state()) + .get_or_init(|| { + self.data_source + .create_sibling_state(context.session_config().options()) + }) .clone(); let args = OpenArgs::new(partition, Arc::clone(&context)) .with_shared_state(shared_state); diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index 37683b8e3095a..3a0a7fc4fcde8 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -218,6 +218,7 @@ datafusion.execution.batch_size 8192 datafusion.execution.coalesce_batches true datafusion.execution.collect_statistics true datafusion.execution.enable_ansi_mode false +datafusion.execution.enable_file_stream_work_stealing true datafusion.execution.enable_recursive_ctes true datafusion.execution.enforce_batch_size_in_joins false datafusion.execution.hash_join_buffering_capacity 0 @@ -372,6 +373,7 @@ datafusion.execution.batch_size 8192 Default batch size while creating new batch datafusion.execution.coalesce_batches true When set to true, record batches will be examined between each operator and small batches will be coalesced into larger batches. This is helpful when there are highly selective filters or joins that could produce tiny output batches. The target batch size is determined by the configuration setting datafusion.execution.collect_statistics true Should DataFusion collect statistics when first creating a table. Has no effect after the table is created. Applies to the default `ListingTableProvider` in DataFusion. Defaults to true. datafusion.execution.enable_ansi_mode false Whether to enable ANSI SQL mode. The flag is experimental and relevant only for DataFusion Spark built-in functions When `enable_ansi_mode` is set to `true`, the query engine follows ANSI SQL semantics for expressions, casting, and error handling. This means: - **Strict type coercion rules:** implicit casts between incompatible types are disallowed. - **Standard SQL arithmetic behavior:** operations such as division by zero, numeric overflow, or invalid casts raise runtime errors rather than returning `NULL` or adjusted values. - **Consistent ANSI behavior** for string concatenation, comparisons, and `NULL` handling. When `enable_ansi_mode` is `false` (the default), the engine uses a more permissive, non-ANSI mode designed for user convenience and backward compatibility. In this mode: - Implicit casts between types are allowed (e.g., string to integer when possible). - Arithmetic operations are more lenient — for example, `abs()` on the minimum representable integer value returns the input value instead of raising overflow. - Division by zero or invalid casts may return `NULL` instead of failing. # Default `false` — ANSI SQL mode is disabled by default. +datafusion.execution.enable_file_stream_work_stealing true When `true` (the default), sibling partition streams of a single file scan share a work queue: whichever output partition goes idle first steals the next unopened file (or byte-range morsel) from the shared queue. This balances work when all output partitions are polled concurrently in one process. Executors that run each output partition as an isolated task in a separate process (for example Ballista and datafusion-distributed) never poll the sibling partitions, so the single polled partition drains the whole queue and reads files belonging to other partitions, inflating the scan output by the partition count. Such executors should set this to `false`, which makes each partition read only its own file group. datafusion.execution.enable_recursive_ctes true Should DataFusion support recursive CTEs datafusion.execution.enforce_batch_size_in_joins false Should DataFusion enforce batch size in joins or not. By default, DataFusion will not enforce batch size in joins. Enforcing batch size in joins can reduce memory usage when joining large tables with a highly-selective join filter, but is also slightly slower. datafusion.execution.hash_join_buffering_capacity 0 How many bytes to buffer in the probe side of hash joins while the build side is concurrently being built. Without this, hash joins will wait until the full materialization of the build side before polling the probe side. This is useful in scenarios where the query is not completely CPU bounded, allowing to do some early work concurrently and reducing the latency of the query. Note that when hash join buffering is enabled, the probe side will start eagerly polling data, not giving time for the producer side of dynamic filters to produce any meaningful predicate. Queries with dynamic filters might see performance degradation. Disabled by default, set to a number greater than 0 for enabling it. diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index 6decc17d14ea9..ce305bf5e22ae 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -133,6 +133,7 @@ The following configuration settings are available: | datafusion.execution.enable_recursive_ctes | true | Should DataFusion support recursive CTEs | | datafusion.execution.split_file_groups_by_statistics | false | Attempt to eliminate sorts by packing & sorting files with non-overlapping statistics into the same file groups. Currently experimental | | datafusion.execution.keep_partition_by_columns | false | Should DataFusion keep the columns used for partition_by in the output RecordBatches | +| datafusion.execution.enable_file_stream_work_stealing | true | When `true` (the default), sibling partition streams of a single file scan share a work queue: whichever output partition goes idle first steals the next unopened file (or byte-range morsel) from the shared queue. This balances work when all output partitions are polled concurrently in one process. Executors that run each output partition as an isolated task in a separate process (for example Ballista and datafusion-distributed) never poll the sibling partitions, so the single polled partition drains the whole queue and reads files belonging to other partitions, inflating the scan output by the partition count. Such executors should set this to `false`, which makes each partition read only its own file group. | | datafusion.execution.skip_partial_aggregation_probe_ratio_threshold | 0.8 | Aggregation ratio (number of distinct groups / number of input rows) threshold for skipping partial aggregation. If the value is greater then partial aggregation will skip aggregation for further input | | datafusion.execution.skip_partial_aggregation_probe_rows_threshold | 100000 | Number of input rows partial aggregation partition should process, before aggregation ratio check and trying to switch to skipping aggregation mode | | datafusion.execution.use_row_number_estimates_to_optimize_partitioning | false | Should DataFusion use row number estimates at the input to decide whether increasing parallelism is beneficial or not. By default, only exact row numbers (not estimates) are used for this decision. Setting this flag to `true` will likely produce better plans. if the source of statistics is accurate. We plan to make this the default in the future. | From 35bd5754908e9e6e1d914a043df21928f26a064e Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 2 Jul 2026 08:39:46 -0600 Subject: [PATCH 3/4] refactor: simplify work-stealing regression test and document config param Drop the redundant bool from the test's drive_partition0 helper (assert the shared-queue state directly via create_sibling_state) and document the config parameter on DataSource::create_sibling_state. (cherry picked from commit b836a6d0cd197d7625b101473fe0d6e9ebc85bde) --- datafusion/datasource/src/file_stream/mod.rs | 17 ++++++++--------- datafusion/datasource/src/source.rs | 4 ++++ 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/datafusion/datasource/src/file_stream/mod.rs b/datafusion/datasource/src/file_stream/mod.rs index c5c18e37bc6a6..68f0b616d0378 100644 --- a/datafusion/datasource/src/file_stream/mod.rs +++ b/datafusion/datasource/src/file_stream/mod.rs @@ -1175,12 +1175,11 @@ mod tests { test: &FileStreamMorselTest, config: &FileScanConfig, options: &ConfigOptions, - ) -> Result<(bool, String)> { + ) -> Result { let shared_work_source = config.create_sibling_state(options).and_then(|state| { state.as_ref().downcast_ref::().cloned() }); - let shared = shared_work_source.is_some(); let metrics_set = ExecutionPlanMetricsSet::new(); let partition0 = FileStreamBuilder::new(config) .with_partition(0) @@ -1188,7 +1187,7 @@ mod tests { .with_morselizer(Box::new(test.morselizer.clone())) .with_metrics(&metrics_set) .build()?; - Ok((shared, drain_stream_output(partition0).await?)) + drain_stream_output(partition0).await } // With work stealing enabled (the default), the single polled partition @@ -1197,13 +1196,13 @@ mod tests { assert!(options.execution.enable_file_stream_work_stealing); let test = make_test(); let config = test.test_config(); - let (shared, output) = drive_partition0(&test, &config, &options).await?; assert!( - shared, + config.create_sibling_state(&options).is_some(), "the default config shares one work queue across siblings" ); assert_eq!( - output, "Batch: 101\nBatch: 201", + drive_partition0(&test, &config, &options).await?, + "Batch: 101\nBatch: 201", "with work stealing on, an isolated partition drains the shared queue" ); @@ -1212,13 +1211,13 @@ mod tests { options.execution.enable_file_stream_work_stealing = false; let test = make_test(); let config = test.test_config(); - let (shared, output) = drive_partition0(&test, &config, &options).await?; assert!( - !shared, + config.create_sibling_state(&options).is_none(), "disabling work stealing drops the shared work source" ); assert_eq!( - output, "Batch: 101", + drive_partition0(&test, &config, &options).await?, + "Batch: 101", "with work stealing off, partition 0 reads only its own file group" ); diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs index d68271a05947d..6bb574fcadff3 100644 --- a/datafusion/datasource/src/source.rs +++ b/datafusion/datasource/src/source.rs @@ -243,6 +243,10 @@ pub trait DataSource: Any + Send + Sync + Debug { /// Create per execution state to share across sibling instances of this /// data source during one execution. /// + /// `config` is the session configuration, so implementations can honor + /// options that disable sibling sharing (returning `None`) for consumers + /// that cannot poll all partitions in one process. + /// /// Returns `None` (the default) if this data source has /// no sibling-shared execution state. fn create_sibling_state( From d22cd3de7c103eb1a2c104b7c2116ace54b1fed3 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 2 Jul 2026 09:31:28 -0600 Subject: [PATCH 4/4] refactor: address review feedback on work-stealing config - Reword the config docstring in user-visible terms (dynamic runtime rebalancing) rather than internal work-queue mechanics. - Note on FileScanConfig::file_groups that files may be reassigned across partitions at runtime when work stealing is enabled unless preserve_order or partitioned_by_file_group is set. - Replace the bespoke isolated-partition test with the file's standard morsel snapshot harness: FileStreamMorselTest gains with_enable_file_stream_work_stealing, and the test reuses two_partition_morsel_test to show that disabling the flag keeps each partition's files local. Regenerate configs.md and information_schema.slt for the reworded docstring. (cherry picked from commit 608915e84dd2273626df733579207ba532a566a8) --- datafusion/common/src/config.rs | 22 ++- .../datasource/src/file_scan_config/mod.rs | 6 + datafusion/datasource/src/file_stream/mod.rs | 130 ++++++------------ .../test_files/information_schema.slt | 2 +- docs/source/user-guide/configs.md | 2 +- 5 files changed, 60 insertions(+), 102 deletions(-) diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 3b2cde43f695b..cbfa971055f5e 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -639,19 +639,17 @@ config_namespace! { /// Should DataFusion keep the columns used for partition_by in the output RecordBatches pub keep_partition_by_columns: bool, default = false - /// When `true` (the default), sibling partition streams of a single file - /// scan share a work queue: whichever output partition goes idle first - /// steals the next unopened file (or byte-range morsel) from the shared - /// queue. This balances work when all output partitions are polled - /// concurrently in one process. + /// When `true` (the default), DataFusion's built-in file scans + /// dynamically rebalance files across partitions at query execution + /// time: a partition that goes idle reads files (or byte-range morsels) + /// originally assigned to a sibling partition, which keeps all + /// partitions busy in a single process. /// - /// Executors that run each output partition as an isolated task in a - /// separate process (for example Ballista and datafusion-distributed) - /// never poll the sibling partitions, so the single polled partition - /// drains the whole queue and reads files belonging to other partitions, - /// inflating the scan output by the partition count. Such executors - /// should set this to `false`, which makes each partition read only its - /// own file group. + /// Executors that depend on the plan-time partition assignment — such as + /// Ballista and datafusion-distributed, which run each partition as an + /// isolated task and never poll the siblings — should set this to + /// `false` so each partition reads only its own file group and no + /// runtime reassignment occurs. pub enable_file_stream_work_stealing: bool, default = true /// Aggregation ratio (number of distinct groups / number of input rows) diff --git a/datafusion/datasource/src/file_scan_config/mod.rs b/datafusion/datasource/src/file_scan_config/mod.rs index 700488e3ae707..050cd54bb0adf 100644 --- a/datafusion/datasource/src/file_scan_config/mod.rs +++ b/datafusion/datasource/src/file_scan_config/mod.rs @@ -159,6 +159,12 @@ pub struct FileScanConfig { /// DataFusion may attempt to read each partition of files /// concurrently, however files *within* a partition will be read /// sequentially, one after the next. + /// + /// Note that when `datafusion.execution.enable_file_stream_work_stealing` + /// is enabled (the default), files may be reassigned to a different + /// partition at runtime unless `preserve_order` or + /// `partitioned_by_file_group` is set, so a file is not guaranteed to be + /// read by the partition it is grouped under here. pub file_groups: Vec, /// Table constraints pub constraints: Constraints, diff --git a/datafusion/datasource/src/file_stream/mod.rs b/datafusion/datasource/src/file_stream/mod.rs index 68f0b616d0378..9ea9a88e393c4 100644 --- a/datafusion/datasource/src/file_stream/mod.rs +++ b/datafusion/datasource/src/file_stream/mod.rs @@ -1132,94 +1132,36 @@ mod tests { Ok(()) } - /// Covers . + /// Verifies that disabling `enable_file_stream_work_stealing` keeps each + /// stream's files local, so a sibling cannot steal them at runtime. /// - /// Executors that run each output partition as an isolated task in a - /// separate process (Ballista, datafusion-distributed) build one plan - /// instance and call `execute(partition)` for a single partition. The - /// sibling partitions are never polled in that process. A `DataSourceExec` - /// seeds one shared work queue from every file group, so the single polled - /// partition drains the whole queue and reads files that belong to other - /// partitions, inflating the scan output by the partition count. - /// - /// `datafusion.execution.enable_file_stream_work_stealing` turns the shared - /// queue off: `create_sibling_state` then returns `None` and each partition - /// reads only its own file group. This test drives only partition 0 (as an - /// isolated task would) and checks both behaviors: with the default - /// (stealing on) partition 0 also reads partition 1's file, and with the - /// flag off it reads only its own. + /// Covers : executors + /// that run each output partition as an isolated task in a separate process + /// (Ballista, datafusion-distributed) poll only their own partition, so the + /// shared work queue would let that one partition drain files belonging to + /// its siblings. Disabling the flag falls back to per-partition file groups. #[tokio::test] - async fn isolated_partition_respects_work_stealing_config() -> Result<()> { - // A two-file scan: partition 0 owns file1 (batch 101), partition 1 owns - // file2 (batch 201). Rebuilt per phase because the mock planners are - // consumed as files are opened. - let make_test = || { - FileStreamMorselTest::new() - .with_file_in_partition( - PartitionId(0), - MockPlanner::builder("file1.parquet") - .add_plan(MockPlanBuilder::new().with_morsel(MorselId(10), 101)) - .return_none(), - ) - .with_file_in_partition( - PartitionId(1), - MockPlanner::builder("file2.parquet") - .add_plan(MockPlanBuilder::new().with_morsel(MorselId(11), 201)) - .return_none(), - ) - }; - - // Drive only partition 0, exactly as an isolated per-task executor does. - // Partition 1's stream is never created. - async fn drive_partition0( - test: &FileStreamMorselTest, - config: &FileScanConfig, - options: &ConfigOptions, - ) -> Result { - let shared_work_source = - config.create_sibling_state(options).and_then(|state| { - state.as_ref().downcast_ref::().cloned() - }); - let metrics_set = ExecutionPlanMetricsSet::new(); - let partition0 = FileStreamBuilder::new(config) - .with_partition(0) - .with_shared_work_source(shared_work_source) - .with_morselizer(Box::new(test.morselizer.clone())) - .with_metrics(&metrics_set) - .build()?; - drain_stream_output(partition0).await - } - - // With work stealing enabled (the default), the single polled partition - // drains the shared queue and reads a sibling's file (issue #23293). - let mut options = ConfigOptions::default(); - assert!(options.execution.enable_file_stream_work_stealing); - let test = make_test(); - let config = test.test_config(); - assert!( - config.create_sibling_state(&options).is_some(), - "the default config shares one work queue across siblings" - ); - assert_eq!( - drive_partition0(&test, &config, &options).await?, - "Batch: 101\nBatch: 201", - "with work stealing on, an isolated partition drains the shared queue" - ); + async fn morsel_disabled_work_stealing_keeps_files_local() -> Result<()> { + // same fixture as `morsel_shared_files_can_be_stolen`, but with work + // stealing disabled via config + let test = two_partition_morsel_test() + .with_enable_file_stream_work_stealing(false) + .with_file_stream_events(false); - // Disabling work stealing drops the shared queue, so partition 0 reads - // only its own file group. - options.execution.enable_file_stream_work_stealing = false; - let test = make_test(); - let config = test.test_config(); - assert!( - config.create_sibling_state(&options).is_none(), - "disabling work stealing drops the shared work source" - ); - assert_eq!( - drive_partition0(&test, &config, &options).await?, - "Batch: 101", - "with work stealing off, partition 0 reads only its own file group" - ); + // Even though Partition 1 is polled first, it cannot steal the three + // files assigned to Partition 0; each partition reads only its own. + insta::assert_snapshot!(test.run().await.unwrap(), @r" + ----- Partition 0 ----- + Batch: 101 + Batch: 102 + Batch: 103 + Done + ----- Partition 1 ----- + Batch: 201 + Done + ----- File Stream Events ----- + (omitted due to with_file_stream_events(false)) + "); Ok(()) } @@ -1425,6 +1367,7 @@ mod tests { partition_files: BTreeMap>, preserve_order: bool, partitioned_by_file_group: bool, + enable_file_stream_work_stealing: bool, file_stream_events: bool, build_streams_on_first_read: bool, reads: Vec, @@ -1439,6 +1382,7 @@ mod tests { partition_files: BTreeMap::new(), preserve_order: false, partitioned_by_file_group: false, + enable_file_stream_work_stealing: true, file_stream_events: true, build_streams_on_first_read: false, reads: vec![], @@ -1484,6 +1428,14 @@ mod tests { self } + /// Sets `datafusion.execution.enable_file_stream_work_stealing`. When + /// disabled, each stream keeps its own files local instead of sharing a + /// work queue with its siblings. + fn with_enable_file_stream_work_stealing(mut self, enable: bool) -> Self { + self.enable_file_stream_work_stealing = enable; + self + } + /// Controls whether scheduler events are included in the snapshot. /// /// When disabled, `run()` still includes the event section header but @@ -1561,9 +1513,11 @@ mod tests { // `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(&ConfigOptions::default()) - .and_then(|state| { + let mut options = ConfigOptions::default(); + options.execution.enable_file_stream_work_stealing = + self.enable_file_stream_work_stealing; + let shared_work_source = + config.create_sibling_state(&options).and_then(|state| { state.as_ref().downcast_ref::().cloned() }); if !self.build_streams_on_first_read { diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index 3a0a7fc4fcde8..6f33483031ace 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -373,7 +373,7 @@ datafusion.execution.batch_size 8192 Default batch size while creating new batch datafusion.execution.coalesce_batches true When set to true, record batches will be examined between each operator and small batches will be coalesced into larger batches. This is helpful when there are highly selective filters or joins that could produce tiny output batches. The target batch size is determined by the configuration setting datafusion.execution.collect_statistics true Should DataFusion collect statistics when first creating a table. Has no effect after the table is created. Applies to the default `ListingTableProvider` in DataFusion. Defaults to true. datafusion.execution.enable_ansi_mode false Whether to enable ANSI SQL mode. The flag is experimental and relevant only for DataFusion Spark built-in functions When `enable_ansi_mode` is set to `true`, the query engine follows ANSI SQL semantics for expressions, casting, and error handling. This means: - **Strict type coercion rules:** implicit casts between incompatible types are disallowed. - **Standard SQL arithmetic behavior:** operations such as division by zero, numeric overflow, or invalid casts raise runtime errors rather than returning `NULL` or adjusted values. - **Consistent ANSI behavior** for string concatenation, comparisons, and `NULL` handling. When `enable_ansi_mode` is `false` (the default), the engine uses a more permissive, non-ANSI mode designed for user convenience and backward compatibility. In this mode: - Implicit casts between types are allowed (e.g., string to integer when possible). - Arithmetic operations are more lenient — for example, `abs()` on the minimum representable integer value returns the input value instead of raising overflow. - Division by zero or invalid casts may return `NULL` instead of failing. # Default `false` — ANSI SQL mode is disabled by default. -datafusion.execution.enable_file_stream_work_stealing true When `true` (the default), sibling partition streams of a single file scan share a work queue: whichever output partition goes idle first steals the next unopened file (or byte-range morsel) from the shared queue. This balances work when all output partitions are polled concurrently in one process. Executors that run each output partition as an isolated task in a separate process (for example Ballista and datafusion-distributed) never poll the sibling partitions, so the single polled partition drains the whole queue and reads files belonging to other partitions, inflating the scan output by the partition count. Such executors should set this to `false`, which makes each partition read only its own file group. +datafusion.execution.enable_file_stream_work_stealing true When `true` (the default), DataFusion's built-in file scans dynamically rebalance files across partitions at query execution time: a partition that goes idle reads files (or byte-range morsels) originally assigned to a sibling partition, which keeps all partitions busy in a single process. Executors that depend on the plan-time partition assignment — such as Ballista and datafusion-distributed, which run each partition as an isolated task and never poll the siblings — should set this to `false` so each partition reads only its own file group and no runtime reassignment occurs. datafusion.execution.enable_recursive_ctes true Should DataFusion support recursive CTEs datafusion.execution.enforce_batch_size_in_joins false Should DataFusion enforce batch size in joins or not. By default, DataFusion will not enforce batch size in joins. Enforcing batch size in joins can reduce memory usage when joining large tables with a highly-selective join filter, but is also slightly slower. datafusion.execution.hash_join_buffering_capacity 0 How many bytes to buffer in the probe side of hash joins while the build side is concurrently being built. Without this, hash joins will wait until the full materialization of the build side before polling the probe side. This is useful in scenarios where the query is not completely CPU bounded, allowing to do some early work concurrently and reducing the latency of the query. Note that when hash join buffering is enabled, the probe side will start eagerly polling data, not giving time for the producer side of dynamic filters to produce any meaningful predicate. Queries with dynamic filters might see performance degradation. Disabled by default, set to a number greater than 0 for enabling it. diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index ce305bf5e22ae..e8efee5e00368 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -133,7 +133,7 @@ The following configuration settings are available: | datafusion.execution.enable_recursive_ctes | true | Should DataFusion support recursive CTEs | | datafusion.execution.split_file_groups_by_statistics | false | Attempt to eliminate sorts by packing & sorting files with non-overlapping statistics into the same file groups. Currently experimental | | datafusion.execution.keep_partition_by_columns | false | Should DataFusion keep the columns used for partition_by in the output RecordBatches | -| datafusion.execution.enable_file_stream_work_stealing | true | When `true` (the default), sibling partition streams of a single file scan share a work queue: whichever output partition goes idle first steals the next unopened file (or byte-range morsel) from the shared queue. This balances work when all output partitions are polled concurrently in one process. Executors that run each output partition as an isolated task in a separate process (for example Ballista and datafusion-distributed) never poll the sibling partitions, so the single polled partition drains the whole queue and reads files belonging to other partitions, inflating the scan output by the partition count. Such executors should set this to `false`, which makes each partition read only its own file group. | +| datafusion.execution.enable_file_stream_work_stealing | true | When `true` (the default), DataFusion's built-in file scans dynamically rebalance files across partitions at query execution time: a partition that goes idle reads files (or byte-range morsels) originally assigned to a sibling partition, which keeps all partitions busy in a single process. Executors that depend on the plan-time partition assignment — such as Ballista and datafusion-distributed, which run each partition as an isolated task and never poll the siblings — should set this to `false` so each partition reads only its own file group and no runtime reassignment occurs. | | datafusion.execution.skip_partial_aggregation_probe_ratio_threshold | 0.8 | Aggregation ratio (number of distinct groups / number of input rows) threshold for skipping partial aggregation. If the value is greater then partial aggregation will skip aggregation for further input | | datafusion.execution.skip_partial_aggregation_probe_rows_threshold | 100000 | Number of input rows partial aggregation partition should process, before aggregation ratio check and trying to switch to skipping aggregation mode | | datafusion.execution.use_row_number_estimates_to_optimize_partitioning | false | Should DataFusion use row number estimates at the input to decide whether increasing parallelism is beneficial or not. By default, only exact row numbers (not estimates) are used for this decision. Setting this flag to `true` will likely produce better plans. if the source of statistics is accurate. We plan to make this the default in the future. |