fix

Author: forsaken628

Cargo.lock

@@ -57,11 +57,7 @@ pub fn try_add_multi_sort_merge(
     remove_order_col: bool,
     enable_loser_tree: bool,
 ) -> Result<()> {
-    debug_assert!(if !remove_order_col {
-        schema.has_field(ORDER_COL_NAME)
-    } else {
-        !schema.has_field(ORDER_COL_NAME)
-    });
+    debug_assert!(remove_order_col == schema.has_field(ORDER_COL_NAME));
 
     if pipeline.is_empty() {
         return Err(ErrorCode::Internal("Cannot resize empty pipe."));

src/query/pipeline/transforms/src/processors/transforms/transform_multi_sort_merge.rs

@@ -116,6 +116,7 @@ databend-storages-common-session = { workspace = true }
 databend-storages-common-stage = { workspace = true }
 databend-storages-common-table-meta = { workspace = true }
 derive-visitor = { workspace = true }
+enum-as-inner = { workspace = true }
 ethnum = { workspace = true }
 fastrace = { workspace = true }
 flatbuffers = { workspace = true }

src/query/service/Cargo.toml

@@ -39,6 +39,7 @@
 #![feature(get_mut_unchecked)]
 #![allow(clippy::diverging_sub_expression)]
 #![allow(clippy::arc_with_non_send_sync)]
+#![feature(debug_closure_helpers)]
 
 extern crate core;
 

src/query/service/src/lib.rs

@@ -37,10 +37,10 @@ use databend_common_pipeline_transforms::TransformSortMergeLimit;
 
 use super::sort_spill::create_memory_merger;
 use super::sort_spill::MemoryMerger;
+use super::sort_spill::OutputData;
 use super::sort_spill::SortSpill;
 use super::Base;
 use super::MemoryRows;
-use crate::pipelines::processors::transforms::sort::sort_spill::OutputData;
 use crate::spillers::Spiller;
 
 #[derive(Debug)]

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

@@ -21,6 +21,7 @@ use databend_common_expression::BlockMetaInfoDowncast;
 use databend_common_expression::DataBlock;
 use databend_common_expression::DataSchemaRef;
 use databend_common_pipeline_transforms::SortSpillParams;
+use enum_as_inner::EnumAsInner;
 use sort_spill::SpillableBlock;
 
 use crate::spillers::Spiller;
@@ -95,11 +96,29 @@ impl BlockMetaInfo for SortExchangeMeta {
 #[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
 pub struct SortBound {
     index: u32,
-    next: Option<u32>,
+    next: SortBoundNext,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize, EnumAsInner)]
+pub enum SortBoundNext {
+    Uinit,
+    Next(u32),
+    More,
+    Skip,
+    Last,
+}
+
+impl SortBound {
+    pub fn has_more(&self) -> bool {
+        match self.next {
+            SortBoundNext::Next(next) => self.index == next,
+            _ => unreachable!(),
+        }
+    }
 }
 
 impl SortBound {
-    fn create(index: u32, next: Option<u32>) -> Box<dyn BlockMetaInfo> {
+    fn create(index: u32, next: SortBoundNext) -> Box<dyn BlockMetaInfo> {
         debug_assert!(index != u32::MAX);
         SortBound { index, next }.boxed()
     }

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

@@ -20,16 +20,16 @@ use databend_common_exception::Result;
 use databend_common_expression::BlockMetaInfoDowncast;
 use databend_common_expression::DataBlock;
 use databend_common_pipeline_core::processors::Event;
+use databend_common_pipeline_core::processors::InputPort;
+use databend_common_pipeline_core::processors::OutputPort;
 use databend_common_pipeline_core::processors::Processor;
 use databend_common_pipeline_transforms::processors::sort::Rows;
 use databend_common_pipeline_transforms::HookTransform;
 use databend_common_pipeline_transforms::HookTransformer;
 
 use super::bounds::Bounds;
 use super::SortCollectedMeta;
-use crate::pipelines::processors::transforms::sort::SortExchangeMeta;
-use crate::pipelines::processors::InputPort;
-use crate::pipelines::processors::OutputPort;
+use super::SortExchangeMeta;
 use crate::sessions::QueryContext;
 
 pub struct TransformSortBoundBroadcast<R: Rows> {

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

@@ -20,7 +20,7 @@ use databend_common_expression::DataBlock;
 use databend_common_pipeline_core::Pipeline;
 use databend_common_settings::FlightCompression;
 
-use crate::pipelines::processors::transforms::SortBound;
+use super::SortBound;
 use crate::servers::flight::v1::exchange::DataExchange;
 use crate::servers::flight::v1::exchange::DefaultExchangeInjector;
 use crate::servers::flight::v1::exchange::ExchangeInjector;

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

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::any::Any;
+use std::collections::VecDeque;
 use std::marker::PhantomData;
 use std::sync::Arc;
 
@@ -31,6 +32,7 @@ use databend_common_pipeline_transforms::sort::Rows;
 use databend_common_pipeline_transforms::sort::SortedStream;
 
 use super::SortBound;
+use super::SortBoundNext;
 
 type Stream<A> = BoundedInputStream<<A as SortAlgorithm>::Rows>;
 
@@ -42,8 +44,9 @@ where A: SortAlgorithm
     schema: DataSchemaRef,
     block_size: usize,
 
-    output_data: Option<DataBlock>,
+    output_data: VecDeque<DataBlock>,
     cur_index: u32,
+    cur_output: bool,
     inner: std::result::Result<Merger<A, Stream<A>>, Vec<Stream<A>>>,
 }
 
@@ -56,6 +59,8 @@ where A: SortAlgorithm
         schema: DataSchemaRef,
         block_size: usize,
     ) -> Result<Self> {
+        assert!(inputs.len() != 1);
+
         let streams = inputs
             .iter()
             .map(|input| BoundedInputStream {
@@ -73,8 +78,9 @@ where A: SortAlgorithm
             output,
             schema,
             block_size,
-            output_data: None,
+            output_data: VecDeque::new(),
             cur_index: 0,
+            cur_output: false,
             inner: Err(streams),
         })
     }
@@ -103,7 +109,7 @@ where A: SortAlgorithm + 'static
             return Ok(Event::NeedConsume);
         }
 
-        if let Some(block) = self.output_data.take() {
+        if let Some(block) = self.output_data.pop_front() {
             self.output.push_data(Ok(block));
             return Ok(Event::NeedConsume);
         }
@@ -114,10 +120,16 @@ where A: SortAlgorithm + 'static
     fn process(&mut self) -> Result<()> {
         let merger = self.inner.as_mut().ok().unwrap();
         if let Some(block) = merger.next_block()? {
-            self.output_data = Some(block.add_meta(Some(SortBound::create(
-                self.cur_index,
-                (!merger.is_finished()).then_some(self.cur_index),
-            )))?);
+            self.cur_output = true;
+            self.output_data
+                .push_back(block.add_meta(Some(SortBound::create(
+                    self.cur_index,
+                    if merger.is_finished() {
+                        SortBoundNext::Last
+                    } else {
+                        SortBoundNext::More
+                    },
+                )))?);
         };
         Ok(())
     }
@@ -133,6 +145,19 @@ where A: SortAlgorithm + 'static
                 if !merger.is_finished() {
                     return Ok(Event::Sync);
                 }
+
+                if !self.cur_output {
+                    let empty = DataBlock::empty_with_schema(self.schema.clone());
+                    self.output_data.push_back(
+                        empty.add_meta(Some(SortBound::create(
+                            self.cur_index,
+                            SortBoundNext::Skip,
+                        )))?,
+                    );
+                } else {
+                    self.cur_output = false;
+                }
+
                 self.cur_index += 1;
                 let merger = std::mem::replace(inner, Err(vec![])).ok().unwrap();
                 self.inner = Err(merger.streams());
@@ -235,7 +260,7 @@ impl<R: Rows> BoundedInputStream<R> {
             "meta: {meta:?}, bound: {bound:?}",
         );
         if meta.index == bound.bound_index {
-            bound.more = meta.next.is_some_and(|next| next == meta.index);
+            bound.more = meta.has_more();
             self.data.take().map(|mut data| {
                 data.take_meta().unwrap();
                 data
@@ -268,11 +293,11 @@ mod tests {
 
     use super::*;
 
-    fn create_block(empty: bool, index: u32, next: Option<u32>) -> DataBlock {
+    fn create_block(empty: bool, index: u32, next: u32) -> DataBlock {
         let block = DataBlock::new_from_columns(vec![Int32Type::from_data(vec![1, 2, 3])]);
         let block = if empty { block.slice(0..0) } else { block };
         block
-            .add_meta(Some(SortBound::create(index, next)))
+            .add_meta(Some(SortBound::create(index, SortBoundNext::Next(next))))
             .unwrap()
     }
 
@@ -311,7 +336,7 @@ mod tests {
         }
 
         {
-            let block = create_block(true, 0, Some(0));
+            let block = create_block(true, 0, 0);
             output.push_data(Ok(block));
 
             let (_, pending) = stream.next().unwrap();
@@ -321,7 +346,7 @@ mod tests {
         }
 
         {
-            let block = create_block(false, 0, Some(0));
+            let block = create_block(false, 0, 0);
             output.push_data(Ok(block));
 
             let (data, pending) = stream.next().unwrap();
@@ -332,7 +357,7 @@ mod tests {
         }
 
         {
-            let block = create_block(true, 0, Some(1));
+            let block = create_block(true, 0, 1);
             output.push_data(Ok(block));
 
             let (data, pending) = stream.next().unwrap();
@@ -341,7 +366,7 @@ mod tests {
             assert!(!stream.bound.unwrap().more);
             assert!(pending);
 
-            let block = create_block(false, 1, Some(1));
+            let block = create_block(false, 1, 1);
             output.push_data(Ok(block));
 
             let (data, pending) = stream.next().unwrap();

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

@@ -26,11 +26,12 @@ use databend_common_pipeline_transforms::processors::sort::algorithm::SortAlgori
 use databend_common_pipeline_transforms::HookTransform;
 use databend_common_pipeline_transforms::HookTransformer;
 
+use super::sort_spill::OutputData;
 use super::sort_spill::SortSpill;
 use super::Base;
 use super::SortBound;
+use super::SortBoundNext;
 use super::SortCollectedMeta;
-use crate::pipelines::processors::transforms::sort::sort_spill::OutputData;
 
 pub struct TransformSortRestore<A: SortAlgorithm> {
     input: Vec<SortCollectedMeta>,
@@ -121,7 +122,8 @@ where
             finish,
         } = spill_sort.on_restore().await?;
         if let Some(block) = block {
-            let mut block = block.add_meta(Some(SortBound::create(bound_index, None)))?;
+            let mut block =
+                block.add_meta(Some(SortBound::create(bound_index, SortBoundNext::Uinit)))?;
             if self.remove_order_col {
                 block.pop_columns(1);
             }
@@ -188,7 +190,7 @@ impl Processor for SortBoundEdge {
                 .take_meta()
                 .and_then(SortBound::downcast_from)
                 .expect("require a SortBound");
-            meta.next = None;
+            meta.next = SortBoundNext::Uinit;
             self.output
                 .push_data(Ok(block.add_meta(Some(meta.boxed()))?));
             self.output.finish();
@@ -212,7 +214,7 @@ impl Processor for SortBoundEdge {
             .and_then(SortBound::downcast_mut)
             .expect("require a SortBound");
 
-        output_meta.next = Some(incoming_index);
+        output_meta.next = SortBoundNext::Next(incoming_index);
 
         log::debug!("output_meta: {:?}", output_meta);