fix

Signed-off-by: coldWater <forsaken628@gmail.com>

Author: forsaken628

src/query/service/src/pipelines/builders/builder_sort.rs

@@ -33,8 +33,8 @@ use databend_common_storage::DataOperator;
 use databend_common_storages_fuse::TableContext;
 
 use crate::pipelines::memory_settings::MemorySettingsExt;
-use crate::pipelines::processors::transforms::add_range_shuffle_exchange;
 use crate::pipelines::processors::transforms::add_range_shuffle_route;
+use crate::pipelines::processors::transforms::SortRangeExchange;
 use crate::pipelines::processors::transforms::SortSampleState;
 use crate::pipelines::processors::transforms::TransformLimit;
 use crate::pipelines::processors::transforms::TransformSortBuilder;
@@ -137,8 +137,9 @@ impl PipelineBuilder {
             None => {
                 // Build for single node mode.
                 // We build the full sort pipeline for it.
-                let k = self.settings.get_range_shuffle_sort_simple_size()?;
-                if k > 0 && self.main_pipeline.output_len() > 1 {
+                if self.settings.get_enable_range_shuffle_sort()?
+                    && self.main_pipeline.output_len() > 1
+                {
                     builder
                         .remove_order_col_at_last()
                         .build_range_shuffle_sort_pipeline(&mut self.main_pipeline)
@@ -259,7 +260,11 @@ impl SortPipelineBuilder {
 
         builder.add_shuffle(pipeline, state.clone())?;
 
-        add_range_shuffle_exchange(pipeline, max_threads)?;
+        pipeline.exchange(max_threads, Arc::new(SortRangeExchange));
+
+        pipeline.add_transform(|input, output| {
+            Ok(ProcessorPtr::create(builder.build_combine(input, output)?))
+        })?;
 
         pipeline.add_transform(|input, output| {
             Ok(ProcessorPtr::create(builder.build_exec(input, output)?))

src/query/service/src/pipelines/executor/executor_graph.rs

@@ -789,7 +789,7 @@ impl RunningGraph {
             true => Ok(()),
             false => Err(ErrorCode::Internal(format!(
                 "Pipeline graph is not finished, details: {}",
-                self.format_graph_nodes()
+                self.format_graph_nodes(true)
             ))),
         }
     }
@@ -855,7 +855,7 @@ impl RunningGraph {
         self.0.finished_notify.clone()
     }
 
-    pub fn format_graph_nodes(&self) -> String {
+    pub fn format_graph_nodes(&self, pretty: bool) -> String {
         pub struct NodeDisplay {
             id: usize,
             name: String,
@@ -955,7 +955,11 @@ impl RunningGraph {
             }
         }
 
-        format!("{:?}", nodes_display)
+        if pretty {
+            format!("{:#?}", nodes_display)
+        } else {
+            format!("{:?}", nodes_display)
+        }
     }
 
     /// Change the priority

src/query/service/src/pipelines/executor/pipeline_executor.rs

@@ -268,7 +268,7 @@ impl PipelineExecutor {
     pub fn format_graph_nodes(&self) -> String {
         match self {
             PipelineExecutor::QueryPipelineExecutor(executor) => executor.format_graph_nodes(),
-            PipelineExecutor::QueriesPipelineExecutor(v) => v.graph.format_graph_nodes(),
+            PipelineExecutor::QueriesPipelineExecutor(v) => v.graph.format_graph_nodes(false),
         }
     }
 

src/query/service/src/pipelines/executor/processor_async_task.rs

@@ -157,7 +157,7 @@ impl ProcessorAsyncTask {
                                     processor_name,
                                     elapsed,
                                     active_workers,
-                                    graph_clone.format_graph_nodes()
+                                    graph_clone.format_graph_nodes(false)
                                 );
                             }
                         };

src/query/service/src/pipelines/executor/query_pipeline_executor.rs

@@ -431,7 +431,7 @@ impl QueryPipelineExecutor {
     }
 
     pub fn format_graph_nodes(&self) -> String {
-        self.graph.format_graph_nodes()
+        self.graph.format_graph_nodes(false)
     }
 
     pub fn fetch_plans_profile(&self, collect_metrics: bool) -> HashMap<u32, PlanProfile> {

src/query/service/src/pipelines/processors/transforms/sort/mod.rs

@@ -28,6 +28,7 @@ mod bounds;
 mod merge_sort;
 mod sort_builder;
 mod sort_collect;
+mod sort_combine;
 mod sort_exchange;
 mod sort_execute;
 mod sort_route;
@@ -60,7 +61,7 @@ local_block_meta_serde!(SortCollectedMeta);
 impl BlockMetaInfo for SortCollectedMeta {}
 
 #[derive(Debug)]
-struct SortScatteredMeta(pub Vec<SortCollectedMeta>);
+struct SortScatteredMeta(pub Vec<Option<SortCollectedMeta>>);
 
 local_block_meta_serde!(SortScatteredMeta);
 

src/query/service/src/pipelines/processors/transforms/sort/sort_builder.rs

@@ -31,10 +31,12 @@ use databend_common_pipeline_transforms::sort::utils::ORDER_COL_NAME;
 use databend_common_pipeline_transforms::sort::RowConverter;
 use databend_common_pipeline_transforms::sort::Rows;
 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_execute::TransformSortExecute;
 use super::sort_shuffle::SortSampleState;
 use super::sort_shuffle::TransformSortShuffle;
@@ -46,6 +48,7 @@ enum SortType {
     Collect,
     Execute,
     Shuffle,
+    Combine,
 }
 
 pub struct TransformSortBuilder {
@@ -183,6 +186,25 @@ impl TransformSortBuilder {
         select_row_type(&mut build)
     }
 
+    pub fn build_combine(
+        &self,
+        input: Arc<InputPort>,
+        output: Arc<OutputPort>,
+    ) -> Result<Box<dyn Processor>> {
+        self.check();
+
+        let mut build = Build {
+            params: self,
+            input,
+            output,
+            typ: SortType::Combine,
+            id: 0,
+            state: None,
+        };
+
+        select_row_type(&mut build)
+    }
+
     fn should_use_sort_limit(&self) -> bool {
         self.limit.map(|limit| limit < 10000).unwrap_or_default()
     }
@@ -288,6 +310,15 @@ impl Build<'_> {
             self.params.spiller.clone(),
         )))
     }
+
+    fn build_sort_combine<R>(&mut self) -> Result<Box<dyn Processor>>
+    where R: Rows + 'static {
+        Ok(AccumulatingTransformer::create(
+            self.input.clone(),
+            self.output.clone(),
+            TransformSortCombine::<R>::new(self.params.block_size),
+        ))
+    }
 }
 
 impl RowsTypeVisitor for Build<'_> {
@@ -320,6 +351,7 @@ impl RowsTypeVisitor for Build<'_> {
                 false => self.build_sort_exec::<HeapSort<R>>(),
             },
             SortType::Shuffle => self.build_sort_shuffle::<R>(),
+            SortType::Combine => self.build_sort_combine::<R>(),
         }
     }
 }

src/query/service/src/pipelines/processors/transforms/sort/sort_combine.rs

@@ -0,0 +1,80 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use databend_common_exception::Result;
+use databend_common_expression::BlockMetaInfoDowncast;
+use databend_common_expression::DataBlock;
+use databend_common_pipeline_transforms::sort::Rows;
+use databend_common_pipeline_transforms::AccumulatingTransform;
+
+use super::bounds::Bounds;
+use super::SortCollectedMeta;
+
+pub struct TransformSortCombine<R: Rows> {
+    batch_rows: usize,
+    metas: Vec<SortCollectedMeta>,
+    _r: std::marker::PhantomData<R>,
+}
+
+impl<R: Rows> TransformSortCombine<R> {
+    pub fn new(batch_rows: usize) -> Self {
+        Self {
+            batch_rows,
+            metas: vec![],
+            _r: Default::default(),
+        }
+    }
+}
+
+impl<R: Rows> AccumulatingTransform for TransformSortCombine<R> {
+    const NAME: &'static str = "TransformSortCombine";
+
+    fn transform(&mut self, mut data: DataBlock) -> Result<Vec<DataBlock>> {
+        self.metas.push(
+            data.take_meta()
+                .and_then(SortCollectedMeta::downcast_from)
+                .expect("require a SortCollectedMeta"),
+        );
+        Ok(vec![])
+    }
+
+    fn on_finish(&mut self, output: bool) -> Result<Vec<DataBlock>> {
+        if !output || self.metas.is_empty() {
+            return Ok(vec![]);
+        }
+
+        let params = self.metas.first().map(|meta| meta.params).unwrap();
+
+        let bounds = self
+            .metas
+            .iter_mut()
+            .map(|meta| std::mem::take(&mut meta.bounds))
+            .collect();
+        let bounds = Bounds::merge::<R>(bounds, self.batch_rows)?;
+
+        let blocks = self
+            .metas
+            .drain(..)
+            .flat_map(|meta| meta.blocks.into_iter())
+            .collect();
+
+        Ok(vec![DataBlock::empty_with_meta(Box::new(
+            SortCollectedMeta {
+                params,
+                bounds,
+                blocks,
+            },
+        ))])
+    }
+}

src/query/service/src/pipelines/processors/transforms/sort/sort_exchange.rs

@@ -12,76 +12,33 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::iter;
-use std::sync::Arc;
-
 use databend_common_exception::Result;
 use databend_common_expression::BlockMetaInfoDowncast;
 use databend_common_expression::DataBlock;
 use databend_common_pipeline_core::processors::Exchange;
-use databend_common_pipeline_core::processors::InputPort;
-use databend_common_pipeline_core::processors::OutputPort;
-use databend_common_pipeline_core::processors::PartitionProcessor;
-use databend_common_pipeline_core::Pipe;
-use databend_common_pipeline_core::PipeItem;
-use databend_common_pipeline_core::Pipeline;
 
 use super::SortScatteredMeta;
 
-struct SortRangeExchange;
+pub struct SortRangeExchange;
 
 impl Exchange for SortRangeExchange {
     const NAME: &'static str = "SortRange";
     fn partition(&self, mut data: DataBlock, n: usize) -> Result<Vec<DataBlock>> {
-        let Some(meta) = data.take_meta() else {
-            unreachable!();
-        };
-
-        let Some(SortScatteredMeta(scattered)) = SortScatteredMeta::downcast_from(meta) else {
-            unreachable!();
-        };
-
+        let scattered = data
+            .take_meta()
+            .and_then(SortScatteredMeta::downcast_from)
+            .expect("require a SortScatteredMeta")
+            .0;
         assert!(scattered.len() <= n);
 
         let blocks = scattered
             .into_iter()
-            .map(|meta| DataBlock::empty_with_meta(Box::new(meta)))
+            .map(|meta| {
+                meta.map(|meta| DataBlock::empty_with_meta(Box::new(meta)))
+                    .unwrap_or_else(DataBlock::empty)
+            })
             .collect();
 
         Ok(blocks)
     }
 }
-
-fn create_exchange_pipe(num_input: usize, num_output: usize) -> Pipe {
-    let items = iter::repeat_with(|| {
-        let input = InputPort::create();
-        let outputs = iter::repeat_with(OutputPort::create)
-            .take(num_output)
-            .collect::<Vec<_>>();
-
-        PipeItem::create(
-            PartitionProcessor::create(input.clone(), outputs.clone(), Arc::new(SortRangeExchange)),
-            vec![input],
-            outputs,
-        )
-    })
-    .take(num_input)
-    .collect::<Vec<_>>();
-
-    Pipe::create(num_input, num_input * num_output, items)
-}
-
-pub fn add_range_shuffle_exchange(pipeline: &mut Pipeline, num_output: usize) -> Result<()> {
-    let num_input = pipeline.output_len();
-
-    pipeline.add_pipe(create_exchange_pipe(num_input, num_output));
-
-    let n = num_output;
-    let reorder_edges = (0..num_input * n)
-        .map(|i| (i % n) * num_input + (i / n))
-        .collect::<Vec<_>>();
-
-    pipeline.reorder_inputs(reorder_edges);
-
-    Ok(())
-}

src/query/service/src/pipelines/processors/transforms/sort/sort_execute.rs

@@ -91,6 +91,11 @@ where
             return Ok(Event::NeedConsume);
         }
 
+        if self.input.is_finished() && self.inner.is_none() {
+            self.output.finish();
+            return Ok(Event::Finished);
+        }
+
         if let Some(mut block) = self.input.pull_data().transpose()? {
             assert!(self.inner.is_none());
             let meta = block
@@ -102,12 +107,12 @@ where
             return Ok(Event::Async);
         }
 
-        if self.input.is_finished() {
-            Ok(Event::Async)
-        } else {
-            self.input.set_need_data();
-            Ok(Event::NeedData)
+        if self.inner.is_some() {
+            return Ok(Event::Async);
         }
+
+        self.input.set_need_data();
+        Ok(Event::NeedData)
     }
 
     #[async_backtrace::framed]
@@ -122,6 +127,7 @@ where
         }
         if finish {
             self.output.finish();
+            self.inner = None;
         }
         Ok(())
     }