[branch-54] Add datafusion.execution.enable_file_stream_work_stealing config (#23296)

## Which issue does this PR close?

- Backport of #23294 to `branch-54`.
- Closes #23293 for the 54 release line.

## Rationale for this change

`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 unopened 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 and datafusion-distributed — never poll the
sibling partitions. The single polled partition drains the whole shared
queue and reads files belonging to other partitions, so every isolated
task reads the entire input and the scan output is inflated by the
partition count. This is a correctness bug for those executors, not just
a performance one.

The existing escape hatches (`preserve_order`,
`partitioned_by_file_group`) are plan-level flags on `FileScanConfig`,
not something a distributed executor can set centrally through the
session config, and a plain repartitioned scan does not set
`partitioned_by_file_group`. There is no session-level off switch,
unlike `datafusion.optimizer.enable_dynamic_filter_pushdown`, which
exists precisely so consumers that cannot support runtime
cross-partition state can disable it.

## What changes are included in this PR?

This is a backport of #23294. The changes are identical in intent; two
conflicts were resolved for `branch-54`, which does not yet carry the
`output_partitioning` field on `FileScanConfig` or the
`enable_migration_aggregate` config that exist on `main`:

- Add `datafusion.execution.enable_file_stream_work_stealing` (default
`true`). When `false`, `FileScanConfig::create_sibling_state` returns
`None`, so each partition falls back to `WorkSource::Local` and reads
only its own file group.
- Thread `&ConfigOptions` into `DataSource::create_sibling_state` so the
flag is read from the session config at `execute` time. As a session
config value it round-trips through `datafusion-proto` with no proto
schema change.
- Regenerate `configs.md` and add the setting to
`information_schema.slt`.
- Add the regression test
`morsel_disabled_work_stealing_keeps_files_local` using the file's
standard morsel snapshot harness.

## Are these changes tested?

Yes. The `file_stream` test suite (including the new
`morsel_disabled_work_stealing_keeps_files_local`) passes, and `cargo
clippy` is clean. The existing sibling work-stealing tests continue to
pass with the default.

## Are there any user-facing changes?

A new session config,
`datafusion.execution.enable_file_stream_work_stealing` (default
`true`), so existing behavior is unchanged.
`DataSource::create_sibling_state` gains a `&ConfigOptions` parameter
(an API change for anyone implementing the `DataSource` trait directly).

Author: andygrove

datafusion/common/src/config.rs

@@ -639,6 +639,19 @@ 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), 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.
+        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

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<FileGroup>,
     /// Table constraints
     pub constraints: Constraints,
@@ -1038,10 +1044,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<Arc<dyn Any + Send + Sync>> {
-        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<Arc<dyn Any + Send + Sync>> {
+        if self.preserve_order
+            || self.partitioned_by_file_group
+            || !config.execution.enable_file_stream_work_stealing
+        {
             return None;
         }
 

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,6 +1132,40 @@ mod tests {
         Ok(())
     }
 
+    /// Verifies that disabling `enable_file_stream_work_stealing` keeps each
+    /// stream's files local, so a sibling cannot steal them at runtime.
+    ///
+    /// Covers <https://github.com/apache/datafusion/issues/23293>: 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 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);
+
+        // 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(())
+    }
+
     /// Verifies that an empty sibling can immediately steal shared files when
     /// it is polled before the stream that originally owned them.
     #[tokio::test]
@@ -1216,7 +1251,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::<SharedWorkSource>().cloned())
             .expect("shared work source");
         let limited_metrics = ExecutionPlanMetricsSet::new();
@@ -1332,6 +1367,7 @@ mod tests {
         partition_files: BTreeMap<PartitionId, Vec<String>>,
         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<PartitionId>,
@@ -1346,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![],
@@ -1391,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
@@ -1468,9 +1513,13 @@ 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::<SharedWorkSource>().cloned()
-            });
+            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::<SharedWorkSource>().cloned()
+                });
             if !self.build_streams_on_first_read {
                 for partition in build_order {
                     let stream = FileStreamBuilder::new(&config)

datafusion/datasource/src/source.rs

@@ -243,9 +243,16 @@ 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(&self) -> Option<Arc<dyn Any + Send + Sync>> {
+    fn create_sibling_state(
+        &self,
+        _config: &ConfigOptions,
+    ) -> Option<Arc<dyn Any + Send + Sync>> {
         None
     }
 
@@ -391,7 +398,10 @@ impl ExecutionPlan for DataSourceExec {
     ) -> Result<SendableRecordBatchStream> {
         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);

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), 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.