bound_index

Author: forsaken628

Cargo.lock

@@ -171,7 +171,6 @@ tokio = { workspace = true }
 tokio-stream = { workspace = true, features = ["net"] }
 toml = { workspace = true, features = ["parse"] }
 tonic = { workspace = true }
-twox-hash = { workspace = true }
 typetag = { workspace = true }
 url = { workspace = true }
 uuid = { workspace = true }

src/query/service/Cargo.toml

@@ -20,7 +20,6 @@ use databend_common_expression::BlockMetaInfo;
 use databend_common_expression::BlockMetaInfoDowncast;
 use databend_common_expression::DataBlock;
 use databend_common_expression::DataSchemaRef;
-use databend_common_expression::Scalar;
 use databend_common_pipeline_transforms::SortSpillParams;
 use sort_spill::SpillableBlock;
 
@@ -107,7 +106,14 @@ impl BlockMetaInfo for SortExchangeMeta {
 
 #[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
 pub struct SortBound {
-    bound: Option<Scalar>,
+    bound_index: u32,
+}
+
+impl SortBound {
+    fn create(bound_index: u32) -> Box<dyn BlockMetaInfo> {
+        debug_assert!(bound_index != u32::MAX);
+        SortBound { bound_index }.boxed()
+    }
 }
 
 #[typetag::serde(name = "sort_bound")]

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

@@ -12,17 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::hash::Hash;
-use std::hash::Hasher;
 use std::sync::Arc;
 
 use databend_common_exception::Result;
 use databend_common_expression::BlockMetaInfoDowncast;
 use databend_common_expression::DataBlock;
-use databend_common_expression::Scalar;
 use databend_common_pipeline_core::Pipeline;
 use databend_common_settings::FlightCompression;
-use twox_hash::XxHash64;
 
 use crate::pipelines::processors::transforms::SortBound;
 use crate::servers::flight::v1::exchange::DataExchange;
@@ -88,7 +84,7 @@ impl ExchangeInjector for SortInjector {
 }
 
 pub struct SortBoundScatter {
-    partitions: u64,
+    partitions: u32,
 }
 
 impl FlightScatter for SortBoundScatter {
@@ -101,23 +97,13 @@ impl FlightScatter for SortBoundScatter {
     }
 }
 
-pub(super) fn bound_scatter(data_block: DataBlock, n: u64) -> Result<Vec<DataBlock>> {
+pub(super) fn bound_scatter(data_block: DataBlock, n: u32) -> Result<Vec<DataBlock>> {
     let meta = data_block
         .get_meta()
         .and_then(SortBound::downcast_ref_from)
         .unwrap();
 
-    let bound = match &meta.bound {
-        Some(bound) => {
-            debug_assert!(!bound.is_null());
-            bound
-        }
-        None => &Scalar::Null,
-    };
-
-    let mut hasher = XxHash64::default();
-    bound.hash(&mut hasher);
-    let index = hasher.finish() % n;
+    let index = meta.bound_index % n;
 
     Ok(std::iter::repeat_n(DataBlock::empty(), index as _)
         .chain(Some(data_block))

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

@@ -13,7 +13,6 @@
 // limitations under the License.
 
 use std::any::Any;
-use std::cmp::Ordering;
 use std::marker::PhantomData;
 use std::sync::Arc;
 
@@ -22,7 +21,6 @@ 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;
@@ -46,7 +44,7 @@ where A: SortAlgorithm
     limit: Option<usize>,
 
     output_data: Option<DataBlock>,
-    bound: Option<Scalar>,
+    bound_index: u32,
     inner: std::result::Result<Merger<A, Stream<A>>, Vec<Stream<A>>>,
 }
 
@@ -67,7 +65,7 @@ where A: SortAlgorithm
                 data: None,
                 input: input.clone(),
                 remove_order_col,
-                bound: None,
+                bound_index: None,
                 sort_row_offset: schema.fields().len() - 1,
                 _r: PhantomData,
             })
@@ -80,7 +78,7 @@ where A: SortAlgorithm
             block_size,
             limit,
             output_data: None,
-            bound: None,
+            bound_index: u32::MAX,
             inner: Err(streams),
         })
     }
@@ -119,11 +117,7 @@ where A: SortAlgorithm + 'static
 
     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))?);
+            self.output_data = Some(block.add_meta(Some(SortBound::create(self.bound_index)))?);
         };
         Ok(())
     }
@@ -158,20 +152,17 @@ where A: SortAlgorithm + 'static
                 .get_meta()
                 .and_then(SortBound::downcast_ref_from)
                 .expect("require a SortBound");
-            bounds.push(&meta.bound)
+            bounds.push(meta.bound_index)
         }
 
-        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(),
+        let bound_index = match bounds.iter().min() {
+            Some(index) => *index,
             None => return Ok(Event::Finished),
         };
+        assert!(self.bound_index != u32::MAX || bound_index > self.bound_index);
+        self.bound_index = bound_index;
         for stream in streams.iter_mut() {
-            stream.bound = self.bound.clone();
+            stream.bound_index = Some(self.bound_index);
         }
 
         let streams = std::mem::take(streams);
@@ -189,7 +180,7 @@ struct BoundedInputStream<R: Rows> {
     data: Option<DataBlock>,
     input: Arc<InputPort>,
     remove_order_col: bool,
-    bound: Option<Scalar>,
+    bound_index: Option<u32>,
     sort_row_offset: usize,
     _r: PhantomData<R>,
 }
@@ -244,9 +235,9 @@ impl<R: Rows> BoundedInputStream<R> {
             .and_then(SortBound::downcast_ref_from)
             .expect("require a SortBound");
 
-        if meta.bound == self.bound {
+        if meta.bound_index == self.bound_index.unwrap() {
             self.data.take().map(|mut data| {
-                let _ = data.take_meta();
+                data.take_meta().unwrap();
                 data
             })
         } else {

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

@@ -115,11 +115,11 @@ where
 
         let OutputData {
             block,
-            bound,
+            bound: (bound_index, _),
             finish,
         } = spill_sort.on_restore().await?;
         if let Some(block) = block {
-            let mut block = block.add_meta(Some(SortBound { bound }.boxed()))?;
+            let mut block = block.add_meta(Some(SortBound::create(bound_index)))?;
             if self.remove_order_col {
                 block.pop_columns(1);
             }

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

@@ -68,6 +68,7 @@ struct StepSort<A: SortAlgorithm> {
     /// Each boundary represents a cutoff point where data less than or equal to it belongs to one partition.
     bounds: Bounds,
     cur_bound: Option<Scalar>,
+    bound_index: i32,
 
     subsequent: Vec<BoundBlockStream<A::Rows, Arc<Spiller>>>,
     current: Vec<BoundBlockStream<A::Rows, Arc<Spiller>>>,
@@ -116,6 +117,7 @@ where A: SortAlgorithm
                 params,
                 bounds,
                 cur_bound: None,
+                bound_index: -1,
                 subsequent,
                 current: vec![],
                 output_merger: None,
@@ -172,7 +174,7 @@ where A: SortAlgorithm
             sort.merge_current(&self.base).await?;
             Ok(OutputData {
                 block: None,
-                bound: None,
+                bound: (u32::MAX, None),
                 finish: false,
             })
         } else {
@@ -303,6 +305,7 @@ impl<A: SortAlgorithm> StepCollect<A> {
         Ok(StepSort {
             bounds,
             cur_bound: None,
+            bound_index: -1,
             subsequent: std::mem::take(&mut self.streams),
             current: vec![],
             output_merger: None,
@@ -313,7 +316,7 @@ impl<A: SortAlgorithm> StepCollect<A> {
 
 pub struct OutputData {
     pub block: Option<DataBlock>,
-    pub bound: Option<Scalar>,
+    pub bound: (u32, Option<Scalar>),
     pub finish: bool,
 }
 
@@ -323,6 +326,7 @@ impl<A: SortAlgorithm> StepSort<A> {
             Some(bound) => self.cur_bound = Some(bound),
             None => self.cur_bound = None,
         }
+        self.bound_index += 1;
     }
 
     async fn merge_current(&mut self, base: &Base) -> Result<()> {
@@ -372,7 +376,8 @@ impl<A: SortAlgorithm> StepSort<A> {
                     let mut s = self.current.pop().unwrap();
                     s.restore_first().await?;
                     let block = Some(s.take_next_bounded_block());
-                    let bound = s.bound.clone();
+                    assert!(self.bound_index >= 0);
+                    let bound = (self.bound_index as _, s.bound.clone());
 
                     if !s.is_empty() {
                         if s.should_include_first() {
@@ -415,15 +420,16 @@ impl<A: SortAlgorithm> StepSort<A> {
 
             return Ok(OutputData {
                 block: None,
-                bound: None,
+                bound: (u32::MAX, None),
                 finish: self.subsequent.is_empty(),
             });
         };
 
         let mut sorted = base.new_stream([base.new_block(data)].into(), self.cur_bound.clone());
         let (block, bound) = if sorted.should_include_first() {
             let block = Some(sorted.take_next_bounded_block());
-            let bound = sorted.bound.clone();
+            debug_assert!(self.bound_index >= 0);
+            let bound = (self.bound_index as _, sorted.bound.clone());
             if sorted.is_empty() {
                 return Ok(OutputData {
                     block,
@@ -433,7 +439,7 @@ impl<A: SortAlgorithm> StepSort<A> {
             }
             (block, bound)
         } else {
-            (None, None)
+            (None, (u32::MAX, None))
         };
 
         while let Some(data) = merger.async_next_block().await? {