fix(query): the workload_group sort spill needs to be able to be forced to be turned off (#18561)

fix

Author: forsaken628

src/query/pipeline/transforms/src/processors/memory_settings.rs

@@ -22,16 +22,19 @@ use databend_common_base::runtime::ThreadTracker;
 use databend_common_base::runtime::GLOBAL_MEM_STAT;
 
 #[derive(Clone)]
+#[non_exhaustive]
 pub struct MemorySettings {
-    pub max_memory_usage: usize,
+    pub spill_unit_size: usize,
+
     pub enable_global_level_spill: bool,
+    pub max_memory_usage: usize,
     pub global_memory_tracking: &'static MemStat,
 
+    pub enable_group_spill: bool,
+
+    pub enable_query_level_spill: bool,
     pub max_query_memory_usage: usize,
     pub query_memory_tracking: Option<Arc<MemStat>>,
-    pub enable_query_level_spill: bool,
-
-    pub spill_unit_size: usize,
 }
 
 impl Debug for MemorySettings {
@@ -71,28 +74,74 @@ impl Debug for MemorySettings {
     }
 }
 
-impl MemorySettings {
-    pub fn disable_spill() -> MemorySettings {
+pub struct MemorySettingsBuilder {
+    enable_global_level_spill: bool,
+    max_memory_usage: Option<usize>,
+
+    enable_group_spill: bool,
+
+    enable_query_level_spill: bool,
+    max_query_memory_usage: Option<usize>,
+    query_memory_tracking: Option<Arc<MemStat>>,
+
+    spill_unit_size: Option<usize>,
+}
+
+impl MemorySettingsBuilder {
+    pub fn with_max_memory_usage(mut self, max: usize) -> Self {
+        self.enable_global_level_spill = true;
+        self.max_memory_usage = Some(max);
+        self
+    }
+
+    pub fn with_max_query_memory_usage(
+        mut self,
+        max: usize,
+        tracking: Option<Arc<MemStat>>,
+    ) -> Self {
+        self.enable_query_level_spill = true;
+        self.max_query_memory_usage = Some(max);
+        self.query_memory_tracking = tracking;
+        self
+    }
+
+    pub fn with_workload_group(mut self, enable: bool) -> Self {
+        self.enable_group_spill = enable;
+        self
+    }
+
+    pub fn with_spill_unit_size(mut self, spill_unit_size: usize) -> Self {
+        self.spill_unit_size = Some(spill_unit_size);
+        self
+    }
+
+    pub fn build(self) -> MemorySettings {
         MemorySettings {
-            spill_unit_size: 0,
-            max_memory_usage: usize::MAX,
-            enable_global_level_spill: false,
-            max_query_memory_usage: usize::MAX,
-            query_memory_tracking: None,
-            enable_query_level_spill: false,
+            enable_group_spill: self.enable_group_spill,
+            max_memory_usage: self.max_memory_usage.unwrap_or(usize::MAX),
+            enable_global_level_spill: self.enable_global_level_spill,
             global_memory_tracking: &GLOBAL_MEM_STAT,
+            enable_query_level_spill: self.enable_query_level_spill,
+            max_query_memory_usage: self.max_query_memory_usage.unwrap_or(usize::MAX),
+            query_memory_tracking: self.query_memory_tracking,
+            spill_unit_size: self.spill_unit_size.unwrap_or(0),
         }
     }
+}
 
-    pub fn always_spill(spill_unit_size: usize) -> MemorySettings {
-        MemorySettings {
-            spill_unit_size,
-            max_memory_usage: 0,
-            max_query_memory_usage: 0,
-            enable_query_level_spill: true,
-            enable_global_level_spill: true,
-            global_memory_tracking: &GLOBAL_MEM_STAT,
+impl MemorySettings {
+    pub fn builder() -> MemorySettingsBuilder {
+        MemorySettingsBuilder {
+            enable_global_level_spill: false,
+            max_memory_usage: None,
+
+            enable_group_spill: true,
+
+            enable_query_level_spill: false,
+            max_query_memory_usage: None,
             query_memory_tracking: None,
+
+            spill_unit_size: None,
         }
     }
 
@@ -103,7 +152,9 @@ impl MemorySettings {
             return true;
         }
 
-        if let Some(workload_group) = ThreadTracker::workload_group() {
+        if self.enable_group_spill
+            && let Some(workload_group) = ThreadTracker::workload_group()
+        {
             let workload_group_memory_usage = workload_group.mem_stat.get_memory_usage();
             let max_memory_usage = workload_group.max_memory_usage.load(Ordering::Relaxed);
 
@@ -132,6 +183,10 @@ impl MemorySettings {
     }
 
     fn check_workload_group(&self) -> Option<isize> {
+        if !self.enable_group_spill {
+            return None;
+        }
+
         let workload_group = ThreadTracker::workload_group()?;
         let usage = workload_group.mem_stat.get_memory_usage();
         let max_memory_usage = workload_group.max_memory_usage.load(Ordering::Relaxed);
@@ -162,7 +217,7 @@ impl MemorySettings {
         })
     }
 
-    pub fn check_spill_remain(&self) -> isize {
+    pub fn check_spill_remain(&self) -> Option<isize> {
         [
             self.check_global(),
             self.check_workload_group(),
@@ -171,7 +226,6 @@ impl MemorySettings {
         .into_iter()
         .flatten()
         .reduce(|a, b| a.min(b))
-        .unwrap_or(1)
     }
 }
 
@@ -186,6 +240,7 @@ mod tests {
     impl Default for MemorySettings {
         fn default() -> Self {
             Self {
+                enable_group_spill: true,
                 max_memory_usage: usize::MAX,
                 enable_global_level_spill: false,
                 global_memory_tracking: create_static_mem_stat(0),

src/query/service/src/pipelines/memory_settings.rs

@@ -12,7 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use databend_common_base::runtime::GLOBAL_MEM_STAT;
 use databend_common_catalog::table_context::TableContext;
 use databend_common_exception::Result;
 use databend_common_pipeline_transforms::MemorySettings;
@@ -33,148 +32,128 @@ pub trait MemorySettingsExt: Sized {
 impl MemorySettingsExt for MemorySettings {
     fn from_join_settings(ctx: &QueryContext) -> Result<Self> {
         let settings = ctx.get_settings();
+        let mut builder = MemorySettings::builder();
 
         if settings.get_force_join_data_spill()? {
-            return Ok(MemorySettings::always_spill(0));
+            return Ok(builder.with_max_memory_usage(0).build());
         }
 
-        let mut enable_global_level_spill = false;
-        let mut max_memory_usage = settings.get_max_memory_usage()? as usize;
-
+        let max_memory_usage = settings.get_max_memory_usage()? as usize;
         let max_memory_ratio = settings.get_join_spilling_memory_ratio()?;
         if max_memory_usage != 0 && max_memory_ratio != 0 {
-            enable_global_level_spill = true;
-            let max_memory_ratio = (max_memory_ratio as f64 / 100_f64).min(1_f64);
-            max_memory_usage = (max_memory_usage as f64 * max_memory_ratio) as usize;
+            let ratio = (max_memory_ratio as f64 / 100.0).min(1.0);
+            let global_limit = (max_memory_usage as f64 * ratio) as usize;
+            builder = builder.with_max_memory_usage(global_limit);
         }
 
         let max_query_memory_usage = settings.get_max_query_memory_usage()? as usize;
-        let enable_query_level_spill = match settings.get_query_out_of_memory_behavior()? {
-            OutofMemoryBehavior::Throw => false,
-            OutofMemoryBehavior::Spilling => max_query_memory_usage != 0,
-        };
-
-        Ok(MemorySettings {
-            max_memory_usage,
-            max_query_memory_usage,
-            enable_query_level_spill,
-            enable_global_level_spill,
-            spill_unit_size: 0,
-            query_memory_tracking: ctx.get_query_memory_tracking(),
-            global_memory_tracking: &GLOBAL_MEM_STAT,
-        })
+        let out_of_memory_behavior = settings.get_query_out_of_memory_behavior()?;
+        if matches!(out_of_memory_behavior, OutofMemoryBehavior::Spilling)
+            && max_query_memory_usage != 0
+        {
+            builder = builder.with_max_query_memory_usage(
+                max_query_memory_usage,
+                ctx.get_query_memory_tracking(),
+            );
+        }
+
+        Ok(builder.build())
     }
 
     fn from_sort_settings(ctx: &QueryContext) -> Result<Self> {
+        let settings = ctx.get_settings();
+        let mut builder = MemorySettings::builder()
+            .with_spill_unit_size(settings.get_sort_spilling_batch_bytes()?);
+
         if !ctx.get_enable_sort_spill() {
-            return Ok(MemorySettings::disable_spill());
+            return Ok(builder.with_workload_group(false).build());
         }
 
-        let settings = ctx.get_settings();
-
         if settings.get_force_sort_data_spill()? {
-            let spilling_batch_bytes = settings.get_sort_spilling_batch_bytes()?;
-            return Ok(MemorySettings::always_spill(spilling_batch_bytes));
+            return Ok(builder.with_max_memory_usage(0).build());
         }
 
-        let mut enable_global_level_spill = false;
-        let mut max_memory_usage = settings.get_max_memory_usage()? as usize;
-
+        let max_memory_usage = settings.get_max_memory_usage()? as usize;
         let max_memory_ratio = settings.get_sort_spilling_memory_ratio()?;
         if max_memory_usage != 0 && max_memory_ratio != 0 {
-            enable_global_level_spill = true;
-            let max_memory_ratio = (max_memory_ratio as f64 / 100_f64).min(1_f64);
-            max_memory_usage = (max_memory_usage as f64 * max_memory_ratio) as usize;
+            let ratio = (max_memory_ratio as f64 / 100.0).min(1.0);
+            let global_limit = (max_memory_usage as f64 * ratio) as usize;
+            builder = builder.with_max_memory_usage(global_limit);
         }
 
         let max_query_memory_usage = settings.get_max_query_memory_usage()? as usize;
-        let enable_query_level_spill = match settings.get_query_out_of_memory_behavior()? {
-            OutofMemoryBehavior::Throw => false,
-            OutofMemoryBehavior::Spilling => max_query_memory_usage != 0,
-        };
-
-        let spilling_batch_bytes = settings.get_sort_spilling_batch_bytes()?;
-
-        Ok(MemorySettings {
-            max_memory_usage,
-            max_query_memory_usage,
-            enable_query_level_spill,
-            enable_global_level_spill,
-            query_memory_tracking: ctx.get_query_memory_tracking(),
-            spill_unit_size: spilling_batch_bytes,
-            global_memory_tracking: &GLOBAL_MEM_STAT,
-        })
+        let out_of_memory_behavior = settings.get_query_out_of_memory_behavior()?;
+        if matches!(out_of_memory_behavior, OutofMemoryBehavior::Spilling)
+            && max_query_memory_usage != 0
+        {
+            builder = builder.with_max_query_memory_usage(
+                max_query_memory_usage,
+                ctx.get_query_memory_tracking(),
+            );
+        }
+
+        Ok(builder.build())
     }
 
     fn from_window_settings(ctx: &QueryContext) -> Result<Self> {
         let settings = ctx.get_settings();
+        let spill_unit_size = settings.get_window_spill_unit_size_mb()? * 1024 * 1024;
+        let mut builder = MemorySettings::builder().with_spill_unit_size(spill_unit_size);
 
         if settings.get_force_window_data_spill()? {
-            let spill_unit_size = settings.get_window_spill_unit_size_mb()? * 1024 * 1024;
-            return Ok(MemorySettings::always_spill(spill_unit_size));
+            return Ok(builder.with_max_memory_usage(0).build());
         }
 
-        let mut enable_global_level_spill = false;
-        let mut max_memory_usage = settings.get_max_memory_usage()? as usize;
-
+        let max_memory_usage = settings.get_max_memory_usage()? as usize;
         let max_memory_ratio = settings.get_window_partition_spilling_memory_ratio()?;
         if max_memory_usage != 0 && max_memory_ratio != 0 {
-            enable_global_level_spill = true;
-            let max_memory_ratio = (max_memory_ratio as f64 / 100_f64).min(1_f64);
-            max_memory_usage = (max_memory_usage as f64 * max_memory_ratio) as usize;
+            let ratio = (max_memory_ratio as f64 / 100.0).min(1.0);
+            let global_limit = (max_memory_usage as f64 * ratio) as usize;
+            builder = builder.with_max_memory_usage(global_limit);
         }
 
         let max_query_memory_usage = settings.get_max_query_memory_usage()? as usize;
-        let enable_query_level_spill = match settings.get_query_out_of_memory_behavior()? {
-            OutofMemoryBehavior::Throw => false,
-            OutofMemoryBehavior::Spilling => max_query_memory_usage != 0,
-        };
-
-        let spill_unit_size = settings.get_window_spill_unit_size_mb()? * 1024 * 1024;
+        let out_of_memory_behavior = settings.get_query_out_of_memory_behavior()?;
+        if matches!(out_of_memory_behavior, OutofMemoryBehavior::Spilling)
+            && max_query_memory_usage != 0
+        {
+            builder = builder.with_max_query_memory_usage(
+                max_query_memory_usage,
+                ctx.get_query_memory_tracking(),
+            );
+        }
 
-        Ok(MemorySettings {
-            spill_unit_size,
-            max_memory_usage,
-            max_query_memory_usage,
-            enable_query_level_spill,
-            enable_global_level_spill,
-            query_memory_tracking: ctx.get_query_memory_tracking(),
-            global_memory_tracking: &GLOBAL_MEM_STAT,
-        })
+        Ok(builder.build())
     }
 
     fn from_aggregate_settings(ctx: &QueryContext) -> Result<Self> {
         let settings = ctx.get_settings();
+        let mut builder = MemorySettings::builder();
 
         if settings.get_force_aggregate_data_spill()? {
-            return Ok(MemorySettings::always_spill(0));
+            return Ok(builder.with_max_memory_usage(0).build());
         }
 
-        let mut enable_global_level_spill = false;
-        let mut max_memory_usage = settings.get_max_memory_usage()? as usize;
-
+        let max_memory_usage = settings.get_max_memory_usage()? as usize;
         let max_memory_ratio = settings.get_aggregate_spilling_memory_ratio()?;
         if max_memory_usage != 0 && max_memory_ratio != 0 {
-            enable_global_level_spill = true;
-            let max_memory_ratio = (max_memory_ratio as f64 / 100_f64).min(1_f64);
-            max_memory_usage = (max_memory_usage as f64 * max_memory_ratio) as usize;
+            let ratio = (max_memory_ratio as f64 / 100.0).min(1.0);
+            let global_limit = (max_memory_usage as f64 * ratio) as usize;
+            builder = builder.with_max_memory_usage(global_limit);
         }
 
         let max_query_memory_usage = settings.get_max_query_memory_usage()? as usize;
-        let enable_query_level_spill = match settings.get_query_out_of_memory_behavior()? {
-            OutofMemoryBehavior::Throw => false,
-            OutofMemoryBehavior::Spilling => max_query_memory_usage != 0,
-        };
-
-        Ok(MemorySettings {
-            max_memory_usage,
-            max_query_memory_usage,
-            enable_query_level_spill,
-            enable_global_level_spill,
-            spill_unit_size: 0,
-            query_memory_tracking: ctx.get_query_memory_tracking(),
-            global_memory_tracking: &GLOBAL_MEM_STAT,
-        })
+        let out_of_memory_behavior = settings.get_query_out_of_memory_behavior()?;
+        if matches!(out_of_memory_behavior, OutofMemoryBehavior::Spilling)
+            && max_query_memory_usage != 0
+        {
+            builder = builder.with_max_query_memory_usage(
+                max_query_memory_usage,
+                ctx.get_query_memory_tracking(),
+            );
+        }
+
+        Ok(builder.build())
     }
 }
 
@@ -201,7 +180,6 @@ mod tests {
         let memory_settings = MemorySettings::from_join_settings(&ctx)?;
 
         assert!(!memory_settings.enable_global_level_spill);
-        assert_eq!(memory_settings.max_memory_usage, 0);
         assert!(memory_settings.enable_query_level_spill);
         Ok(())
     }