memory

Author: forsaken628

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

@@ -75,12 +75,10 @@ impl Debug for 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>>,
 
@@ -89,7 +87,6 @@ pub struct MemorySettingsBuilder {
 
 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
     }
@@ -99,7 +96,6 @@ impl MemorySettingsBuilder {
         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
@@ -118,12 +114,15 @@ impl MemorySettingsBuilder {
     pub fn build(self) -> MemorySettings {
         MemorySettings {
             enable_group_spill: self.enable_group_spill,
+
+            enable_global_level_spill: self.max_memory_usage.is_some(),
             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,
+
+            enable_query_level_spill: self.max_query_memory_usage.is_some(),
             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),
         }
     }
@@ -132,12 +131,10 @@ impl MemorySettingsBuilder {
 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,
 
@@ -146,29 +143,25 @@ impl MemorySettings {
     }
 
     pub fn check_spill(&self) -> bool {
-        if self.enable_global_level_spill
-            && self.global_memory_tracking.get_memory_usage() >= self.max_memory_usage
+        if let Some(remain) = self.check_global()
+            && remain <= 0
         {
             return true;
         }
 
-        if self.enable_group_spill
-            && let Some(workload_group) = ThreadTracker::workload_group()
+        if let Some(remain) = self.check_workload_group()
+            && remain <= 0
         {
-            let workload_group_memory_usage = workload_group.mem_stat.get_memory_usage();
-            let max_memory_usage = workload_group.max_memory_usage.load(Ordering::Relaxed);
-
-            if max_memory_usage != 0 && workload_group_memory_usage >= max_memory_usage {
-                return true;
-            }
+            return true;
         }
 
-        let Some(query_memory_tracking) = self.query_memory_tracking.as_ref() else {
-            return false;
-        };
-
-        self.enable_query_level_spill
-            && query_memory_tracking.get_memory_usage() >= self.max_query_memory_usage
+        if let Some(remain) = self.check_query()
+            && remain <= 0
+        {
+            true
+        } else {
+            false
+        }
     }
 
     fn check_global(&self) -> Option<isize> {
@@ -207,8 +200,7 @@ impl MemorySettings {
             return None;
         }
 
-        let query_memory_tracking = self.query_memory_tracking.as_ref()?;
-        let usage = query_memory_tracking.get_memory_usage();
+        let usage = self.query_memory_tracking.as_ref()?.get_memory_usage();
 
         Some(if usage >= self.max_query_memory_usage {
             -((usage - self.max_query_memory_usage) as isize)

src/query/pipeline/transforms/src/processors/traits/spill.rs

@@ -16,6 +16,8 @@ use databend_common_exception::Result;
 use databend_common_expression::DataBlock;
 use databend_storages_common_cache::TempPath;
 
+use crate::MemorySettings;
+
 #[derive(Debug, Clone, Hash, PartialEq, Eq)]
 pub enum Location {
     Remote(String),
@@ -32,6 +34,14 @@ impl Location {
     }
 }
 
+#[async_trait::async_trait]
+pub trait SortSpiller: Clone + Send + Sync + 'static {
+    async fn spill(&self, data_block: DataBlock) -> Result<Location>;
+    async fn restore(&self, location: &Location) -> Result<DataBlock>;
+    fn remove_local_file(&self, local: &TempPath);
+    fn memory_settings(&self) -> &MemorySettings;
+}
+
 #[async_trait::async_trait]
 pub trait DataBlockSpill: Clone + Send + Sync + 'static {
     async fn spill(&self, data_block: DataBlock) -> Result<Location> {

src/query/pipeline/transforms/src/processors/transforms/sorts/sort_collect.rs

@@ -35,18 +35,17 @@ use super::TransformSortMergeLimit;
 use super::core::RowConverter;
 use super::core::Rows;
 use super::core::algorithm::SortAlgorithm;
-use crate::MemorySettings;
-use crate::traits::DataBlockSpill;
+use crate::traits::SortSpiller;
 
 #[allow(clippy::large_enum_variant)]
-enum Inner<A: SortAlgorithm, S: DataBlockSpill> {
+enum Inner<A: SortAlgorithm, S: SortSpiller> {
     Collect(Vec<DataBlock>),
     Limit(TransformSortMergeLimit<A::Rows>),
     Spill(Vec<DataBlock>, SortSpill<A, S>),
     None,
 }
 
-pub struct TransformSortCollect<A: SortAlgorithm, C, S: DataBlockSpill> {
+pub struct TransformSortCollect<A: SortAlgorithm, C, S: SortSpiller> {
     name: &'static str,
     input: Arc<InputPort>,
     output: Arc<OutputPort>,
@@ -68,14 +67,13 @@ pub struct TransformSortCollect<A: SortAlgorithm, C, S: DataBlockSpill> {
     aborting: AtomicBool,
 
     enable_restore_prefetch: bool,
-    memory_settings: MemorySettings,
 }
 
 impl<A, C, S> TransformSortCollect<A, C, S>
 where
     A: SortAlgorithm,
     C: RowConverter<A::Rows>,
-    S: DataBlockSpill,
+    S: SortSpiller,
 {
     pub fn new(
         input: Arc<InputPort>,
@@ -87,7 +85,6 @@ where
         sort_limit: bool,
         order_col_generated: bool,
         enable_restore_prefetch: bool,
-        memory_settings: MemorySettings,
     ) -> Result<Self> {
         let row_converter = C::create(&sort_desc, base.schema.clone())?;
         let (name, inner) = match base.limit {
@@ -108,7 +105,6 @@ where
             base,
             inner,
             aborting: AtomicBool::new(false),
-            memory_settings,
             max_block_size,
             default_num_merge,
             enable_restore_prefetch,
@@ -184,7 +180,7 @@ where
         SortSpillParams::determine(
             bytes,
             rows,
-            ByteSize(self.memory_settings.spill_unit_size as _),
+            ByteSize(self.base.spiller.memory_settings().spill_unit_size as _),
             self.max_block_size,
             self.enable_restore_prefetch,
         )
@@ -217,22 +213,23 @@ where
     }
 
     fn check_spill(&self) -> bool {
+        let memory_settings = self.base.spiller.memory_settings();
         match &self.inner {
             Inner::Limit(limit_sort) => {
-                self.memory_settings.check_spill()
+                memory_settings.check_spill()
                     && limit_sort.num_bytes()
-                        >= ByteSize(self.memory_settings.spill_unit_size as _) * 2_u64
+                        >= ByteSize(memory_settings.spill_unit_size as _) * 2_u64
             }
             Inner::Collect(input_data) => {
-                self.memory_settings.check_spill()
+                memory_settings.check_spill()
                     && input_data.iter().map(|b| b.memory_size()).sum::<usize>()
-                        >= self.memory_settings.spill_unit_size * 2
+                        >= memory_settings.spill_unit_size * 2
             }
             Inner::Spill(input_data, sort_spill) => {
-                let rows = input_data.in_memory_rows();
-                let params = sort_spill.params();
-                self.memory_settings.check_spill() && rows >= params.batch_rows * 2
-                    || input_data.in_memory_rows() >= params.max_rows()
+                input_data.in_memory_rows() >= sort_spill.params().batch_rows * 2 && {
+                    let remain = memory_settings.check_spill_remain().unwrap();
+                    remain < memory_settings.spill_unit_size as isize * 2
+                }
             }
             _ => unreachable!(),
         }
@@ -257,7 +254,7 @@ where
     A: SortAlgorithm + 'static,
     A::Rows: 'static,
     C: RowConverter<A::Rows> + Send + 'static,
-    S: DataBlockSpill,
+    S: SortSpiller,
 {
     fn name(&self) -> String {
         self.name.to_string()

src/query/pipeline/transforms/src/processors/transforms/sorts/sort_local_merge.rs

@@ -41,8 +41,7 @@ use super::core::RowConverter;
 use super::core::Rows;
 use super::core::algorithm::SortAlgorithm;
 use super::create_memory_merger;
-use crate::MemorySettings;
-use crate::traits::DataBlockSpill;
+use crate::traits::SortSpiller;
 
 #[derive(Debug)]
 enum State {
@@ -55,14 +54,14 @@ enum State {
 }
 
 #[allow(clippy::large_enum_variant)]
-enum Inner<A: SortAlgorithm, S: DataBlockSpill> {
+enum Inner<A: SortAlgorithm, S: SortSpiller> {
     Collect(Vec<DataBlock>),
     Limit(TransformSortMergeLimit<A::Rows>),
     Memory(MemoryMerger<A>),
     Spill(Vec<DataBlock>, SortSpill<A, S>),
 }
 
-pub struct TransformSort<A: SortAlgorithm, C, S: DataBlockSpill> {
+pub struct TransformSort<A: SortAlgorithm, C, S: SortSpiller> {
     name: &'static str,
     input: Arc<InputPort>,
     output: Arc<OutputPort>,
@@ -87,14 +86,13 @@ pub struct TransformSort<A: SortAlgorithm, C, S: DataBlockSpill> {
 
     max_block_size: usize,
     enable_restore_prefetch: bool,
-    memory_settings: MemorySettings,
 }
 
 impl<A, C, S> TransformSort<A, C, S>
 where
     A: SortAlgorithm,
     C: RowConverter<A::Rows>,
-    S: DataBlockSpill,
+    S: SortSpiller,
 {
     #[allow(clippy::too_many_arguments)]
     pub fn new(
@@ -108,7 +106,6 @@ where
         output_order_col: bool,
         order_col_generated: bool,
         enable_restore_prefetch: bool,
-        memory_settings: MemorySettings,
     ) -> Result<Self> {
         assert!(max_block_size > 0);
         let sort_row_offset = schema.fields().len() - 1;
@@ -142,7 +139,6 @@ where
             max_block_size,
             aborting: AtomicBool::new(false),
             enable_restore_prefetch,
-            memory_settings,
         })
     }
 
@@ -198,7 +194,7 @@ where
         SortSpillParams::determine(
             bytes,
             rows,
-            ByteSize(self.memory_settings.spill_unit_size as _),
+            ByteSize(self.base.spiller.memory_settings().spill_unit_size as _),
             self.max_block_size,
             self.enable_restore_prefetch,
         )
@@ -276,32 +272,25 @@ where
         self.output.push_data(Ok(block));
     }
 
-    fn input_rows(&self) -> (usize, usize) {
-        match &self.inner {
-            Inner::Collect(input_data) | Inner::Spill(input_data, _) => {
-                (input_data.len(), input_data.in_memory_rows())
-            }
-            _ => (0, 0),
-        }
-    }
-
     fn check_spill(&self) -> bool {
+        let memory_settings = self.base.spiller.memory_settings();
         match &self.inner {
             Inner::Limit(limit_sort) => {
-                self.memory_settings.check_spill()
+                memory_settings.check_spill()
                     && limit_sort.num_bytes()
-                        >= ByteSize(self.memory_settings.spill_unit_size as _) * 2_u64
+                        >= ByteSize(memory_settings.spill_unit_size as _) * 2_u64
             }
             Inner::Collect(input_data) => {
-                self.memory_settings.check_spill()
+                memory_settings.check_spill()
                     && input_data.iter().map(|b| b.memory_size()).sum::<usize>()
-                        >= self.memory_settings.spill_unit_size * 2
+                        >= memory_settings.spill_unit_size * 2
             }
             Inner::Spill(input_data, sort_spill) => {
-                let rows = input_data.in_memory_rows();
                 let params = sort_spill.params();
-                self.memory_settings.check_spill() && rows >= params.batch_rows * 2
-                    || input_data.in_memory_rows() >= params.max_rows()
+                input_data.in_memory_rows() >= params.batch_rows * 2 && {
+                    let remain = memory_settings.check_spill_remain().unwrap();
+                    remain < memory_settings.spill_unit_size as isize * 2
+                }
             }
             _ => unreachable!(),
         }
@@ -314,7 +303,7 @@ where
     A: SortAlgorithm + 'static,
     A::Rows: 'static,
     C: RowConverter<A::Rows> + Send + 'static,
-    S: DataBlockSpill,
+    S: SortSpiller,
 {
     fn name(&self) -> String {
         self.name.to_string()
@@ -428,12 +417,12 @@ where
                 let finished = self.input.is_finished();
                 self.trans_to_spill()?;
 
-                let (incoming_block, incoming) = self.input_rows();
                 let Inner::Spill(input_data, spill_sort) = &mut self.inner else {
                     unreachable!()
                 };
-
+                let incoming = input_data.in_memory_rows();
                 if incoming > 0 {
+                    let incoming_block = input_data.len();
                     let total_rows = spill_sort.collect_total_rows();
                     log::debug!(incoming_block, incoming_rows = incoming, total_rows, finished; "sort_input_data");
                     spill_sort
@@ -450,8 +439,9 @@ where
                     unreachable!()
                 };
                 assert!(input_data.is_empty());
-                let OutputData { block, finish, .. } =
-                    spill_sort.on_restore(&self.memory_settings).await?;
+                let OutputData { block, finish, .. } = spill_sort
+                    .on_restore(self.base.spiller.memory_settings())
+                    .await?;
                 self.output_data.extend(block);
                 if finish {
                     self.state = State::Finish