feat(query): support shuffle sort (#17853)

* reservoir_sampling

Signed-off-by: coldWater <[email protected]>

* Simpler

Signed-off-by: coldWater <[email protected]>

* TransformSortSimple

Signed-off-by: coldWater <[email protected]>

* rename

* fix

* trait Spill

* execute

* move

* update

* bounds

* wait

* remove

* scalar

* exchange

* update

* test

Signed-off-by: coldWater <[email protected]>

* x

Signed-off-by: coldWater <[email protected]>

* route

Signed-off-by: coldWater <[email protected]>

* builder

Signed-off-by: coldWater <[email protected]>

* update

Signed-off-by: coldWater <[email protected]>

* build

Signed-off-by: coldWater <[email protected]>

* rename

Signed-off-by: coldWater <[email protected]>

* fix

Signed-off-by: coldWater <[email protected]>

* fix

Signed-off-by: coldWater <[email protected]>

* fix

Signed-off-by: coldWater <[email protected]>

* fix

Signed-off-by: coldWater <[email protected]>

* fix

Signed-off-by: coldWater <[email protected]>

* update

Signed-off-by: coldWater <[email protected]>

* fix

Signed-off-by: coldWater <[email protected]>

* fix

Signed-off-by: coldWater <[email protected]>

* fix

Signed-off-by: coldWater <[email protected]>

* x

* logical plan

* x

* x

* x

* x

* x

* x

* x

* x

* x

* x

* x

* format

* x

* x

* allow_adjust_parallelism

* InputBoundStream

* BoundedMultiSortMergeProcessor

* build_bounded_merge_sort

* bound_index

* SortBoundEdge

* route

* clean up

* enable

* fix

* fix

* fix

* test

* fix

* fix

* fix

* fix

* fix

* fix

* force_disable_distributed_optimization

* trace and refine

* default_num_merge

---------

Signed-off-by: coldWater <[email protected]>

Author: forsaken628

.github/workflows/reuse.sqllogic.yml

@@ -219,8 +219,8 @@ jobs:
         tests:
           - { dirs: "query", runner: "4c16g" }
           - { dirs: "duckdb", runner: "4c16g" }
-          - { dirs: "crdb", runner: "2c8g" }
-          - { dirs: "base", runner: "2c8g" }
+          - { dirs: "crdb", runner: "2c8g", parallel: "2" }
+          - { dirs: "base", runner: "2c8g", parallel: "2" }
           - { dirs: "ydb", runner: "2c8g" }
           - { dirs: "tpcds", runner: "2c8g" }
           - { dirs: "tpch", runner: "2c8g" }

Cargo.lock

@@ -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/common/base/src/base/watch_notify.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/block.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/array.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/boolean.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/expression/src/types/map.rs

@@ -460,15 +460,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(
@@ -493,14 +507,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/core/src/pipeline.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/mod.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 {