fix

Author: forsaken628

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

@@ -51,12 +51,12 @@ use crate::spillers::SpillerType;
 
 impl PipelineBuilder {
     pub(crate) fn build_sort(&mut self, sort: &Sort) -> Result<()> {
-        let plan_schema = sort.output_schema()?;
+        let output_schema = sort.output_schema()?;
         let sort_desc = sort
             .order_by
             .iter()
             .map(|desc| {
-                let offset = plan_schema.index_of(&desc.order_by.to_string())?;
+                let offset = output_schema.index_of(&desc.order_by.to_string())?;
                 Ok(SortColumnDescription {
                     offset,
                     asc: desc.asc,
@@ -103,7 +103,7 @@ impl PipelineBuilder {
 
         let builder = SortPipelineBuilder::create(
             self.ctx.clone(),
-            plan_schema,
+            output_schema,
             sort_desc,
             sort.broadcast_id,
         )?
@@ -155,7 +155,9 @@ impl PipelineBuilder {
                     self.build_pipeline(&sort.input)?;
                 }
 
-                builder.build_bounded_merge_sort(&mut self.main_pipeline)
+                builder
+                    .remove_order_col_at_last()
+                    .build_bounded_merge_sort(&mut self.main_pipeline)
             }
             SortStep::Route => TransformSortBuilder::add_route(&mut self.main_pipeline),
         }
@@ -164,7 +166,7 @@ impl PipelineBuilder {
 
 pub struct SortPipelineBuilder {
     ctx: Arc<QueryContext>,
-    schema: DataSchemaRef,
+    output_schema: DataSchemaRef,
     sort_desc: Arc<[SortColumnDescription]>,
     limit: Option<usize>,
     block_size: usize,
@@ -176,7 +178,7 @@ pub struct SortPipelineBuilder {
 impl SortPipelineBuilder {
     pub fn create(
         ctx: Arc<QueryContext>,
-        schema: DataSchemaRef,
+        output_schema: DataSchemaRef,
         sort_desc: Arc<[SortColumnDescription]>,
         broadcast_id: Option<u32>,
     ) -> Result<Self> {
@@ -185,7 +187,7 @@ impl SortPipelineBuilder {
         let enable_loser_tree = settings.get_enable_loser_tree_merge_sort()?;
         Ok(Self {
             ctx,
-            schema,
+            output_schema,
             sort_desc,
             limit: None,
             block_size,
@@ -223,65 +225,6 @@ impl SortPipelineBuilder {
         self.build_merge_sort_pipeline(pipeline, false)
     }
 
-    fn build_sample(self, pipeline: &mut Pipeline) -> Result<()> {
-        let settings = self.ctx.get_settings();
-        let max_block_size = settings.get_max_block_size()? as usize;
-
-        // Partial sort
-        pipeline.add_transformer(|| {
-            TransformSortPartial::new(
-                LimitType::from_limit_rows(self.limit),
-                self.sort_desc.clone(),
-            )
-        });
-
-        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)?)
-        };
-
-        let memory_settings = MemorySettings::from_sort_settings(&self.ctx)?;
-        let enable_loser_tree = settings.get_enable_loser_tree_merge_sort()?;
-
-        let builder =
-            TransformSortBuilder::new(self.schema.clone(), self.sort_desc.clone(), max_block_size)
-                .with_spiller(spiller)
-                .with_limit(self.limit)
-                .with_order_column(false, true)
-                .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)?))
-        })?;
-
-        builder.add_bound_broadcast(
-            pipeline,
-            max_block_size,
-            self.ctx.clone(),
-            self.broadcast_id.unwrap(),
-        )?;
-
-        pipeline.add_transform(|input, output| {
-            Ok(ProcessorPtr::create(builder.build_restore(input, output)?))
-        })?;
-
-        pipeline.add_transform(|input, output| {
-            Ok(ProcessorPtr::create(
-                builder.build_bound_edge(input, output)?,
-            ))
-        })?;
-
-        Ok(())
-    }
-
     fn build_merge_sort(&self, pipeline: &mut Pipeline, order_col_generated: bool) -> Result<()> {
         // Merge sort
         let need_multi_merge = pipeline.output_len() > 1;
@@ -293,8 +236,8 @@ impl SortPipelineBuilder {
 
         let memory_settings = MemorySettings::from_sort_settings(&self.ctx)?;
         let sort_merge_output_schema = match output_order_col {
-            true => add_order_field(self.schema.clone(), &self.sort_desc),
-            false => self.schema.clone(),
+            true => add_order_field(self.output_schema.clone(), &self.sort_desc),
+            false => self.output_schema.clone(),
         };
 
         let settings = self.ctx.get_settings();
@@ -336,7 +279,7 @@ impl SortPipelineBuilder {
         })
     }
 
-    pub fn build_merge_sort_pipeline(
+    fn build_merge_sort_pipeline(
         self,
         pipeline: &mut Pipeline,
         order_col_generated: bool,
@@ -358,7 +301,7 @@ impl SortPipelineBuilder {
             let max_threads = settings.get_max_threads()? as usize;
             add_k_way_merge_sort(
                 pipeline,
-                self.schema.clone(),
+                self.output_schema.clone(),
                 max_threads,
                 self.block_size,
                 self.limit,
@@ -369,7 +312,7 @@ impl SortPipelineBuilder {
         } else {
             try_add_multi_sort_merge(
                 pipeline,
-                self.schema.clone(),
+                self.output_schema.clone(),
                 self.block_size,
                 self.limit,
                 self.sort_desc,
@@ -379,12 +322,77 @@ impl SortPipelineBuilder {
         }
     }
 
-    pub fn build_bounded_merge_sort(self, pipeline: &mut Pipeline) -> Result<()> {
-        let builder =
-            TransformSortBuilder::new(self.schema.clone(), self.sort_desc.clone(), self.block_size)
-                .with_limit(self.limit)
-                .with_order_column(true, !self.remove_order_col_at_last)
-                .with_enable_loser_tree(self.enable_loser_tree);
+    fn build_sample(self, pipeline: &mut Pipeline) -> Result<()> {
+        let settings = self.ctx.get_settings();
+        let max_block_size = settings.get_max_block_size()? as usize;
+
+        // Partial sort
+        pipeline.add_transformer(|| {
+            TransformSortPartial::new(
+                LimitType::from_limit_rows(self.limit),
+                self.sort_desc.clone(),
+            )
+        });
+
+        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)?)
+        };
+
+        let memory_settings = MemorySettings::from_sort_settings(&self.ctx)?;
+        let enable_loser_tree = settings.get_enable_loser_tree_merge_sort()?;
+
+        let builder = TransformSortBuilder::new(
+            self.output_schema.clone(),
+            self.sort_desc.clone(),
+            max_block_size,
+        )
+        .with_spiller(spiller)
+        .with_limit(self.limit)
+        .with_order_column(false, true)
+        .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)?))
+        })?;
+
+        builder.add_bound_broadcast(
+            pipeline,
+            max_block_size,
+            self.ctx.clone(),
+            self.broadcast_id.unwrap(),
+        )?;
+
+        pipeline.add_transform(|input, output| {
+            Ok(ProcessorPtr::create(builder.build_restore(input, output)?))
+        })?;
+
+        pipeline.add_transform(|input, output| {
+            Ok(ProcessorPtr::create(
+                builder.build_bound_edge(input, output)?,
+            ))
+        })?;
+
+        Ok(())
+    }
+
+    fn build_bounded_merge_sort(self, pipeline: &mut Pipeline) -> Result<()> {
+        let builder = TransformSortBuilder::new(
+            self.output_schema.clone(),
+            self.sort_desc.clone(),
+            self.block_size,
+        )
+        .with_limit(self.limit)
+        .with_order_column(true, !self.remove_order_col_at_last)
+        .with_enable_loser_tree(self.enable_loser_tree);
 
         let inputs_port: Vec<_> = (0..pipeline.output_len())
             .map(|_| InputPort::create())

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

@@ -31,6 +31,9 @@ pub struct Bounds(
 
 impl Bounds {
     pub fn new_unchecked(column: Column) -> Bounds {
+        if column.len() == 0 {
+            return Self::default();
+        }
         Bounds(vec![column])
     }
 
@@ -49,7 +52,6 @@ impl Bounds {
     }
 
     pub fn merge<R: Rows>(mut vector: Vec<Bounds>, batch_rows: usize) -> Result<Self> {
-        debug_assert!(vector.iter().all(|bounds| !bounds.is_empty()));
         match vector.len() {
             0 => Ok(Bounds(vec![])),
             1 => Ok(vector.pop().unwrap()),
@@ -68,7 +70,7 @@ impl Bounds {
                     blocks
                         .iter()
                         .rev()
-                        .map(|b| b.get_last_column().clone())
+                        .map(|b| b.get_by_offset(0).to_column())
                         .collect(),
                 ))
             }
@@ -136,6 +138,7 @@ impl Bounds {
         )]))
     }
 
+    #[allow(dead_code)]
     pub fn dedup_reduce<R: Rows>(&self, n: usize) -> Self {
         if n == 0 {
             return Self::default();
@@ -178,11 +181,8 @@ impl Bounds {
             last = Some((cur_rows, r_idx));
         }
 
-        Bounds(vec![Column::take_column_indices(
-            &self.0,
-            &indices,
-            indices.len(),
-        )])
+        let col = Column::take_column_indices(&self.0, &indices, indices.len());
+        Bounds::new_unchecked(col)
     }
 
     pub fn dedup<R: Rows>(&self) -> Self {

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

@@ -138,9 +138,7 @@ impl<R: Rows + 'static> HookTransform for TransformSortBoundBroadcast<R> {
         let bounds = Bounds::merge::<R>(
             self.input_data
                 .iter_mut()
-                .filter_map(|meta| {
-                    (!meta.bounds.is_empty()).then(|| std::mem::take(&mut meta.bounds))
-                })
+                .map(|meta| std::mem::take(&mut meta.bounds))
                 .collect(),
             self.state.batch_rows,
         )?;
@@ -163,42 +161,43 @@ impl<R: Rows + 'static> HookTransform for TransformSortBoundBroadcast<R> {
             sequences,
         };
 
+        let bounds = local.normalize_bounds();
         let global = self
             .state
-            .commit_sample(Some(SortExchangeMeta {
-                params,
-                bounds: local.normalize_bounds::<R>(),
-            }))
+            .commit_sample(Some(SortExchangeMeta { params, bounds }))
             .await?;
 
         let bounds_vec = global
             .into_iter()
-            .filter_map(|meta| (!meta.bounds.is_empty()).then_some(meta.bounds))
+            .map(|meta| {
+                assert!(!meta.bounds.is_empty());
+                meta.bounds.clone()
+            })
             .collect();
-        self.output_data = Some(SortCollectedMeta {
-            bounds: Bounds::merge::<R>(bounds_vec, self.state.batch_rows)?.dedup::<R>(),
-            ..local
-        });
+
+        let bounds = Bounds::merge::<R>(bounds_vec, self.state.batch_rows)?.dedup::<R>();
+        log::debug!("global_bounds.len: {}", bounds.len());
+        self.output_data = Some(SortCollectedMeta { bounds, ..local });
         Ok(())
     }
 }
 
 impl SortCollectedMeta {
-    fn normalize_bounds<R: Rows>(&self) -> Bounds {
+    fn normalize_bounds(&self) -> Bounds {
         if self.bounds.len() > 1 {
-            return self.bounds.dedup::<R>();
+            return self.bounds.clone();
         }
 
         let Some(seq) = self.sequences.get(self.sequences.len() / 2) else {
-            return Bounds::default();
+            unreachable!()
         };
 
         seq.get(seq.len() / 2)
             .map(|block| match block.domain.len() {
                 0 => Bounds::default(),
                 1 => Bounds::new_unchecked(block.domain.clone()),
-                _ => Bounds::new_unchecked(block.domain.slice(0..1)).dedup::<R>(),
+                _ => Bounds::new_unchecked(block.domain.slice(0..1)),
             })
-            .unwrap_or_default()
+            .unwrap()
     }
 }

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

@@ -355,9 +355,8 @@ impl Build<'_> {
         Ok(Box::new(BoundedMultiSortMergeProcessor::<A>::new(
             inputs,
             self.output.clone(),
-            self.schema().clone(),
+            self.params.schema.clone(),
             self.params.block_size,
-            !self.params.output_order_col,
         )?))
     }
 }