build_bounded_merge_sort

forsaken628 · forsaken628 · commit 9cbe20750f76 · 2025-07-09T11:11:58.000+08:00
diff --git a/src/query/service/src/pipelines/builders/builder_sort.rs b/src/query/service/src/pipelines/builders/builder_sort.rs
@@ -19,7 +19,11 @@ use databend_common_exception::Result;
 use databend_common_expression::DataSchemaRef;
 use databend_common_expression::LimitType;
 use databend_common_expression::SortColumnDescription;
+use databend_common_pipeline_core::processors::InputPort;
+use databend_common_pipeline_core::processors::OutputPort;
 use databend_common_pipeline_core::processors::ProcessorPtr;
+use databend_common_pipeline_core::Pipe;
+use databend_common_pipeline_core::PipeItem;
 use databend_common_pipeline_core::Pipeline;
 use databend_common_pipeline_transforms::processors::add_k_way_merge_sort;
 use databend_common_pipeline_transforms::processors::sort::utils::add_order_field;
@@ -38,6 +42,7 @@ use databend_storages_common_cache::TempDirManager;
 
 use crate::pipelines::memory_settings::MemorySettingsExt;
 use crate::pipelines::processors::transforms::add_range_shuffle_route;
+use crate::pipelines::processors::transforms::BoundedMergeSortBuilder;
 use crate::pipelines::processors::transforms::SortInjector;
 use crate::pipelines::processors::transforms::SortRangeExchange;
 use crate::pipelines::processors::transforms::TransformLimit;
@@ -156,7 +161,7 @@ impl PipelineBuilder {
                     self.build_pipeline(&sort.input)?;
                 }
 
-                self.main_pipeline.resize(1, false)
+                builder.build_bounded_merge_sort(&mut self.main_pipeline)
             }
             SortStep::Route => self.main_pipeline.resize(1, false),
         }
@@ -253,7 +258,7 @@ impl SortPipelineBuilder {
         let memory_settings = MemorySettings::from_sort_settings(&self.ctx)?;
         let enable_loser_tree = settings.get_enable_loser_tree_merge_sort()?;
 
-        let builder = TransformSortBuilder::create(
+        let builder = TransformSortBuilder::new(
             self.schema.clone(),
             self.sort_desc.clone(),
             max_block_size,
@@ -325,7 +330,7 @@ impl SortPipelineBuilder {
         let memory_settings = MemorySettings::from_sort_settings(&self.ctx)?;
         let enable_loser_tree = settings.get_enable_loser_tree_merge_sort()?;
 
-        let builder = TransformSortBuilder::create(
+        let builder = TransformSortBuilder::new(
             self.schema.clone(),
             self.sort_desc.clone(),
             max_block_size,
@@ -377,12 +382,10 @@ impl SortPipelineBuilder {
         // Merge sort
         let need_multi_merge = pipeline.output_len() > 1;
         let output_order_col = need_multi_merge || !self.remove_order_col_at_last;
-        debug_assert!(if order_col_generated {
+        debug_assert!(
             // If `order_col_generated`, it means this transform is the last processor in the distributed sort pipeline.
-            !output_order_col
-        } else {
-            true
-        });
+            !order_col_generated || !output_order_col
+        );
 
         let memory_settings = MemorySettings::from_sort_settings(&self.ctx)?;
         let sort_merge_output_schema = match output_order_col {
@@ -414,7 +417,7 @@ impl SortPipelineBuilder {
         };
 
         pipeline.add_transform(|input, output| {
-            let builder = TransformSortBuilder::create(
+            let builder = TransformSortBuilder::new(
                 sort_merge_output_schema.clone(),
                 self.sort_desc.clone(),
                 self.block_size,
@@ -476,4 +479,32 @@ impl SortPipelineBuilder {
     pub fn exchange_injector(&self) -> Arc<dyn ExchangeInjector> {
         Arc::new(SortInjector {})
     }
+
+    pub fn build_bounded_merge_sort(self, pipeline: &mut Pipeline) -> Result<()> {
+        let inputs_port: Vec<_> = (0..pipeline.output_len())
+            .map(|_| InputPort::create())
+            .collect();
+        let output_port = OutputPort::create();
+
+        let processor = ProcessorPtr::create(
+            BoundedMergeSortBuilder::new(
+                inputs_port.clone(),
+                output_port.clone(),
+                self.schema.clone(),
+                self.sort_desc.clone(),
+                self.block_size,
+                self.limit,
+                self.remove_order_col_at_last,
+                self.enable_loser_tree,
+            )
+            .build()?,
+        );
+
+        pipeline.add_pipe(Pipe::create(inputs_port.len(), 1, vec![PipeItem::create(
+            processor,
+            inputs_port,
+            vec![output_port],
+        )]));
+        Ok(())
+    }
 }
diff --git a/src/query/service/src/pipelines/processors/transforms/sort/mod.rs b/src/query/service/src/pipelines/processors/transforms/sort/mod.rs
@@ -39,9 +39,14 @@ mod sort_route;
 mod sort_shuffle;
 mod sort_spill;
 
+pub use merge_sort::*;
 pub use sort_builder::*;
+pub use sort_collect::*;
+pub use sort_combine::*;
 pub use sort_exchange::*;
 pub use sort_exchange_injector::*;
+pub use sort_merge_stream::*;
+pub use sort_restore::*;
 pub use sort_route::*;
 pub use sort_shuffle::*;
 
diff --git a/src/query/service/src/pipelines/processors/transforms/sort/sort_builder.rs b/src/query/service/src/pipelines/processors/transforms/sort/sort_builder.rs
@@ -34,13 +34,7 @@ use databend_common_pipeline_transforms::sort::RowsTypeVisitor;
 use databend_common_pipeline_transforms::AccumulatingTransformer;
 use databend_common_pipeline_transforms::MemorySettings;
 
-use super::merge_sort::TransformSort;
-use super::sort_collect::TransformSortCollect;
-use super::sort_combine::TransformSortCombine;
-use super::sort_restore::TransformSortRestore;
-use super::sort_shuffle::SortSampleState;
-use super::sort_shuffle::TransformSortBoundBroadcast;
-use super::Base;
+use super::*;
 use crate::sessions::QueryContext;
 use crate::spillers::Spiller;
 
@@ -65,7 +59,7 @@ pub struct TransformSortBuilder {
 }
 
 impl TransformSortBuilder {
-    pub fn create(
+    pub fn new(
         schema: DataSchemaRef,
         sort_desc: Arc<[SortColumnDescription]>,
         block_size: usize,
@@ -357,3 +351,83 @@ impl RowsTypeVisitor for Build<'_> {
         }
     }
 }
+
+pub struct BoundedMergeSortBuilder {
+    inputs: Vec<Arc<InputPort>>,
+    output: Arc<OutputPort>,
+    schema: DataSchemaRef,
+    sort_desc: Arc<[SortColumnDescription]>,
+    block_size: usize,
+    limit: Option<usize>,
+    remove_order_col: bool,
+    enable_loser_tree: bool,
+}
+
+impl BoundedMergeSortBuilder {
+    pub fn new(
+        inputs: Vec<Arc<InputPort>>,
+        output: Arc<OutputPort>,
+        schema: DataSchemaRef,
+        sort_desc: Arc<[SortColumnDescription]>,
+        block_size: usize,
+        limit: Option<usize>,
+        remove_order_col: bool,
+        enable_loser_tree: bool,
+    ) -> Self {
+        Self {
+            inputs,
+            output,
+            schema,
+            sort_desc,
+            block_size,
+            limit,
+            remove_order_col,
+            enable_loser_tree,
+        }
+    }
+
+    pub fn build(mut self) -> Result<Box<dyn Processor>> {
+        select_row_type(&mut self)
+    }
+}
+
+impl RowsTypeVisitor for BoundedMergeSortBuilder {
+    type Result = Result<Box<dyn Processor>>;
+
+    fn schema(&self) -> DataSchemaRef {
+        self.schema.clone()
+    }
+
+    fn sort_desc(&self) -> &[SortColumnDescription] {
+        &self.sort_desc
+    }
+
+    fn visit_type<R, C>(&mut self) -> Self::Result
+    where
+        R: Rows + 'static,
+        C: RowConverter<R> + Send + 'static,
+    {
+        match self.enable_loser_tree {
+            true => Ok(Box::new(
+                BoundedMultiSortMergeProcessor::<LoserTreeSort<R>>::new(
+                    self.inputs.clone(),
+                    self.output.clone(),
+                    self.schema.clone(),
+                    self.block_size,
+                    self.limit,
+                    self.remove_order_col,
+                )?,
+            )),
+            false => Ok(Box::new(
+                BoundedMultiSortMergeProcessor::<HeapSort<R>>::new(
+                    self.inputs.clone(),
+                    self.output.clone(),
+                    self.schema.clone(),
+                    self.block_size,
+                    self.limit,
+                    self.remove_order_col,
+                )?,
+            )),
+        }
+    }
+}
diff --git a/src/query/service/src/pipelines/processors/transforms/sort/sort_merge_stream.rs b/src/query/service/src/pipelines/processors/transforms/sort/sort_merge_stream.rs
@@ -53,7 +53,7 @@ where A: SortAlgorithm
 impl<A> BoundedMultiSortMergeProcessor<A>
 where A: SortAlgorithm
 {
-    pub fn create(
+    pub fn new(
         inputs: Vec<Arc<InputPort>>,
         output: Arc<OutputPort>,
         schema: DataSchemaRef,

Original file line number	Diff line number	Diff line change
`@@ -53,7 +53,7 @@ where A: SortAlgorithm`
`53`	`53`	`impl<A> BoundedMultiSortMergeProcessor<A>`
`54`	`54`	`where A: SortAlgorithm`
`55`	`55`	`{`
`56`		`- pub fn create(`
	`56`	`+ pub fn new(`
`57`	`57`	`inputs: Vec<Arc<InputPort>>,`
`58`	`58`	`output: Arc<OutputPort>,`
`59`	`59`	`schema: DataSchemaRef,`