From 9583c7d09fc38eaab3ae6735550fad3750d05fa4 Mon Sep 17 00:00:00 2001 From: coldWater Date: Thu, 30 Jul 2026 17:01:38 +0800 Subject: [PATCH] x --- .../transforms/sorts/sort_collect.rs | 57 ++++++++++++++++--- .../transforms/sorts/sort_local_merge.rs | 52 ++++++++++++++--- 2 files changed, 91 insertions(+), 18 deletions(-) diff --git a/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_collect.rs b/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_collect.rs index 47d14605583f2..f6ee9f8c211a7 100644 --- a/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_collect.rs +++ b/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_collect.rs @@ -22,6 +22,7 @@ use databend_common_pipeline::core::Event; use databend_common_pipeline::core::InputPort; use databend_common_pipeline::core::OutputPort; use databend_common_pipeline::core::Processor; +use databend_common_pipeline::core::check_interrupt; use super::Base; use super::MergeSort; @@ -32,6 +33,7 @@ use super::TransformSortMergeLimit; use super::core::RowConverter; use super::core::Rows; use super::core::algorithm::SortAlgorithm; +use super::create_memory_merger; use crate::traits::SortSpiller; #[allow(clippy::large_enum_variant)] @@ -121,7 +123,11 @@ where Ok(()) } - fn collect_trans_to_spill(&mut self, input_data: Vec, no_spill: bool) { + fn collect_trans_to_spill( + &mut self, + input_data: Vec, + no_spill: bool, + ) -> Result { let (num_rows, num_bytes) = input_data .iter() .map(|block| (block.num_rows(), ByteSize(block.memory_size() as _))) @@ -139,19 +145,50 @@ where } else { self.determine_params(num_bytes, num_rows) }; - let spill_sort = SortSpill::new(self.base.clone(), params); - self.inner = Inner::Spill(input_data, spill_sort); + let limit = if !no_spill + && let Some(limit) = self.base.limit + && input_data.len() > 1 + { + limit + } else { + self.inner = Inner::Spill(input_data, SortSpill::new(self.base.clone(), params)); + return Ok(true); + }; + + let mut merger = create_memory_merger::( + input_data, + self.base.sort_row_offset, + Some(limit), + params.batch_rows, + ); + let mut merged = Vec::new(); + while let Some(block) = merger.next_block()? { + check_interrupt()?; + merged.push(block); + } + drop(merger); + + let memory_settings = self.base.spiller.memory_settings(); + if memory_settings.check_spill() { + self.inner = Inner::Spill(merged, SortSpill::new(self.base.clone(), params)); + Ok(true) + } else { + self.inner = Inner::Collect(merged); + Ok(false) + } } - fn trans_to_spill(&mut self, no_spill: bool) -> Result<()> { + fn trans_to_spill(&mut self, no_spill: bool) -> Result { match &mut self.inner { - Inner::Limit(_) => self.limit_trans_to_spill(no_spill), + Inner::Limit(_) => { + self.limit_trans_to_spill(no_spill)?; + Ok(true) + } Inner::Collect(input_data) => { let input_data = std::mem::take(input_data); - self.collect_trans_to_spill(input_data, no_spill); - Ok(()) + self.collect_trans_to_spill(input_data, no_spill) } - Inner::Spill(_, _) => Ok(()), + Inner::Spill(_, _) => Ok(true), Inner::None => unreachable!(), } } @@ -316,7 +353,9 @@ where #[async_backtrace::framed] async fn async_process(&mut self) -> Result<()> { let finished = self.input.is_finished(); - self.trans_to_spill(finished)?; + if !self.trans_to_spill(finished)? { + return Ok(()); + } let Inner::Spill(input_data, spill_sort) = &mut self.inner else { unreachable!() diff --git a/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_local_merge.rs b/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_local_merge.rs index 51891cacf841d..a5bed70788dc0 100644 --- a/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_local_merge.rs +++ b/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_local_merge.rs @@ -24,6 +24,7 @@ use databend_common_pipeline::core::Event; use databend_common_pipeline::core::InputPort; use databend_common_pipeline::core::OutputPort; use databend_common_pipeline::core::Processor; +use databend_common_pipeline::core::check_interrupt; use super::Base; use super::MemoryMerger; @@ -146,7 +147,7 @@ where Ok(()) } - fn collect_trans_to_spill(&mut self, input_data: Vec) { + fn collect_trans_to_spill(&mut self, input_data: Vec) -> Result { let (num_rows, num_bytes) = input_data .iter() .map(|block| (block.num_rows(), ByteSize(block.memory_size() as _))) @@ -154,19 +155,50 @@ where (acc_rows + rows, acc_bytes + bytes) }); let params = self.determine_params(num_bytes, num_rows); - let spill_sort = SortSpill::new(self.base.clone(), params); - self.inner = Inner::Spill(input_data, spill_sort); + + let limit = if let Some(limit) = self.base.limit + && input_data.len() > 1 + { + limit + } else { + self.inner = Inner::Spill(input_data, SortSpill::new(self.base.clone(), params)); + return Ok(true); + }; + + let mut merger = create_memory_merger::( + input_data, + self.base.sort_row_offset, + Some(limit), + params.batch_rows, + ); + let mut merged = Vec::new(); + while let Some(block) = merger.next_block()? { + check_interrupt()?; + merged.push(block.maybe_gc()); + } + drop(merger); + + let memory_settings = self.base.spiller.memory_settings(); + if !memory_settings.check_spill() { + self.inner = Inner::Collect(merged); + Ok(false) + } else { + self.inner = Inner::Spill(merged, SortSpill::new(self.base.clone(), params)); + Ok(true) + } } - fn trans_to_spill(&mut self) -> Result<()> { + fn trans_to_spill(&mut self) -> Result { match &mut self.inner { - Inner::Limit(_) => self.limit_trans_to_spill(), + Inner::Limit(_) => { + self.limit_trans_to_spill()?; + Ok(true) + } Inner::Collect(input_data) => { let input_data = std::mem::take(input_data); - self.collect_trans_to_spill(input_data); - Ok(()) + self.collect_trans_to_spill(input_data) } - Inner::Spill(_, _) => Ok(()), + Inner::Spill(_, _) => Ok(true), Inner::Memory(_) => unreachable!(), } } @@ -401,7 +433,9 @@ where match &self.state { State::Collect => { let finished = self.input.is_finished(); - self.trans_to_spill()?; + if !self.trans_to_spill()? { + return Ok(()); + } let Inner::Spill(input_data, spill_sort) = &mut self.inner else { unreachable!()