feat(query): range shuffle sort for standalone mode

src/common/base/src/base/watch_notify.rs

@@ -41,6 +41,16 @@ impl WatchNotify {
         let _ = rx.changed().await;
     }
 
+    pub fn has_notified(&self) -> bool {
+        match self.rx.has_changed() {
+            Ok(b) => b,
+            Err(_) => {
+                // The sender has never dropped before
+                unreachable!()
+            }
+        }
+    }
+
     pub fn notify_waiters(&self) {
         let _ = self.tx.send_replace(true);
     }
@@ -61,11 +71,18 @@ mod tests {
     #[tokio::test]
     async fn test_notify_waiters_ahead() {
         let notify = WatchNotify::new();
+        assert!(!notify.has_notified());
+        let notified1 = notify.notified();
+        assert!(!notify.has_notified());
+
         // notify_waiters ahead of notified being instantiated and awaited
         notify.notify_waiters();
-
+        assert!(notify.has_notified());
         // this should not await indefinitely
-        let notified = notify.notified();
-        notified.await;
+        let notified2 = notify.notified();
+        notified2.await;
+
+        notified1.await;
+        assert!(notify.has_notified());
     }
 }

src/query/expression/src/block.rs

@@ -332,10 +332,34 @@ pub trait BlockMetaInfoDowncast: Sized + BlockMetaInfo {
         let boxed = boxed.as_ref() as &dyn Any;
         boxed.downcast_ref()
     }
+
+    fn downcast_from_err(boxed: BlockMetaInfoPtr) -> std::result::Result<Self, BlockMetaInfoPtr> {
+        if (boxed.as_ref() as &dyn Any).is::<Self>() {
+            Ok(*(boxed as Box<dyn Any>).downcast().unwrap())
+        } else {
+            Err(boxed)
+        }
+    }
+
+    fn downcast_mut(boxed: &mut BlockMetaInfoPtr) -> Option<&mut Self> {
+        let boxed = boxed.as_mut() as &mut dyn Any;
+        boxed.downcast_mut()
+    }
 }
 
 impl<T: BlockMetaInfo> BlockMetaInfoDowncast for T {}
 
+#[typetag::serde(name = "empty")]
+impl BlockMetaInfo for () {
+    fn equals(&self, info: &Box<dyn BlockMetaInfo>) -> bool {
+        <() as BlockMetaInfoDowncast>::downcast_ref_from(info).is_some()
+    }
+
+    fn clone_self(&self) -> Box<dyn BlockMetaInfo> {
+        Box::new(())
+    }
+}
+
 impl DataBlock {
     #[inline]
     pub fn new(entries: Vec<BlockEntry>, num_rows: usize) -> Self {
@@ -431,11 +455,6 @@ impl DataBlock {
         DataBlock::new_with_meta(vec![], 0, Some(meta))
     }
 
-    #[inline]
-    pub fn take_meta(&mut self) -> Option<BlockMetaInfoPtr> {
-        self.meta.take()
-    }
-
     #[inline]
     pub fn columns(&self) -> &[BlockEntry] {
         &self.entries
@@ -702,6 +721,16 @@ impl DataBlock {
         })
     }
 
+    #[inline]
+    pub fn take_meta(&mut self) -> Option<BlockMetaInfoPtr> {
+        self.meta.take()
+    }
+
+    #[inline]
+    pub fn mut_meta(&mut self) -> Option<&mut BlockMetaInfoPtr> {
+        self.meta.as_mut()
+    }
+
     #[inline]
     pub fn replace_meta(&mut self, meta: BlockMetaInfoPtr) {
         self.meta.replace(meta);

src/query/expression/src/types/array.rs

@@ -54,7 +54,7 @@ impl<T: AccessType> AccessType for ArrayType<T> {
         scalar.clone()
     }
 
-    fn try_downcast_scalar<'a>(scalar: &'a ScalarRef) -> Option<Self::ScalarRef<'a>> {
+    fn try_downcast_scalar<'a>(scalar: &ScalarRef<'a>) -> Option<Self::ScalarRef<'a>> {
         match scalar {
             ScalarRef::Array(array) => T::try_downcast_column(array),
             _ => None,

src/query/expression/src/types/boolean.rs

@@ -50,7 +50,7 @@ impl AccessType for BooleanType {
         *scalar
     }
 
-    fn try_downcast_scalar<'a>(scalar: &'a ScalarRef) -> Option<Self::ScalarRef<'a>> {
+    fn try_downcast_scalar<'a>(scalar: &ScalarRef<'a>) -> Option<Self::ScalarRef<'a>> {
         match scalar {
             ScalarRef::Boolean(scalar) => Some(*scalar),
             _ => None,

src/query/expression/src/types/map.rs

@@ -451,7 +451,7 @@ impl<K: AccessType, V: AccessType> AccessType for MapType<K, V> {
         MapInternal::<K, V>::to_scalar_ref(scalar)
     }
 
-    fn try_downcast_scalar<'a>(scalar: &'a ScalarRef) -> Option<Self::ScalarRef<'a>> {
+    fn try_downcast_scalar<'a>(scalar: &ScalarRef<'a>) -> Option<Self::ScalarRef<'a>> {
         match scalar {
             ScalarRef::Map(array) => KvPair::<K, V>::try_downcast_column(array),
             _ => None,

src/query/pipeline/core/src/pipeline.rs

@@ -458,15 +458,29 @@ impl Pipeline {
         self.sinks = new_sinks;
     }
 
-    pub fn exchange<T: Exchange>(&mut self, n: usize, exchange: Arc<T>) {
+    pub fn exchange<T: Exchange>(&mut self, n: usize, exchange: Arc<T>) -> Result<()> {
+        self.exchange_with_merge(n, exchange.clone(), |inputs, output| {
+            Ok(MergePartitionProcessor::create(
+                inputs,
+                output,
+                exchange.clone(),
+            ))
+        })
+    }
+
+    pub fn exchange_with_merge<T, F>(&mut self, n: usize, exchange: Arc<T>, f: F) -> Result<()>
+    where
+        T: Exchange,
+        F: Fn(Vec<Arc<InputPort>>, Arc<OutputPort>) -> Result<ProcessorPtr>,
+    {
         if self.sinks.is_empty() {
-            return;
+            return Ok(());
         }
 
         let input_len = self.sinks.len();
         let mut items = Vec::with_capacity(input_len);
 
-        for _index in 0..input_len {
+        for _ in 0..input_len {
             let input = InputPort::create();
             let outputs: Vec<_> = (0..n).map(|_| OutputPort::create()).collect();
             items.push(PipeItem::create(
@@ -491,14 +505,15 @@ impl Pipeline {
             let output = OutputPort::create();
             let inputs: Vec<_> = (0..input_len).map(|_| InputPort::create()).collect();
             items.push(PipeItem::create(
-                MergePartitionProcessor::create(inputs.clone(), output.clone(), exchange.clone()),
+                f(inputs.clone(), output.clone())?,
                 inputs,
                 vec![output],
             ));
         }
 
         // merge partition
-        self.add_pipe(Pipe::create(input_len * n, n, items))
+        self.add_pipe(Pipe::create(input_len * n, n, items));
+        Ok(())
     }
 
     #[track_caller]

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

@@ -18,10 +18,12 @@ mod transform_accumulating;
 mod transform_accumulating_async;
 mod transform_async;
 mod transform_blocking;
+mod transform_blocking_async;
 mod transform_compact_block;
 mod transform_compact_builder;
 mod transform_compact_no_split_builder;
 mod transform_dummy;
+mod transform_hook;
 mod transform_k_way_merge_sort;
 mod transform_multi_sort_merge;
 mod transform_pipeline_helper;
@@ -36,12 +38,14 @@ pub use transform_accumulating::*;
 pub use transform_accumulating_async::*;
 pub use transform_async::*;
 pub use transform_blocking::*;
+pub use transform_blocking_async::*;
 pub use transform_compact_block::*;
 pub use transform_compact_builder::*;
 pub use transform_compact_no_split_builder::*;
 pub use transform_dummy::*;
+pub use transform_hook::*;
 pub use transform_k_way_merge_sort::*;
-pub use transform_multi_sort_merge::try_add_multi_sort_merge;
+pub use transform_multi_sort_merge::*;
 pub use transform_pipeline_helper::TransformPipelineHelper;
 pub use transform_retry_async::*;
 pub use transform_sort_merge::sort_merge;

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

@@ -58,6 +58,17 @@ impl Rows for BinaryColumn {
     fn slice(&self, range: Range<usize>) -> Self {
         self.slice(range)
     }
+
+    fn scalar_as_item<'a>(s: &'a Scalar) -> Self::Item<'a> {
+        match s {
+            Scalar::Binary(s) => s,
+            _ => unreachable!(),
+        }
+    }
+
+    fn owned_item(item: Self::Item<'_>) -> Scalar {
+        Scalar::Binary(Vec::from(item))
+    }
 }
 
 impl RowConverter<BinaryColumn> for CommonRowConverter {

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

@@ -17,6 +17,7 @@ mod simple;
 mod utils;
 
 use std::fmt::Debug;
+use std::ops::Range;
 
 pub use common::*;
 use databend_common_exception::ErrorCode;
@@ -25,7 +26,9 @@ use databend_common_expression::types::ArgType;
 use databend_common_expression::types::DataType;
 use databend_common_expression::BlockEntry;
 use databend_common_expression::Column;
+use databend_common_expression::DataBlock;
 use databend_common_expression::DataSchemaRef;
+use databend_common_expression::Scalar;
 use databend_common_expression::SortColumnDescription;
 pub use simple::*;
 pub use utils::*;
@@ -39,6 +42,18 @@ where Self: Sized + Debug
         output_schema: DataSchemaRef,
     ) -> Result<Self>;
     fn convert(&self, columns: &[BlockEntry], num_rows: usize) -> Result<T>;
+
+    fn convert_data_block(
+        &self,
+        sort_desc: &[SortColumnDescription],
+        data_block: &DataBlock,
+    ) -> Result<T> {
+        let order_by_cols = sort_desc
+            .iter()
+            .map(|desc| data_block.get_by_offset(desc.offset).clone())
+            .collect::<Vec<_>>();
+        self.convert(&order_by_cols, data_block.num_rows())
+    }
 }
 
 /// Rows can be compared.
@@ -82,5 +97,9 @@ where Self: Sized + Clone + Debug + Send
         self.row(self.len() - 1)
     }
 
-    fn slice(&self, range: std::ops::Range<usize>) -> Self;
+    fn slice(&self, range: Range<usize>) -> Self;
+
+    fn scalar_as_item<'a>(s: &'a Scalar) -> Self::Item<'a>;
+
+    fn owned_item(item: Self::Item<'_>) -> Scalar;
 }

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

@@ -23,6 +23,7 @@ use databend_common_expression::types::ValueType;
 use databend_common_expression::BlockEntry;
 use databend_common_expression::Column;
 use databend_common_expression::DataSchemaRef;
+use databend_common_expression::Scalar;
 use databend_common_expression::SortColumnDescription;
 
 use super::RowConverter;
@@ -68,6 +69,15 @@ where
             inner: T::slice_column(&self.inner, range),
         }
     }
+
+    fn scalar_as_item<'a>(s: &'a Scalar) -> Self::Item<'a> {
+        let s = &s.as_ref();
+        T::try_downcast_scalar(s).unwrap()
+    }
+
+    fn owned_item(item: Self::Item<'_>) -> Scalar {
+        T::upcast_scalar(T::to_owned_scalar(item))
+    }
 }
 
 /// Rows structure for single simple types. (numbers, date, timestamp)
@@ -111,6 +121,15 @@ where
             inner: T::slice_column(&self.inner, range),
         }
     }
+
+    fn scalar_as_item<'a>(s: &'a Scalar) -> Self::Item<'a> {
+        let s = &s.as_ref();
+        Reverse(T::try_downcast_scalar(s).unwrap())
+    }
+
+    fn owned_item(item: Self::Item<'_>) -> Scalar {
+        T::upcast_scalar(T::to_owned_scalar(item.0))
+    }
 }
 
 /// If there is only one sort field and its type is a primitive type,