feat: cap spill merge fan-in

datafusion/core/src/execution/context/mod.rs

@@ -82,7 +82,7 @@ use datafusion_execution::cache::cache_manager::{
 };
 pub use datafusion_execution::config::SessionConfig;
 use datafusion_execution::disk_manager::{
-    DEFAULT_MAX_TEMP_DIRECTORY_SIZE, DiskManagerBuilder,
+    DEFAULT_MAX_SPILL_MERGE_FAN_IN, DEFAULT_MAX_TEMP_DIRECTORY_SIZE, DiskManagerBuilder,
 };
 use datafusion_execution::registry::SerializerRegistry;
 use datafusion_expr::HigherOrderUDF;
@@ -1208,6 +1208,14 @@ impl SessionContext {
                 let limit = Self::parse_capacity_limit(variable, value)?;
                 builder.with_file_statistics_cache_limit(limit)
             }
+            "max_spill_merge_fan_in" => {
+                let fan_in = value.parse::<usize>().map_err(|e| {
+                    DataFusionError::Plan(format!(
+                        "Failed to parse non-negative integer from '{variable}', value '{value}': {e}"
+                    ))
+                })?;
+                builder.with_max_spill_merge_fan_in(fan_in)
+            }
             _ => return plan_err!("Unknown runtime configuration: {variable}"),
             // Remember to update `reset_runtime_variable()` when adding new options
         };
@@ -1252,6 +1260,10 @@ impl SessionContext {
                     DEFAULT_FILE_STATISTICS_MEMORY_LIMIT,
                 );
             }
+            "max_spill_merge_fan_in" => {
+                builder =
+                    builder.with_max_spill_merge_fan_in(DEFAULT_MAX_SPILL_MERGE_FAN_IN);
+            }
             _ => return plan_err!("Unknown runtime configuration: {variable}"),
         };
         *state = SessionStateBuilder::from(state.clone())

datafusion/core/tests/sql/runtime_config.rs

@@ -227,6 +227,34 @@ async fn test_max_temp_directory_size_enforcement() {
     );
 }
 
+#[tokio::test]
+async fn test_max_spill_merge_fan_in_runtime_config() {
+    let ctx = SessionContext::new();
+
+    ctx.sql("SET datafusion.runtime.max_spill_merge_fan_in = '8'")
+        .await
+        .unwrap()
+        .collect()
+        .await
+        .unwrap();
+    assert_eq!(ctx.runtime_env().disk_manager.max_spill_merge_fan_in(), 8);
+
+    ctx.sql("RESET datafusion.runtime.max_spill_merge_fan_in")
+        .await
+        .unwrap()
+        .collect()
+        .await
+        .unwrap();
+    assert_eq!(ctx.runtime_env().disk_manager.max_spill_merge_fan_in(), 0);
+
+    let error = ctx
+        .sql("SET datafusion.runtime.max_spill_merge_fan_in = '-1'")
+        .await
+        .unwrap_err()
+        .to_string();
+    assert!(error.contains("Failed to parse non-negative integer"));
+}
+
 #[tokio::test]
 async fn test_test_metadata_cache_limit() {
     let ctx = SessionContext::new();

datafusion/execution/src/disk_manager.rs

@@ -32,6 +32,7 @@ use std::sync::Arc;
 use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
 use tempfile::{Builder, NamedTempFile, TempDir};
 pub const DEFAULT_MAX_TEMP_DIRECTORY_SIZE: u64 = 100 * 1024 * 1024 * 1024; // 100GB
+pub const DEFAULT_MAX_SPILL_MERGE_FAN_IN: usize = 0;
 
 /// Builder pattern for the [DiskManager] structure
 #[derive(Clone)]
@@ -41,6 +42,9 @@ pub struct DiskManagerBuilder {
     /// The maximum amount of data (in bytes) stored inside the temporary directories.
     /// Default to 100GB
     max_temp_directory_size: u64,
+    /// Maximum number of spill files opened by one external merge pass.
+    /// A value of 0 means unlimited.
+    max_spill_merge_fan_in: usize,
 }
 impl Debug for DiskManagerBuilder {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
@@ -55,6 +59,7 @@ impl Default for DiskManagerBuilder {
         Self {
             mode: DiskManagerMode::OsTmpDirectory,
             max_temp_directory_size: DEFAULT_MAX_TEMP_DIRECTORY_SIZE,
+            max_spill_merge_fan_in: DEFAULT_MAX_SPILL_MERGE_FAN_IN,
         }
     }
 }
@@ -78,12 +83,22 @@ impl DiskManagerBuilder {
         self
     }
 
+    pub fn set_max_spill_merge_fan_in(&mut self, value: usize) {
+        self.max_spill_merge_fan_in = value;
+    }
+
+    pub fn with_max_spill_merge_fan_in(mut self, value: usize) -> Self {
+        self.set_max_spill_merge_fan_in(value);
+        self
+    }
+
     /// Create a DiskManager given the builder
     pub fn build(self) -> Result<DiskManager> {
         match self.mode {
             DiskManagerMode::OsTmpDirectory => Ok(DiskManager {
                 local_dirs: Mutex::new(Some(vec![])),
                 max_temp_directory_size: AtomicU64::new(self.max_temp_directory_size),
+                max_spill_merge_fan_in: AtomicUsize::new(self.max_spill_merge_fan_in),
                 used_disk_space: Arc::new(AtomicU64::new(0)),
                 active_files_count: Arc::new(AtomicUsize::new(0)),
                 factory: None,
@@ -96,6 +111,7 @@ impl DiskManagerBuilder {
                 Ok(DiskManager {
                     local_dirs: Mutex::new(Some(local_dirs)),
                     max_temp_directory_size: AtomicU64::new(self.max_temp_directory_size),
+                    max_spill_merge_fan_in: AtomicUsize::new(self.max_spill_merge_fan_in),
                     used_disk_space: Arc::new(AtomicU64::new(0)),
                     active_files_count: Arc::new(AtomicUsize::new(0)),
                     factory: None,
@@ -104,13 +120,15 @@ impl DiskManagerBuilder {
             DiskManagerMode::Disabled => Ok(DiskManager {
                 local_dirs: Mutex::new(None),
                 max_temp_directory_size: AtomicU64::new(self.max_temp_directory_size),
+                max_spill_merge_fan_in: AtomicUsize::new(self.max_spill_merge_fan_in),
                 used_disk_space: Arc::new(AtomicU64::new(0)),
                 active_files_count: Arc::new(AtomicUsize::new(0)),
                 factory: None,
             }),
             DiskManagerMode::Custom(factory) => Ok(DiskManager {
                 local_dirs: Mutex::new(None),
                 max_temp_directory_size: AtomicU64::new(self.max_temp_directory_size),
+                max_spill_merge_fan_in: AtomicUsize::new(self.max_spill_merge_fan_in),
                 used_disk_space: Arc::new(AtomicU64::new(0)),
                 active_files_count: Arc::new(AtomicUsize::new(0)),
                 factory: Some(factory),
@@ -161,6 +179,9 @@ pub struct DiskManager {
     /// Default to 100GB. Stored as `AtomicU64` so it can be adjusted at runtime
     /// without requiring exclusive (`&mut`) access to the `DiskManager`.
     max_temp_directory_size: AtomicU64,
+    /// Maximum number of spill files opened by one external merge pass.
+    /// A value of 0 preserves the memory-driven, unbounded behavior.
+    max_spill_merge_fan_in: AtomicUsize,
     /// Used disk space in the temporary directories. Now only spilled data for
     /// external executors are counted.
     used_disk_space: Arc<AtomicU64>,
@@ -250,6 +271,22 @@ impl DiskManager {
         self.max_temp_directory_size.load(Ordering::Relaxed)
     }
 
+    /// Atomically set the maximum spill merge fan-in.
+    ///
+    /// A value of 0 disables the cap. Values of 1 are accepted but external
+    /// merge code will still merge at least two spill streams to make progress.
+    pub fn set_max_spill_merge_fan_in(&self, max_spill_merge_fan_in: usize) {
+        self.max_spill_merge_fan_in
+            .store(max_spill_merge_fan_in, Ordering::Relaxed);
+    }
+
+    /// Returns the maximum number of spill files opened by one merge pass.
+    ///
+    /// A value of 0 means unlimited.
+    pub fn max_spill_merge_fan_in(&self) -> usize {
+        self.max_spill_merge_fan_in.load(Ordering::Relaxed)
+    }
+
     /// Returns the current spilling progress
     pub fn spilling_progress(&self) -> SpillingProgress {
         SpillingProgress {
@@ -905,6 +942,25 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_max_spill_merge_fan_in_builder_and_dynamic_update() -> Result<()> {
+        let dm = Arc::new(
+            DiskManager::builder()
+                .with_max_spill_merge_fan_in(8)
+                .build()?,
+        );
+
+        assert_eq!(dm.max_spill_merge_fan_in(), 8);
+
+        dm.set_max_spill_merge_fan_in(4);
+        assert_eq!(dm.max_spill_merge_fan_in(), 4);
+
+        dm.set_max_spill_merge_fan_in(0);
+        assert_eq!(dm.max_spill_merge_fan_in(), 0);
+
+        Ok(())
+    }
+
     #[test]
     fn test_disabled_disk_manager_rejects_nonzero_limit() -> Result<()> {
         let dm = DiskManager::builder()