update

Signed-off-by: coldWater <forsaken628@gmail.com>

Author: forsaken628

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

@@ -31,6 +31,10 @@ pub struct Bounds(
 );
 
 impl Bounds {
+    pub fn new_unchecked(column: Column) -> Bounds {
+        Bounds(vec![column])
+    }
+
     pub fn from_column<R: Rows>(column: Column) -> Result<Bounds> {
         let block = DataBlock::sort(
             &DataBlock::new_from_columns(vec![column]),
@@ -100,12 +104,12 @@ impl Bounds {
         if n == 0 {
             return Some(Self::default());
         }
-        let count = self.len();
-        if n >= count {
+        let total = self.len();
+        if n >= total {
             return None;
         }
 
-        let step = count / n;
+        let step = total / n;
         let offset = step / 2;
         let indices = self
             .0
@@ -131,6 +135,56 @@ impl Bounds {
             indices.len(),
         )]))
     }
+
+    pub fn dedup_reduce<R: Rows>(&self, n: usize) -> Self {
+        if n == 0 {
+            return Self::default();
+        }
+        let total = self.len();
+        let mut step = total as f64 / n as f64;
+        let mut target = step / 2.0;
+        let mut indices = Vec::with_capacity(n);
+        let mut last: Option<(R, _)> = None;
+        for (i, (b_idx, r_idx)) in self
+            .0
+            .iter()
+            .enumerate()
+            .rev()
+            .flat_map(|(b_idx, col)| std::iter::repeat_n(b_idx, col.len()).zip(0..col.len()))
+            .enumerate()
+        {
+            if indices.len() >= n {
+                break;
+            }
+            if (i as f64) < target {
+                continue;
+            }
+
+            let cur_rows = R::from_column(&self.0[b_idx]).unwrap();
+            if last
+                .as_ref()
+                .map(|(last_rows, last_idx)| cur_rows.row(r_idx) == last_rows.row(*last_idx))
+                .unwrap_or_default()
+            {
+                continue;
+            }
+
+            indices.push((b_idx as u32, r_idx as u32, 1));
+            target += step;
+            if (i as f64) > target && indices.len() < n {
+                step = (total - i) as f64 / (n - indices.len()) as f64;
+                target = i as f64 + step / 2.0;
+            }
+            last = Some((cur_rows, r_idx));
+        }
+
+        Bounds(vec![Column::take_column_indices(
+            &self.0,
+            R::data_type(),
+            &indices,
+            indices.len(),
+        )])
+    }
 }
 
 impl SortedStream for Bounds {
@@ -233,4 +287,35 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_dedup_reduce() -> Result<()> {
+        let column = Int32Type::from_data(vec![1, 2, 2, 3, 3, 3, 4, 5, 5]);
+        let bounds = Bounds::new_unchecked(column);
+        let reduced = bounds.dedup_reduce::<SimpleRowsAsc<Int32Type>>(3);
+        assert_eq!(reduced, Bounds(vec![Int32Type::from_data(vec![2, 3, 5])]));
+
+        let column = Int32Type::from_data(vec![5, 5, 4, 3, 3, 3, 2, 2, 1]);
+        let bounds = Bounds::new_unchecked(column);
+        let reduced = bounds.dedup_reduce::<SimpleRowsDesc<Int32Type>>(3);
+        assert_eq!(reduced, Bounds(vec![Int32Type::from_data(vec![4, 3, 1])]));
+
+        let bounds_vec = [vec![5, 6, 7, 7], vec![3, 3, 4, 5], vec![1, 2, 2, 3]]
+            .into_iter()
+            .map(|v| Int32Type::from_data(v))
+            .collect::<Vec<_>>();
+        let bounds = Bounds(bounds_vec);
+        let reduced = bounds.dedup_reduce::<SimpleRowsAsc<Int32Type>>(5);
+        assert_eq!(
+            reduced,
+            Bounds(vec![Int32Type::from_data(vec![2, 3, 4, 6, 7])])
+        );
+
+        let column = Int32Type::from_data(vec![1, 1, 1, 1, 1]);
+        let bounds = Bounds(vec![column]);
+        let reduced = bounds.dedup_reduce::<SimpleRowsAsc<Int32Type>>(3);
+        assert_eq!(reduced, Bounds(vec![Int32Type::from_data(vec![1])]));
+
+        Ok(())
+    }
 }

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

@@ -191,7 +191,7 @@ impl<R: Rows + 'static> Processor for TransformSortShuffle<R> {
     async fn async_process(&mut self) -> Result<()> {
         let bounds = match &self.step {
             Step::None if self.input.is_finished() => Bounds::default(),
-            Step::Meta(meta) => meta.bounds.clone(),
+            Step::Meta(meta) => meta.generate_bounds(),
             _ => unreachable!(),
         };
         self.state.commit_sample::<R>(self.id, bounds)?;
@@ -201,6 +201,27 @@ impl<R: Rows + 'static> Processor for TransformSortShuffle<R> {
     }
 }
 
+impl SortCollectedMeta {
+    fn generate_bounds(&self) -> Bounds {
+        if self.bounds.len() > 1 {
+            return self.bounds.clone();
+        }
+
+        let Some(blocks) = self.blocks.get(self.blocks.len() / 2) else {
+            return Bounds::default();
+        };
+
+        blocks
+            .get(blocks.len() / 2)
+            .map(|block| match block.domain.len() {
+                0 => Bounds::default(),
+                1 => Bounds::new_unchecked(block.domain.clone()),
+                _ => Bounds::new_unchecked(block.domain.slice(0..1)),
+            })
+            .unwrap_or_default()
+    }
+}
+
 pub struct SortSampleState {
     inner: RwLock<StateInner>,
     pub(super) done: WatchNotify,
@@ -261,9 +282,13 @@ impl StateInner {
     fn determine_bounds<R: Rows>(&mut self) -> Result<()> {
         let v = self.partial.drain(..).map(Option::unwrap).collect();
         let bounds = Bounds::merge::<R>(v, self.batch_rows)?;
-        let bounds = bounds
-            .reduce(self.partitions - 1, R::data_type())
-            .unwrap_or(bounds);
+
+        let n = self.partitions - 1;
+        let bounds = if bounds.len() < n {
+            bounds
+        } else {
+            bounds.dedup_reduce::<R>(n)
+        };
         assert!(bounds.len() < self.partitions);
 
         self.bounds = Some(bounds);

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

@@ -596,7 +596,7 @@ pub struct SpillableBlock {
     data: Option<DataBlock>,
     rows: usize,
     location: Option<Location>,
-    domain: Column,
+    pub(super) domain: Column,
     processed: usize,
 }
 

tests/sqllogictests/suites/mode/standalone/explain/window.test

@@ -47,6 +47,9 @@ set sort_spilling_memory_ratio = 0;
 statement ok
 set enable_parallel_multi_merge_sort = 0;
 
+statement ok
+set enable_range_shuffle_sort = 0;
+
 query T
 explain pipeline SELECT depname, empno, salary, sum(salary) OVER (PARTITION BY depname ORDER BY empno) FROM empsalary ORDER BY depname, empno;
 ----

tests/sqllogictests/suites/stage/formats/parquet/read_policy.test

@@ -135,7 +135,7 @@ select t, t:a from @data/parquet/tuple.parquet order by id desc limit 2;
 
 # topk contains output
 query TT
-select id, t:b, t:a from @data/parquet/tuple.parquet order by t:a desc limit 2;
+select id, t:b, t:a from @data/parquet/tuple.parquet order by t:a desc, id desc limit 2;
 ----
 3 c 3
 2 b 3