BoundedMultiSortMergeProcessor

Author: forsaken628

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

@@ -12,22 +12,180 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::any::Any;
+use std::cmp::Ordering;
 use std::marker::PhantomData;
 use std::sync::Arc;
 
 use databend_common_exception::Result;
 use databend_common_expression::BlockMetaInfoDowncast;
 use databend_common_expression::Column;
 use databend_common_expression::DataBlock;
+use databend_common_expression::DataSchemaRef;
 use databend_common_expression::Scalar;
+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::sort::algorithm::SortAlgorithm;
+use databend_common_pipeline_transforms::sort::Merger;
 use databend_common_pipeline_transforms::sort::Rows;
 use databend_common_pipeline_transforms::sort::SortedStream;
 
 use crate::pipelines::processors::transforms::SortBound;
 
-/// InputBoundStream is a stream of blocks that are cutoff less or equal than bound.
-struct InputBoundStream<R: Rows> {
+type Stream<A> = BoundedInputStream<<A as SortAlgorithm>::Rows>;
+
+pub struct BoundedMultiSortMergeProcessor<A>
+where A: SortAlgorithm
+{
+    inputs: Vec<Arc<InputPort>>,
+    output: Arc<OutputPort>,
+    schema: DataSchemaRef,
+    block_size: usize,
+    limit: Option<usize>,
+
+    output_data: Option<DataBlock>,
+    bound: Option<Scalar>,
+    inner: std::result::Result<Merger<A, Stream<A>>, Vec<Stream<A>>>,
+}
+
+impl<A> BoundedMultiSortMergeProcessor<A>
+where A: SortAlgorithm
+{
+    pub fn create(
+        inputs: Vec<Arc<InputPort>>,
+        output: Arc<OutputPort>,
+        schema: DataSchemaRef,
+        block_size: usize,
+        limit: Option<usize>,
+        remove_order_col: bool,
+    ) -> Result<Self> {
+        let streams = inputs
+            .iter()
+            .map(|input| BoundedInputStream {
+                data: None,
+                input: input.clone(),
+                remove_order_col,
+                bound: None,
+                sort_row_offset: schema.fields().len() - 1,
+                _r: PhantomData,
+            })
+            .collect();
+
+        Ok(Self {
+            inputs,
+            output,
+            schema,
+            block_size,
+            limit,
+            output_data: None,
+            bound: None,
+            inner: Err(streams),
+        })
+    }
+}
+
+impl<A> Processor for BoundedMultiSortMergeProcessor<A>
+where A: SortAlgorithm + 'static
+{
+    fn name(&self) -> String {
+        "BoundedMultiSortMerge".to_string()
+    }
+
+    fn as_any(&mut self) -> &mut dyn Any {
+        self
+    }
+
+    fn event(&mut self) -> Result<Event> {
+        if self.output.is_finished() {
+            for input in self.inputs.iter() {
+                input.finish();
+            }
+            return Ok(Event::Finished);
+        }
+
+        if !self.output.can_push() {
+            return Ok(Event::NeedConsume);
+        }
+
+        if let Some(block) = self.output_data.take() {
+            self.output.push_data(Ok(block));
+            return Ok(Event::NeedConsume);
+        }
+
+        self.next_event()
+    }
+
+    fn process(&mut self) -> Result<()> {
+        if let Some(block) = self.inner.as_mut().ok().unwrap().next_block()? {
+            let meta = SortBound {
+                bound: self.bound.clone(),
+            }
+            .boxed();
+            self.output_data = Some(block.add_meta(Some(meta))?);
+        };
+        Ok(())
+    }
+}
+
+impl<A> BoundedMultiSortMergeProcessor<A>
+where A: SortAlgorithm + 'static
+{
+    fn next_event(&mut self) -> Result<Event> {
+        let streams = match &mut self.inner {
+            inner @ Ok(_) => {
+                let merger = inner.as_ref().ok().unwrap();
+                if !merger.is_finished() {
+                    return Ok(Event::Sync);
+                }
+                let merger = std::mem::replace(inner, Err(vec![])).ok().unwrap();
+                self.inner = Err(merger.streams());
+                self.inner.as_mut().err().unwrap()
+            }
+            Err(streams) => streams,
+        };
+
+        let mut bounds = Vec::with_capacity(streams.len());
+        for stream in streams.iter_mut() {
+            if stream.pull()? {
+                return Ok(Event::NeedData);
+            }
+            let Some(data) = &stream.data else {
+                continue;
+            };
+            let meta = data
+                .get_meta()
+                .and_then(SortBound::downcast_ref_from)
+                .expect("require a SortBound");
+            bounds.push(&meta.bound)
+        }
+
+        self.bound = match bounds.iter().min_by(|a, b| match (a, b) {
+            (None, None) => Ordering::Equal,
+            (None, Some(_)) => Ordering::Greater,
+            (Some(_), None) => Ordering::Less,
+            (Some(a), Some(b)) => A::Rows::scalar_as_item(a).cmp(&A::Rows::scalar_as_item(b)),
+        }) {
+            Some(bound) => (*bound).clone(),
+            None => return Ok(Event::Finished),
+        };
+        for stream in streams.iter_mut() {
+            stream.bound = self.bound.clone();
+        }
+
+        let streams = std::mem::take(streams);
+        self.inner = Ok(Merger::create(
+            self.schema.clone(),
+            streams,
+            self.block_size,
+            self.limit,
+        ));
+        Ok(Event::Sync)
+    }
+}
+
+struct BoundedInputStream<R: Rows> {
     data: Option<DataBlock>,
     input: Arc<InputPort>,
     remove_order_col: bool,
@@ -36,7 +194,7 @@ struct InputBoundStream<R: Rows> {
     _r: PhantomData<R>,
 }
 
-impl<R: Rows> SortedStream for InputBoundStream<R> {
+impl<R: Rows> SortedStream for BoundedInputStream<R> {
     fn next(&mut self) -> Result<(Option<(DataBlock, Column)>, bool)> {
         if self.pull()? {
             return Ok((None, true));
@@ -59,7 +217,7 @@ fn sort_column(data: &DataBlock, sort_row_offset: usize) -> &Column {
     data.get_by_offset(sort_row_offset).as_column().unwrap()
 }
 
-impl<R: Rows> InputBoundStream<R> {
+impl<R: Rows> BoundedInputStream<R> {
     fn pull(&mut self) -> Result<bool> {
         if self.data.is_some() {
             return Ok(false);
@@ -79,19 +237,18 @@ impl<R: Rows> InputBoundStream<R> {
     }
 
     fn take_next_bounded_block(&mut self) -> Option<DataBlock> {
-        let Some(bound) = &self.bound else {
-            return self.data.take();
-        };
-
         let meta = self
             .data
             .as_ref()?
             .get_meta()
             .and_then(SortBound::downcast_ref_from)
             .expect("require a SortBound");
 
-        if meta.bound.as_ref() == Some(bound) {
-            self.data.take()
+        if meta.bound == self.bound {
+            self.data.take().map(|mut data| {
+                let _ = data.take_meta();
+                data
+            })
         } else {
             None
         }