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/types/array.rs

@@ -52,7 +52,7 @@ impl<T: ValueType> ValueType 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

@@ -48,7 +48,7 @@ impl ValueType 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/date.rs

@@ -78,7 +78,7 @@ impl ValueType for DateType {
         *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::Date(scalar) => Some(*scalar),
             _ => None,

src/query/expression/src/types/decimal.rs

@@ -72,7 +72,7 @@ impl<Num: Decimal> ValueType for DecimalType<Num> {
         *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>> {
         Num::try_downcast_scalar(scalar.as_decimal()?)
     }
 

src/query/expression/src/types/empty_array.rs

@@ -45,7 +45,7 @@ impl ValueType for EmptyArrayType {
         *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::EmptyArray => Some(()),
             _ => None,

src/query/expression/src/types/empty_map.rs

@@ -45,7 +45,7 @@ impl ValueType for EmptyMapType {
         *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::EmptyMap => Some(()),
             _ => None,

src/query/expression/src/types/interval.rs

@@ -52,7 +52,7 @@ impl ValueType for IntervalType {
         *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::Interval(scalar) => Some(*scalar),
             _ => None,

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

@@ -341,7 +341,7 @@ impl<K: ValueType, V: ValueType> ValueType for MapType<K, V> {
         <MapInternal<K, V> as ValueType>::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/expression/src/types/null.rs

@@ -46,7 +46,7 @@ impl ValueType for NullType {
         *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::Null => Some(()),
             _ => None,

src/query/expression/src/types/number.rs

@@ -117,7 +117,7 @@ impl<Num: Number> ValueType for NumberType<Num> {
         *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>> {
         Num::try_downcast_scalar(scalar.as_number()?)
     }
 

src/query/expression/src/types/timestamp.rs

@@ -83,7 +83,7 @@ impl ValueType for TimestampType {
         *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::Timestamp(scalar) => Some(*scalar),
             _ => None,

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

@@ -41,7 +41,7 @@ pub use transform_compact_builder::*;
 pub use transform_compact_no_split_builder::*;
 pub use transform_dummy::*;
 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

@@ -60,6 +60,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

@@ -24,6 +24,7 @@ use databend_common_expression::BlockEntry;
 use databend_common_expression::Column;
 use databend_common_expression::ColumnBuilder;
 use databend_common_expression::DataSchemaRef;
+use databend_common_expression::Scalar;
 use databend_common_expression::SortColumnDescription;
 use databend_common_expression::Value;
 
@@ -70,6 +71,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)
@@ -113,6 +123,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,