refactor(query): reduce memory usage during aggregate spill

src/common/base/src/runtime/memory/mem_stat.rs

@@ -40,14 +40,14 @@ pub struct MemStat {
     name: Option<String>,
 
     pub(crate) used: AtomicI64,
-    pub(crate) peek_used: AtomicI64,
+    pub(crate) peak_used: AtomicI64,
 
     /// The limit of max used memory for this tracker.
     ///
     /// Set to 0 to disable the limit.
     limit: AtomicI64,
 
-    parent_memory_stat: Vec<Arc<MemStat>>,
+    parent_memory_stat: Option<Arc<MemStat>>,
 }
 
 impl MemStat {
@@ -56,17 +56,17 @@ impl MemStat {
             id: 0,
             name: None,
             used: AtomicI64::new(0),
-            peek_used: AtomicI64::new(0),
+            peak_used: AtomicI64::new(0),
             limit: AtomicI64::new(0),
-            parent_memory_stat: vec![],
+            parent_memory_stat: None,
         }
     }
 
     pub fn create(name: String) -> Arc<MemStat> {
-        MemStat::create_child(name, vec![])
+        MemStat::create_child(name, None)
     }
 
-    pub fn create_child(name: String, parent_memory_stat: Vec<Arc<MemStat>>) -> Arc<MemStat> {
+    pub fn create_child(name: String, parent_memory_stat: Option<Arc<MemStat>>) -> Arc<MemStat> {
         let id = match GlobalSequence::next() {
             0 => GlobalSequence::next(),
             id => id,
@@ -76,16 +76,12 @@ impl MemStat {
             id,
             name: Some(name),
             used: AtomicI64::new(0),
-            peek_used: AtomicI64::new(0),
+            peak_used: AtomicI64::new(0),
             limit: AtomicI64::new(0),
             parent_memory_stat,
         })
     }
 
-    pub fn get_parent_memory_stat(&self) -> Vec<Arc<MemStat>> {
-        self.parent_memory_stat.clone()
-    }
-
     pub fn set_limit(&self, mut size: i64) {
         // It may cause the process unable to run if memory limit is too low.
         if size > 0 && size < MINIMUM_MEMORY_LIMIT {
@@ -107,19 +103,15 @@ impl MemStat {
         let mut used = self.used.fetch_add(batch_memory_used, Ordering::Relaxed);
 
         used += batch_memory_used;
-        self.peek_used.fetch_max(used, Ordering::Relaxed);
+        self.peak_used.fetch_max(used, Ordering::Relaxed);
 
-        for (idx, parent_memory_stat) in self.parent_memory_stat.iter().enumerate() {
+        if let Some(parent_memory_stat) = self.parent_memory_stat.as_ref() {
             if let Err(cause) = parent_memory_stat
                 .record_memory::<NEED_ROLLBACK>(batch_memory_used, current_memory_alloc)
             {
                 if NEED_ROLLBACK {
                     // We only roll back the memory that alloc failed
                     self.used.fetch_sub(current_memory_alloc, Ordering::Relaxed);
-
-                    for index in 0..idx {
-                        self.parent_memory_stat[index].rollback(current_memory_alloc);
-                    }
                 }
 
                 return Err(cause);
@@ -142,8 +134,8 @@ impl MemStat {
     pub fn rollback(&self, memory_usage: i64) {
         self.used.fetch_sub(memory_usage, Ordering::Relaxed);
 
-        for parent_memory_stat in &self.parent_memory_stat {
-            parent_memory_stat.rollback(memory_usage)
+        if let Some(parent_memory_stat) = &self.parent_memory_stat {
+            parent_memory_stat.rollback(memory_usage);
         }
     }
 
@@ -171,7 +163,7 @@ impl MemStat {
 
     #[inline]
     pub fn get_peek_memory_usage(&self) -> i64 {
-        self.peek_used.load(Ordering::Relaxed)
+        self.peak_used.load(Ordering::Relaxed)
     }
 }
 
@@ -268,7 +260,7 @@ mod tests {
     fn test_multiple_level_mem_stat() -> Result<()> {
         let mem_stat = MemStat::create("TEST".to_string());
         let child_mem_stat =
-            MemStat::create_child("TEST_CHILD".to_string(), vec![mem_stat.clone()]);
+            MemStat::create_child("TEST_CHILD".to_string(), Some(mem_stat.clone()));
 
         mem_stat.record_memory::<false>(1, 1).unwrap();
         mem_stat.record_memory::<false>(2, 2).unwrap();
@@ -292,7 +284,7 @@ mod tests {
         let mem_stat = MemStat::create("TEST".to_string());
         mem_stat.set_limit(MINIMUM_MEMORY_LIMIT * 2);
         let child_mem_stat =
-            MemStat::create_child("TEST_CHILD".to_string(), vec![mem_stat.clone()]);
+            MemStat::create_child("TEST_CHILD".to_string(), Some(mem_stat.clone()));
         child_mem_stat.set_limit(MINIMUM_MEMORY_LIMIT);
 
         mem_stat.record_memory::<false>(1, 1).unwrap();
@@ -322,7 +314,7 @@ mod tests {
         let mem_stat = MemStat::create("TEST".to_string());
         mem_stat.set_limit(MINIMUM_MEMORY_LIMIT);
         let child_mem_stat =
-            MemStat::create_child("TEST_CHILD".to_string(), vec![mem_stat.clone()]);
+            MemStat::create_child("TEST_CHILD".to_string(), Some(mem_stat.clone()));
         child_mem_stat.set_limit(MINIMUM_MEMORY_LIMIT * 2);
 
         assert!(child_mem_stat
@@ -335,7 +327,7 @@ mod tests {
         let mem_stat = MemStat::create("TEST".to_string());
         mem_stat.set_limit(MINIMUM_MEMORY_LIMIT * 2);
         let child_mem_stat =
-            MemStat::create_child("TEST_CHILD".to_string(), vec![mem_stat.clone()]);
+            MemStat::create_child("TEST_CHILD".to_string(), Some(mem_stat.clone()));
         child_mem_stat.set_limit(MINIMUM_MEMORY_LIMIT);
 
         assert!(child_mem_stat

src/common/base/src/runtime/memory/stat_buffer_global.rs

@@ -90,7 +90,7 @@ impl GlobalStatBuffer {
                 .used
                 .fetch_add(memory_usage, Ordering::Relaxed);
             self.global_mem_stat
-                .peek_used
+                .peak_used
                 .fetch_max(used + memory_usage, Ordering::Relaxed);
             return Ok(());
         }
@@ -126,7 +126,7 @@ impl GlobalStatBuffer {
                 .used
                 .fetch_add(memory_usage, Ordering::Relaxed);
             self.global_mem_stat
-                .peek_used
+                .peak_used
                 .fetch_max(used + memory_usage, Ordering::Relaxed);
             return;
         }

src/common/base/src/runtime/memory/stat_buffer_mem_stat.rs

@@ -93,7 +93,7 @@ impl MemStatBuffer {
         if self.destroyed_thread_local_macro {
             let used = mem_stat.used.fetch_add(usage, Ordering::Relaxed);
             mem_stat
-                .peek_used
+                .peak_used
                 .fetch_max(used + usage, Ordering::Relaxed);
             return Ok(());
         }
@@ -134,7 +134,7 @@ impl MemStatBuffer {
         if self.destroyed_thread_local_macro {
             let used = mem_stat.used.fetch_add(memory_usage, Ordering::Relaxed);
             mem_stat
-                .peek_used
+                .peak_used
                 .fetch_max(used + memory_usage, Ordering::Relaxed);
             return;
         }

src/query/expression/src/aggregate/partitioned_payload.rs

@@ -16,6 +16,8 @@ use std::sync::Arc;
 
 use bumpalo::Bump;
 use itertools::Itertools;
+use serde::Deserializer;
+use serde::Serializer;
 
 use super::payload::Payload;
 use super::probe_state::ProbeState;
@@ -50,6 +52,18 @@ pub struct PartitionedPayload {
 unsafe impl Send for PartitionedPayload {}
 unsafe impl Sync for PartitionedPayload {}
 
+impl serde::Serialize for PartitionedPayload {
+    fn serialize<S: Serializer>(&self, _: S) -> Result<S::Ok, S::Error> {
+        unreachable!("PartitionedPayload must not be exchanged between multiple nodes.")
+    }
+}
+
+impl<'de> serde::Deserialize<'de> for PartitionedPayload {
+    fn deserialize<D: Deserializer<'de>>(_: D) -> Result<Self, D::Error> {
+        unreachable!("PartitionedPayload must not be exchanged between multiple nodes.")
+    }
+}
+
 impl PartitionedPayload {
     pub fn new(
         group_types: Vec<DataType>,

src/query/expression/src/aggregate/payload.rs

@@ -385,7 +385,11 @@ impl Payload {
         );
     }
 
-    pub fn scatter(&self, state: &mut PayloadFlushState, partition_count: usize) -> bool {
+    pub fn scatter_with_seed<const SEED: u64>(
+        &self,
+        state: &mut PayloadFlushState,
+        partitions: usize,
+    ) -> bool {
         if state.flush_page >= self.pages.len() {
             return false;
         }
@@ -397,23 +401,27 @@ impl Payload {
             state.flush_page += 1;
             state.flush_page_row = 0;
             state.row_count = 0;
-            return self.scatter(state, partition_count);
+            return self.scatter_with_seed::<SEED>(state, partitions);
         }
 
         let end = (state.flush_page_row + BATCH_SIZE).min(page.rows);
         let rows = end - state.flush_page_row;
         state.row_count = rows;
 
-        state.probe_state.reset_partitions(partition_count);
+        state.probe_state.reset_partitions(partitions);
+
+        let mods: StrengthReducedU64 = StrengthReducedU64::new(partitions as u64);
 
-        let mods: StrengthReducedU64 = StrengthReducedU64::new(partition_count as u64);
         for idx in 0..rows {
             state.addresses[idx] = self.data_ptr(page, idx + state.flush_page_row);
 
-            let hash = unsafe { read::<u64>(state.addresses[idx].add(self.hash_offset) as _) };
+            let mut hash = unsafe { read::<u64>(state.addresses[idx].add(self.hash_offset) as _) };
 
-            let partition_idx = (hash % mods) as usize;
+            if SEED != 0 {
+                hash = Self::combine_hash(hash, SEED);
+            }
 
+            let partition_idx = (hash % mods) as usize;
             let sel = &mut state.probe_state.partition_entries[partition_idx];
             sel[state.probe_state.partition_count[partition_idx]] = idx;
             state.probe_state.partition_count[partition_idx] += 1;
@@ -422,6 +430,10 @@ impl Payload {
         true
     }
 
+    pub fn scatter(&self, state: &mut PayloadFlushState, partitions: usize) -> bool {
+        self.scatter_with_seed::<0>(state, partitions)
+    }
+
     pub fn empty_block(&self, fake_rows: Option<usize>) -> DataBlock {
         let fake_rows = fake_rows.unwrap_or(0);
         let columns = (0..self.aggrs.len())
@@ -434,6 +446,18 @@ impl Payload {
             .collect_vec();
         DataBlock::new_from_columns(columns)
     }
+
+    #[allow(unused_parens)]
+    fn combine_hash(hash: u64, seed: u64) -> u64 {
+        static KMUL: u64 = 0x9ddfea08eb382d69;
+
+        let mut a = (seed ^ hash).wrapping_mul(KMUL);
+        a ^= (a >> 47);
+
+        let mut b = (hash ^ a).wrapping_mul(KMUL);
+        b ^= (b >> 47);
+        b.wrapping_mul(KMUL)
+    }
 }
 
 impl Drop for Payload {

src/query/expression/src/aggregate/payload_flush.rs

@@ -164,22 +164,6 @@ impl Payload {
         Ok(Some(DataBlock::new_from_columns(cols)))
     }
 
-    pub fn group_by_flush_all(&self) -> Result<DataBlock> {
-        let mut state = PayloadFlushState::default();
-        let mut blocks = vec![];
-
-        while self.flush(&mut state) {
-            let cols = state.take_group_columns();
-            blocks.push(DataBlock::new_from_columns(cols));
-        }
-
-        if blocks.is_empty() {
-            return Ok(self.empty_block(None));
-        }
-
-        DataBlock::concat(&blocks)
-    }
-
     pub fn flush(&self, state: &mut PayloadFlushState) -> bool {
         if state.flush_page >= self.pages.len() {
             return false;

src/query/pipeline/core/src/lib.rs

@@ -15,6 +15,8 @@
 #![feature(once_cell_try)]
 #![feature(variant_count)]
 #![feature(associated_type_defaults)]
+#![feature(adt_const_params)]
+#![feature(let_chains)]
 #![allow(clippy::arc_with_non_send_sync)]
 #![allow(clippy::useless_asref)]
 

src/query/pipeline/core/src/pipeline.rs

@@ -20,6 +20,7 @@ use std::sync::atomic::Ordering;
 use std::sync::Arc;
 use std::time::Instant;
 
+use databend_common_base::base::tokio::sync::Barrier;
 use databend_common_base::runtime::defer;
 use databend_common_base::runtime::drop_guard;
 use databend_common_exception::ErrorCode;
@@ -456,11 +457,18 @@ impl Pipeline {
             let input_len = pipe.output_length;
             let mut items = Vec::with_capacity(input_len);
 
-            for _index in 0..input_len {
+            let barrier = Arc::new(Barrier::new(input_len));
+            for index in 0..input_len {
                 let input = InputPort::create();
                 let outputs: Vec<_> = (0..n).map(|_| OutputPort::create()).collect();
                 items.push(PipeItem::create(
-                    PartitionProcessor::create(input.clone(), outputs.clone(), exchange.clone()),
+                    PartitionProcessor::create(
+                        input.clone(),
+                        outputs.clone(),
+                        exchange.clone(),
+                        index,
+                        barrier.clone(),
+                    ),
                     vec![input],
                     outputs,
                 ));
@@ -481,7 +489,7 @@ impl Pipeline {
                 let output = OutputPort::create();
                 let inputs: Vec<_> = (0..input_len).map(|_| InputPort::create()).collect();
                 items.push(PipeItem::create(
-                    MergePartitionProcessor::create(
+                    MergePartitionProcessor::<T>::create(
                         inputs.clone(),
                         output.clone(),
                         exchange.clone(),

src/query/pipeline/core/src/processors/mod.rs

@@ -39,5 +39,6 @@ pub use resize_processor::create_resize_item;
 pub use resize_processor::ResizeProcessor;
 pub use shuffle_processor::Exchange;
 pub use shuffle_processor::MergePartitionProcessor;
+pub use shuffle_processor::MultiwayStrategy;
 pub use shuffle_processor::PartitionProcessor;
 pub use shuffle_processor::ShuffleProcessor;

src/query/pipeline/core/src/processors/processor.rs

@@ -80,6 +80,22 @@ pub trait Processor: Send {
         Err(ErrorCode::Unimplemented("Unimplemented async_process."))
     }
 
+    fn prepare_spill_payload(&mut self) -> Result<bool> {
+        Err(ErrorCode::Unimplemented(
+            "Unimplemented prepare_spill_payload",
+        ))
+    }
+
+    async fn flush_spill_payload(&mut self) -> Result<bool> {
+        Err(ErrorCode::Unimplemented(
+            "Unimplemented flush_spill_payload",
+        ))
+    }
+
+    fn configure_peer_nodes(&mut self, _nodes: &[String]) {
+        // do nothing by default
+    }
+
     fn details_status(&self) -> Option<String> {
         None
     }
@@ -198,6 +214,10 @@ impl ProcessorPtr {
         .boxed()
     }
 
+    pub fn configure_peer_nodes(&self, nodes: &[String]) {
+        unsafe { (*self.inner.get()).configure_peer_nodes(nodes) }
+    }
+
     /// # Safety
     pub unsafe fn details_status(&self) -> Option<String> {
         (*self.inner.get()).details_status()