build

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

Author: forsaken628

src/query/pipeline/transforms/src/processors/transforms/sort/rows/utils.rs

@@ -12,7 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_expression::row::RowConverter as CommonConverter;
 use databend_common_expression::types::DataType;
@@ -45,10 +44,10 @@ pub fn convert_rows(
         schema: DataSchemaRef,
         sort_desc: &'a [SortColumnDescription],
         data: DataBlock,
-        result: Result<Column>,
     }
 
     impl RowsTypeVisitor for ConvertRowsVisitor<'_> {
+        type Result = Result<Column>;
         fn schema(&self) -> DataSchemaRef {
             self.schema.clone()
         }
@@ -57,7 +56,7 @@ pub fn convert_rows(
             self.sort_desc
         }
 
-        fn visit_type<R, C>(&mut self)
+        fn visit_type<R, C>(&mut self) -> Self::Result
         where
             R: Rows + 'static,
             C: RowConverter<R> + Send + 'static,
@@ -68,65 +67,63 @@ pub fn convert_rows(
                 .map(|desc| self.data.get_by_offset(desc.offset).to_owned())
                 .collect::<Vec<_>>();
 
-            self.result = try {
-                let converter = C::create(self.sort_desc, self.schema.clone())?;
-                let rows = C::convert(&converter, &columns, self.data.num_rows())?;
-                rows.to_column()
-            }
+            let converter = C::create(self.sort_desc, self.schema.clone())?;
+            let rows = C::convert(&converter, &columns, self.data.num_rows())?;
+            Ok(rows.to_column())
         }
     }
 
     let mut visitor = ConvertRowsVisitor {
         schema: schema.clone(),
         sort_desc,
         data,
-        result: Err(ErrorCode::Internal("unreachable")),
     };
 
-    select_row_type(&mut visitor);
-    visitor.result
+    select_row_type(&mut visitor)
 }
 
-pub fn select_row_type(visitor: &mut impl RowsTypeVisitor) {
-    let sort_desc = visitor.sort_desc();
-    if sort_desc.len() == 1 {
-        let schema = visitor.schema();
-        let sort_type = schema.field(sort_desc[0].offset).data_type();
-        let asc = sort_desc[0].asc;
-
-        match_template! {
-        T = [ Date => DateType, Timestamp => TimestampType, String => StringType ],
-        match sort_type {
-            DataType::T => {
-                if asc {
-                    visitor.visit_type::<SimpleRowsAsc<T>, SimpleRowConverter<T>>()
-                } else {
-                    visitor.visit_type::<SimpleRowsDesc<T>, SimpleRowConverter<T>>()
-                }
-            },
-            DataType::Number(num_ty) => with_number_mapped_type!(|NUM_TYPE| match num_ty {
-                NumberDataType::NUM_TYPE => {
+pub fn select_row_type<V>(visitor: &mut V) -> V::Result
+where V: RowsTypeVisitor {
+    match &visitor.sort_desc() {
+        &[desc] => {
+            let schema = visitor.schema();
+            let sort_type = schema.field(desc.offset).data_type();
+            let asc = desc.asc;
+
+            match_template! {
+            T = [ Date => DateType, Timestamp => TimestampType, String => StringType ],
+            match sort_type {
+                DataType::T => {
                     if asc {
-                        visitor.visit_type::<SimpleRowsAsc<NumberType<NUM_TYPE>>, SimpleRowConverter<NumberType<NUM_TYPE>>>()
+                        visitor.visit_type::<SimpleRowsAsc<T>, SimpleRowConverter<T>>()
                     } else {
-                        visitor.visit_type::<SimpleRowsDesc<NumberType<NUM_TYPE>>, SimpleRowConverter<NumberType<NUM_TYPE>>>()
+                        visitor.visit_type::<SimpleRowsDesc<T>, SimpleRowConverter<T>>()
+                    }
+                },
+                DataType::Number(num_ty) => with_number_mapped_type!(|NUM_TYPE| match num_ty {
+                    NumberDataType::NUM_TYPE => {
+                        if asc {
+                            visitor.visit_type::<SimpleRowsAsc<NumberType<NUM_TYPE>>, SimpleRowConverter<NumberType<NUM_TYPE>>>()
+                        } else {
+                            visitor.visit_type::<SimpleRowsDesc<NumberType<NUM_TYPE>>, SimpleRowConverter<NumberType<NUM_TYPE>>>()
+                        }
                     }
+                }),
+                _ => visitor.visit_type::<CommonRows, CommonConverter>()
                 }
-            }),
-            _ => visitor.visit_type::<CommonRows, CommonConverter>()
             }
         }
-    } else {
-        visitor.visit_type::<CommonRows, CommonConverter>()
+        _ => visitor.visit_type::<CommonRows, CommonConverter>(),
     }
 }
 
 pub trait RowsTypeVisitor {
+    type Result;
     fn schema(&self) -> DataSchemaRef;
 
     fn sort_desc(&self) -> &[SortColumnDescription];
 
-    fn visit_type<R, C>(&mut self)
+    fn visit_type<R, C>(&mut self) -> Self::Result
     where
         R: Rows + 'static,
         C: RowConverter<R> + Send + 'static;
@@ -136,10 +133,10 @@ pub fn order_field_type(schema: &DataSchema, desc: &[SortColumnDescription]) ->
     struct OrderFieldTypeVisitor<'a> {
         schema: DataSchemaRef,
         sort_desc: &'a [SortColumnDescription],
-        result: Option<DataType>,
     }
 
     impl RowsTypeVisitor for OrderFieldTypeVisitor<'_> {
+        type Result = DataType;
         fn schema(&self) -> DataSchemaRef {
             self.schema.clone()
         }
@@ -148,22 +145,20 @@ pub fn order_field_type(schema: &DataSchema, desc: &[SortColumnDescription]) ->
             self.sort_desc
         }
 
-        fn visit_type<R, C>(&mut self)
+        fn visit_type<R, C>(&mut self) -> Self::Result
         where
             R: Rows + 'static,
             C: RowConverter<R> + Send + 'static,
         {
-            self.result = Some(R::data_type());
+            R::data_type()
         }
     }
 
     assert!(!desc.is_empty());
     let mut visitor = OrderFieldTypeVisitor {
         schema: schema.clone().into(),
         sort_desc: desc,
-        result: None,
     };
 
-    select_row_type(&mut visitor);
-    visitor.result.unwrap()
+    select_row_type(&mut visitor)
 }

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

@@ -34,10 +34,9 @@ use databend_common_storages_fuse::TableContext;
 use databend_storages_common_cache::TempDirManager;
 
 use crate::pipelines::memory_settings::MemorySettingsExt;
-use crate::pipelines::processors::transforms::sort::add_range_shuffle;
-use crate::pipelines::processors::transforms::sort::add_range_shuffle_merge;
-use crate::pipelines::processors::transforms::sort::add_sort_sample;
-use crate::pipelines::processors::transforms::sort::SortSampleState;
+use crate::pipelines::processors::transforms::add_range_shuffle_exchange;
+use crate::pipelines::processors::transforms::add_range_shuffle_route;
+use crate::pipelines::processors::transforms::SortSampleState;
 use crate::pipelines::processors::transforms::TransformLimit;
 use crate::pipelines::processors::transforms::TransformSortBuilder;
 use crate::pipelines::PipelineBuilder;
@@ -144,7 +143,7 @@ impl PipelineBuilder {
                 if k > 0 && self.main_pipeline.output_len() > 1 {
                     builder
                         .remove_order_col_at_last()
-                        .build_range_shuffle_sort_pipeline(&mut self.main_pipeline, k)
+                        .build_range_shuffle_sort_pipeline(&mut self.main_pipeline)
                 } else {
                     builder
                         .remove_order_col_at_last()
@@ -213,18 +212,11 @@ impl SortPipelineBuilder {
         self.build_merge_sort_pipeline(pipeline, false)
     }
 
-    fn build_range_shuffle_sort_pipeline(self, pipeline: &mut Pipeline, k: usize) -> Result<()> {
+    fn build_range_shuffle_sort_pipeline(self, pipeline: &mut Pipeline) -> Result<()> {
         let inputs = pipeline.output_len();
         let settings = self.ctx.get_settings();
         let max_threads = settings.get_max_threads()? as usize;
-        let sample = SortSampleState::new(
-            inputs,
-            max_threads,
-            self.schema.clone(),
-            self.sort_desc.clone(),
-        );
-
-        add_sort_sample(pipeline, sample.clone(), self.sort_desc.clone(), k)?;
+        let max_block_size = settings.get_max_block_size()? as usize;
 
         // Partial sort
         pipeline.add_transformer(|| {
@@ -234,20 +226,48 @@ impl SortPipelineBuilder {
             )
         });
 
-        self.build_merge_sort(pipeline, false)?;
+        let spiller = {
+            let location_prefix = self.ctx.query_id_spill_prefix();
+            let config = SpillerConfig {
+                spiller_type: SpillerType::OrderBy,
+                location_prefix,
+                disk_spill: None,
+                use_parquet: settings.get_spilling_file_format()?.is_parquet(),
+            };
+            let op = DataOperator::instance().spill_operator();
+            Arc::new(Spiller::create(self.ctx.clone(), op, config)?)
+        };
 
-        add_range_shuffle(
-            pipeline,
-            sample.clone(),
-            self.sort_desc.clone(),
+        let memory_settings = MemorySettings::from_sort_settings(&self.ctx)?;
+        let enable_loser_tree = settings.get_enable_loser_tree_merge_sort()?;
+
+        let builder = TransformSortBuilder::create(
             self.schema.clone(),
-            self.block_size,
-            self.limit,
-            self.remove_order_col_at_last,
-            self.enable_loser_tree,
-        )?;
+            self.sort_desc.clone(),
+            max_block_size,
+            spiller,
+        )
+        .with_limit(self.limit)
+        .with_order_col_generated(false)
+        .with_output_order_col(false)
+        .with_memory_settings(memory_settings)
+        .with_enable_loser_tree(enable_loser_tree);
+
+        pipeline.add_transform(|input, output| {
+            Ok(ProcessorPtr::create(builder.build_collect(input, output)?))
+        })?;
+
+        let state = SortSampleState::new(inputs, max_threads, self.schema.clone(), max_block_size);
+
+        builder.add_shuffle(pipeline, state.clone())?;
+
+        add_range_shuffle_exchange(pipeline, max_threads)?;
+
+        pipeline.add_transform(|input, output| {
+            Ok(ProcessorPtr::create(builder.build_exec(input, output)?))
+        })?;
 
-        add_range_shuffle_merge(pipeline)?;
+        add_range_shuffle_route(pipeline)?;
 
         if self.limit.is_none() {
             return Ok(());