diff --git a/Cargo.lock b/Cargo.lock index 4bd77dd53fbf2..41229298b0dc7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5239,13 +5239,11 @@ dependencies = [ "databend-common-meta-app", "databend-common-pipeline", "databend-common-pipeline-transforms", - "databend-common-storage", "databend-common-users", "databend-meta-client 260205.13.1", "databend-storages-common-table-meta", "databend_educe", "futures", - "log", "paimon", "pretty_assertions", "serde", diff --git a/src/query/pipeline/transforms/src/processors/transforms/sorts/core/algorithm.rs b/src/query/pipeline/transforms/src/processors/transforms/sorts/core/algorithm.rs index c95d4d5eb6937..31bf240a8bf11 100644 --- a/src/query/pipeline/transforms/src/processors/transforms/sorts/core/algorithm.rs +++ b/src/query/pipeline/transforms/src/processors/transforms/sorts/core/algorithm.rs @@ -14,7 +14,6 @@ use core::fmt; use std::cmp::Ordering; -use std::cmp::Reverse; use std::collections::BinaryHeap; use std::collections::binary_heap; use std::ops::Deref; @@ -26,12 +25,31 @@ use super::Rows; use super::loser_tree; use super::utils::find_bigger_child_of_root; -pub type Cursor = RawCursor; - +pub type Cursor = RawCursor<'static, R, ItemCursorOrder>; + +/// Backend for the k-way merge cursor set. +/// +/// Each input contributes one cursor positioned at its current row. The merge +/// loop performs four steps: +/// +/// 1. Peek the cursor with the smallest current row. +/// 2. Find a consecutive range from that input which can be emitted without +/// another cursor becoming the winner. +/// 3. Emit the range and advance its cursor. +/// 4. Remove an exhausted cursor, or repair the cursor set after its key changes. +/// +/// The loop repeats once per drained range, not once per output row. Its cost can +/// be summarized as: +/// +/// `row work + R * per-range work` +/// +/// `HeapSort` and `LoserTreeTop2Sort` expose the second cursor, allowing the merger to +/// drain up to that key. `LoserTreeSort` keeps the winner in the tree and +/// only coalesces equal keys because that layout does not expose the second cursor. pub trait SortAlgorithm: Send { - const SHOULD_PEEK_TOP2: bool; + const SHOULD_PEEK_TOP2: bool = true; type Rows: Rows; - type PeekMut<'b>: Deref>> + DerefMut + type PeekMut<'b>: Deref> + DerefMut where Self: 'b; fn with_capacity(capacity: usize) -> Self; @@ -43,28 +61,27 @@ pub trait SortAlgorithm: Send { fn rebuild(&mut self); - fn push(&mut self, index: usize, item: Reverse>); + fn push(&mut self, index: usize, item: Cursor); fn pop(&mut self); - fn update_top(&mut self, item: Reverse>); + fn update_top(&mut self, item: Cursor); - fn peek(&self) -> Option<&Reverse>>; + fn peek(&self) -> Option<&Cursor>; - fn peek_top2(&self) -> &Reverse>; + fn peek_top2(&self) -> &Cursor; fn peek_mut(&mut self) -> Self::PeekMut<'_>; fn pop_mut(this: Self::PeekMut<'_>); } -pub type HeapSort = BinaryHeap>>; +pub type HeapSort = BinaryHeap>; -impl SortAlgorithm for BinaryHeap>> { - const SHOULD_PEEK_TOP2: bool = true; +impl SortAlgorithm for BinaryHeap> { type Rows = R; type PeekMut<'a> - = binary_heap::PeekMut<'a, Reverse>> + = binary_heap::PeekMut<'a, Cursor> where R: 'a; fn with_capacity(capacity: usize) -> Self { BinaryHeap::with_capacity(capacity) @@ -76,7 +93,7 @@ impl SortAlgorithm for BinaryHeap>> { fn rebuild(&mut self) {} - fn push(&mut self, _index: usize, item: Reverse>) { + fn push(&mut self, _index: usize, item: Cursor) { BinaryHeap::push(self, item) } @@ -84,17 +101,17 @@ impl SortAlgorithm for BinaryHeap>> { BinaryHeap::pop(self); } - fn update_top(&mut self, item: Reverse>) { + fn update_top(&mut self, item: Cursor) { // `peek_mut` will return a `PeekMut` object which allows us to modify the top element of the heap. // The heap will adjust itself automatically when the `PeekMut` object is dropped (RAII). *BinaryHeap::peek_mut(self).unwrap() = item } - fn peek(&self) -> Option<&Reverse>> { + fn peek(&self) -> Option<&Cursor> { BinaryHeap::peek(self) } - fn peek_top2(&self) -> &Reverse> { + fn peek_top2(&self) -> &Cursor { find_bigger_child_of_root(self) } @@ -107,38 +124,276 @@ impl SortAlgorithm for BinaryHeap>> { } } +pub struct LoserTreeTop2Sort { + top: Option>, + tree: loser_tree::LoserTree>, + exhausted_input: Option, + staged: Vec<(usize, Cursor)>, + length: usize, +} + +impl fmt::Debug for LoserTreeTop2Sort { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let tree = self + .tree + .nodes() + .iter() + .map(|node| { + node.as_ref() + .map(|(slot, cursor)| (*slot, cursor.row_index)) + }) + .collect::>(); + let staged = self + .staged + .iter() + .map(|(index, cursor)| (*index, cursor.row_index)) + .collect::>(); + + f.debug_struct("LoserTreeTop2Sort") + .field("top", &self.top.as_ref().map(|cursor| cursor.row_index)) + .field("tree", &tree) + .field("exhausted_input", &self.exhausted_input) + .field("staged", &staged) + .field("length", &self.length) + .finish() + } +} + +impl LoserTreeTop2Sort { + #[inline(always)] + fn reconcile_top(&mut self) { + if self.exhausted_input.is_some() || !self.staged.is_empty() { + return; + } + if let Some(top) = self.top + && self.tree.peek().is_some_and(|second| top < *second) + { + self.top = Some(self.tree.promote_with(Some(top))); + } + } + + #[inline(always)] + fn pop_top(&mut self) { + let top = self.top.take().expect("loser tree sort must not be empty"); + debug_assert!(self.exhausted_input.is_none()); + self.exhausted_input = Some(top.input_index); + self.length -= 1; + } + + fn rebuild_all(&mut self) { + let mut values = Vec::with_capacity(self.length); + values.extend(self.tree.values()); + values.extend(self.staged.drain(..).map(|(_, cursor)| cursor)); + + if self.top.is_none() { + self.top = values.pop(); + } + debug_assert_eq!(values.len() + usize::from(self.top.is_some()), self.length); + self.tree.rebuild_from(values); + self.reconcile_top(); + } + + fn prepare_for_read(&mut self) { + let Some(exhausted_input) = self.exhausted_input.take() else { + if !self.staged.is_empty() { + self.rebuild_all(); + } + return; + }; + + // Refill normally stages at most one cursor, from the stream whose + // previous block just exhausted. Replay that replacement on the tree + // winner's path while promoting the winner to the external top. + let replacement = match self.staged.as_slice() { + [] => None, + [(index, cursor)] if *index == exhausted_input => Some(*cursor), + _ => { + self.rebuild_all(); + return; + } + }; + self.staged.clear(); + + self.top = if self.tree.peek().is_some() { + Some(self.tree.promote_with(replacement)) + } else { + replacement + }; + self.reconcile_top(); + } +} + +impl SortAlgorithm for LoserTreeTop2Sort { + type Rows = R; + type PeekMut<'a> + = LoserTreeTop2PeekMut<'a, Self::Rows> + where Self: 'a; + fn with_capacity(capacity: usize) -> Self { + LoserTreeTop2Sort { + top: None, + tree: loser_tree::LoserTree::with_capacity(capacity.saturating_sub(1)), + exhausted_input: None, + staged: Vec::with_capacity(capacity), + length: 0, + } + } + + fn len(&self) -> usize { + self.length + } + + fn rebuild(&mut self) { + self.prepare_for_read(); + } + + fn push(&mut self, index: usize, item: Cursor) { + self.length += 1; + self.staged.push((index, item)); + } + + fn pop(&mut self) { + self.pop_top(); + self.prepare_for_read(); + } + + fn update_top(&mut self, item: Cursor) { + self.top = Some(item); + self.reconcile_top(); + } + + fn peek(&self) -> Option<&Cursor> { + debug_assert!(self.exhausted_input.is_none() && self.staged.is_empty()); + self.top.as_ref() + } + + fn peek_top2(&self) -> &Cursor { + debug_assert!(self.exhausted_input.is_none() && self.staged.is_empty()); + self.tree + .peek() + .expect("top2 requires at least two cursors") + } + + fn peek_mut(&mut self) -> Self::PeekMut<'_> { + LoserTreeTop2PeekMut(self) + } + + fn pop_mut(this: Self::PeekMut<'_>) { + this.0.pop_top(); + } +} + +pub struct LoserTreeTop2PeekMut<'a, R: Rows>(&'a mut LoserTreeTop2Sort); + +impl Deref for LoserTreeTop2PeekMut<'_, R> { + type Target = Cursor; + + fn deref(&self) -> &Self::Target { + self.0.top.as_ref().unwrap() + } +} + +impl DerefMut for LoserTreeTop2PeekMut<'_, R> { + fn deref_mut(&mut self) -> &mut Self::Target { + self.0.top.as_mut().unwrap() + } +} + +impl Drop for LoserTreeTop2PeekMut<'_, R> { + fn drop(&mut self) { + self.0.reconcile_top(); + } +} + +/// A loser tree which keeps the winner in the tree itself. +/// +/// Unlike [`LoserTreeTop2Sort`], this layout cannot expose the second cursor without +/// searching the winner path. The merger therefore only coalesces equal keys for +/// this implementation. pub struct LoserTreeSort { - tree: loser_tree::LoserTree>>>, + tree: loser_tree::LoserTree>, + exhausted: Option<(usize, usize)>, + staged: Vec<(usize, Cursor)>, length: usize, } impl fmt::Debug for LoserTreeSort { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - let data = self + let tree = self .tree - .data() + .nodes() + .iter() + .map(|node| { + node.as_ref() + .map(|(slot, cursor)| (*slot, cursor.row_index)) + }) + .collect::>(); + let staged = self + .staged .iter() - .map(|x| x.as_ref().map(|Reverse(cursor)| cursor.row_index)) + .map(|(index, cursor)| (*index, cursor.row_index)) .collect::>(); f.debug_struct("LoserTreeSort") - .field("tree", &self.tree.tree()) - .field("data", &data) + .field("tree", &tree) + .field("exhausted", &self.exhausted) + .field("staged", &staged) .field("length", &self.length) .finish() } } +impl LoserTreeSort { + fn rebuild_all(&mut self) { + let mut values = Vec::with_capacity(self.length); + values.extend(self.tree.values()); + values.extend(self.staged.drain(..).map(|(_, cursor)| cursor)); + debug_assert_eq!(values.len(), self.length); + self.tree.rebuild_from(values); + } + + #[inline(always)] + fn pop_top(&mut self) { + let (slot, top) = self.tree.take_top(); + debug_assert!(self.exhausted.is_none()); + self.exhausted = Some((top.input_index, slot)); + self.length -= 1; + } + + fn prepare_for_read(&mut self) { + let Some((exhausted_input, slot)) = self.exhausted.take() else { + if !self.staged.is_empty() { + self.rebuild_all(); + } + return; + }; + + let replacement = match self.staged.as_slice() { + [] => None, + [(index, cursor)] if *index == exhausted_input => Some(*cursor), + _ => { + self.tree.replace_top(slot, None); + self.rebuild_all(); + return; + } + }; + self.staged.clear(); + + self.tree.replace_top(slot, replacement); + } +} + impl SortAlgorithm for LoserTreeSort { const SHOULD_PEEK_TOP2: bool = false; type Rows = R; type PeekMut<'a> = LoserTreePeekMut<'a, Self::Rows> where Self: 'a; + fn with_capacity(capacity: usize) -> Self { - let data = vec![None; capacity]; - LoserTreeSort { - tree: loser_tree::LoserTree::from(data), + Self { + tree: loser_tree::LoserTree::with_capacity(capacity), + exhausted: None, + staged: Vec::with_capacity(capacity), length: 0, } } @@ -148,64 +403,73 @@ impl SortAlgorithm for LoserTreeSort { } fn rebuild(&mut self) { - self.tree.rebuild() + self.prepare_for_read(); } - fn push(&mut self, index: usize, item: Reverse>) { - self.tree.update(index, Some(item)); - self.length += 1 + fn push(&mut self, index: usize, item: Cursor) { + self.staged.push((index, item)); + self.length += 1; } fn pop(&mut self) { - debug_assert!(self.length >= 1); - self.tree.peek_mut().take(); - self.tree.adjust_top(); - self.length -= 1; + self.pop_top(); + self.prepare_for_read(); } - fn update_top(&mut self, item: Reverse>) { - self.tree.update(self.tree.winner(), Some(item)) + fn update_top(&mut self, item: Cursor) { + *self + .tree + .peek_mut() + .expect("loser tree sort must not be empty") = item; + self.tree.adjust_top(); } - fn peek(&self) -> Option<&Reverse>> { - self.tree.peek().as_ref() + fn peek(&self) -> Option<&Cursor> { + debug_assert!(self.exhausted.is_none() && self.staged.is_empty()); + self.tree.peek() } - fn peek_top2(&self) -> &Reverse> { - self.tree.peek_top2().as_ref().unwrap() + fn peek_top2(&self) -> &Cursor { + unreachable!("embedded loser tree does not expose top2") } fn peek_mut(&mut self) -> Self::PeekMut<'_> { - LoserTreePeekMut(self) + LoserTreePeekMut { + sort: self, + adjust_on_drop: true, + } } - fn pop_mut(this: Self::PeekMut<'_>) { - debug_assert!(this.0.length >= 1); - this.0.tree.peek_mut().take(); - this.0.length -= 1; - // The tree will adjust itself automatically when the `PeekMut` object is dropped (RAII). + fn pop_mut(mut this: Self::PeekMut<'_>) { + this.sort.pop_top(); + this.adjust_on_drop = false; } } -pub struct LoserTreePeekMut<'a, R: Rows>(&'a mut LoserTreeSort); +pub struct LoserTreePeekMut<'a, R: Rows> { + sort: &'a mut LoserTreeSort, + adjust_on_drop: bool, +} impl Deref for LoserTreePeekMut<'_, R> { - type Target = Reverse>; + type Target = Cursor; fn deref(&self) -> &Self::Target { - self.0.tree.peek().as_ref().unwrap() + self.sort.tree.peek().unwrap() } } impl DerefMut for LoserTreePeekMut<'_, R> { fn deref_mut(&mut self) -> &mut Self::Target { - self.0.tree.peek_mut().as_mut().unwrap() + self.sort.tree.peek_mut().unwrap() } } impl Drop for LoserTreePeekMut<'_, R> { fn drop(&mut self) { - self.0.tree.adjust_top(); + if self.adjust_on_drop { + self.sort.tree.adjust_top(); + } } } @@ -213,11 +477,11 @@ impl Drop for LoserTreePeekMut<'_, R> { pub struct ItemCursorOrder; impl CursorOrder for ItemCursorOrder { - fn eq(a: &RawCursor, b: &RawCursor) -> bool { + fn eq<'a>(a: &RawCursor<'a, R, Self>, b: &RawCursor<'a, R, Self>) -> bool { a.current() == b.current() } - fn cmp(a: &RawCursor, b: &RawCursor) -> Ordering { - a.current().cmp(&b.current()) + fn cmp<'a>(a: &RawCursor<'a, R, Self>, b: &RawCursor<'a, R, Self>) -> Ordering { + b.current().cmp(&a.current()) } } diff --git a/src/query/pipeline/transforms/src/processors/transforms/sorts/core/cursor.rs b/src/query/pipeline/transforms/src/processors/transforms/sorts/core/cursor.rs index 40a81517e86d0..f716258c65380 100644 --- a/src/query/pipeline/transforms/src/processors/transforms/sorts/core/cursor.rs +++ b/src/query/pipeline/transforms/src/processors/transforms/sorts/core/cursor.rs @@ -17,82 +17,87 @@ use std::marker::PhantomData; use super::rows::Rows; -/// A cursor point to a certain row in a data block. -#[derive(Clone)] -pub struct Cursor +/// A cursor pointing to a row in a sorted input stream. +#[derive(Debug, Clone)] +pub struct Cursor<'a, R, O> where R: Rows, O: CursorOrder, { pub input_index: usize, pub row_index: usize, - + num_rows: usize, + current: R::Item<'a>, + last: R::Item<'a>, _o: PhantomData, +} - /// rows within [`Cursor`] should be monotonic. - rows: R, +impl<'a, R, O> Copy for Cursor<'a, R, O> +where + R: Rows, + O: CursorOrder, +{ } -impl Cursor +impl<'a, R, O> Cursor<'a, R, O> where R: Rows, O: CursorOrder, { - pub fn new(input_index: usize, rows: R) -> Self { - O::new_cursor(input_index, rows) + pub fn new( + input_index: usize, + num_rows: usize, + current: R::Item<'a>, + last: R::Item<'a>, + ) -> Self { + debug_assert!(num_rows > 0); + Self { + input_index, + row_index: 0, + num_rows, + current, + last, + _o: PhantomData, + } } - #[inline] - pub fn advance(&mut self) -> usize { - let res = self.row_index; - self.row_index += 1; - res + pub fn advance(&mut self, count: usize, current: Option>) { + self.row_index += count; + debug_assert!(self.row_index <= self.num_rows); + debug_assert_eq!(current.is_some(), !self.is_finished()); + if let Some(current) = current { + self.current = current; + } } #[inline] pub fn is_finished(&self) -> bool { - self.rows.len() == self.row_index + self.num_rows == self.row_index } #[inline] - pub fn current(&self) -> R::Item<'_> { - self.rows.row(self.row_index) + pub fn current(&self) -> R::Item<'a> { + self.current } #[inline] - pub fn last(&self) -> R::Item<'_> { - self.rows.last() + pub fn last(&self) -> R::Item<'a> { + self.last } #[inline] pub fn num_rows(&self) -> usize { - self.rows.len() - } - - pub fn cursor_mut(&self) -> CursorMut<'_, R, O> { - CursorMut { - row_index: self.row_index, - cursor: self, - } + self.num_rows } } pub trait CursorOrder: Sized + Copy { - fn eq(a: &Cursor, b: &Cursor) -> bool; - - fn cmp(a: &Cursor, b: &Cursor) -> Ordering; + fn eq<'a>(a: &Cursor<'a, R, Self>, b: &Cursor<'a, R, Self>) -> bool; - fn new_cursor(input_index: usize, rows: R) -> Cursor { - Cursor:: { - input_index, - row_index: 0, - rows, - _o: PhantomData, - } - } + fn cmp<'a>(a: &Cursor<'a, R, Self>, b: &Cursor<'a, R, Self>) -> Ordering; } -impl Ord for Cursor +impl Ord for Cursor<'_, R, O> where R: Rows, O: CursorOrder, @@ -102,7 +107,7 @@ where } } -impl PartialEq for Cursor +impl PartialEq for Cursor<'_, R, O> where R: Rows, O: CursorOrder, @@ -112,14 +117,14 @@ where } } -impl Eq for Cursor +impl Eq for Cursor<'_, R, O> where R: Rows, O: CursorOrder, { } -impl PartialOrd for Cursor +impl PartialOrd for Cursor<'_, R, O> where R: Rows, O: CursorOrder, @@ -128,33 +133,3 @@ where Some(self.cmp(other)) } } - -pub struct CursorMut<'a, R, O> -where - R: Rows, - O: CursorOrder, -{ - pub row_index: usize, - - cursor: &'a Cursor, -} - -impl<'a, R, O> CursorMut<'a, R, O> -where - R: Rows, - O: CursorOrder, -{ - pub fn advance(&mut self) -> usize { - let res = self.row_index; - self.row_index += 1; - res - } - - pub fn is_finished(&self) -> bool { - self.row_index == self.cursor.rows.len() - } - - pub fn current<'b>(&'b self) -> R::Item<'a> { - self.cursor.rows.row(self.row_index) - } -} diff --git a/src/query/pipeline/transforms/src/processors/transforms/sorts/core/loser_tree.rs b/src/query/pipeline/transforms/src/processors/transforms/sorts/core/loser_tree.rs index 014b50073c1e3..bb2f690c0b131 100644 --- a/src/query/pipeline/transforms/src/processors/transforms/sorts/core/loser_tree.rs +++ b/src/query/pipeline/transforms/src/processors/transforms/sorts/core/loser_tree.rs @@ -12,110 +12,112 @@ // See the License for the specific language governing permissions and // limitations under the License. -pub struct LoserTree { - ready: bool, - tree: Vec>, - data: Vec, +type Node = Option<(usize, T)>; + +pub struct LoserTree { + tree: Vec>, } -impl LoserTree { - pub fn from(data: Vec) -> Self { - let length = data.len(); - LoserTree { - ready: false, - tree: vec![None; length], - data, +impl LoserTree { + pub fn with_capacity(capacity: usize) -> Self { + Self { + tree: Vec::with_capacity(capacity), } } - pub fn winner(&self) -> usize { - debug_assert!(self.ready); - self.tree[0].unwrap() + #[inline(always)] + pub fn peek(&self) -> Option<&T> { + self.tree.first()?.as_ref().map(|(_, value)| value) } - pub fn peek(&self) -> &T { - &self.data[self.winner()] + pub fn peek_mut(&mut self) -> Option<&mut T> { + self.tree.first_mut()?.as_mut().map(|(_, value)| value) } - pub fn peek_top2(&self) -> &T { - let top = self.winner(); - let mut top2 = top; - let mut father_loc = (top2 + self.data.len()) / 2; - while father_loc > 0 { - if let Some(father) = self.tree[father_loc] { - if top2 == top || self.data[top2] < self.data[father] { - top2 = father; - } - father_loc /= 2; - } + pub fn rebuild_from(&mut self, values: Vec) { + self.tree.clear(); + self.tree.resize(values.len(), None); + for (slot, value) in values.into_iter().enumerate() { + self.adjust::(slot, value); } - &self.data[top2] } - pub fn peek_mut(&mut self) -> &mut T { - let win = self.winner(); - &mut self.data[win] + pub fn values(&self) -> impl Iterator + '_ { + self.tree.iter().flatten().map(|(_, value)| *value) } - pub fn rebuild(&mut self) { - if self.ready { - return; - } - let length = self.data.len(); - self.tree.fill(None); - for i in 0..length { - self.adjust(i) - } - self.ready = true + pub fn nodes(&self) -> &[Node] { + &self.tree } - pub fn update(&mut self, i: usize, v: T) { - if self.ready && self.winner() == i { - if v >= *self.peek() { - self.data[i] = v; - } else { - self.data[i] = v; - self.adjust(i) - } - } else { - self.data[i] = v; - self.ready = false; - } + #[inline(always)] + pub fn adjust_top(&mut self) { + let (slot, value) = self.tree[0].expect("loser tree must not be empty"); + self.adjust::(slot, value); } - pub fn tree(&self) -> &Vec> { - &self.tree + #[inline(always)] + pub fn take_top(&mut self) -> (usize, T) { + self.tree[0].take().expect("loser tree must not be empty") } - pub fn data(&self) -> &Vec { - &self.data + #[inline(always)] + pub fn replace_top(&mut self, slot: usize, replacement: Option) { + debug_assert!(self.tree[0].is_none()); + match replacement { + Some(replacement) => self.adjust::(slot, replacement), + None => self.remove(slot), + } } - pub fn adjust_top(&mut self) { - let win = self.winner(); - self.adjust(win); + /// Removes the winner while optionally replaying its replacement on the same path. + #[inline(always)] + pub fn promote_with(&mut self, replacement: Option) -> T { + let (slot, value) = self.tree[0].expect("loser tree must not be empty"); + match replacement { + Some(replacement) => self.adjust::(slot, replacement), + None => self.remove(slot), + } + value } - fn adjust(&mut self, index: usize) { - let mut winner: usize = index; - let mut father_loc = (winner + self.data.len()) / 2; + #[inline(always)] + fn adjust(&mut self, slot: usize, value: T) { + let mut winner = (slot, value); + let mut father_loc = (slot + self.tree.len()) / 2; while father_loc > 0 { match self.tree[father_loc] { - None => { + None if BUILD => { self.tree[father_loc] = Some(winner); break; } Some(father) => { - if self.data[winner] < self.data[father] { + if winner.1 < father.1 { self.tree[father_loc] = Some(winner); winner = father; } - father_loc /= 2; } + None => {} } + father_loc /= 2; } self.tree[0] = Some(winner); } + + fn remove(&mut self, slot: usize) { + let mut winner = (slot, None); + let mut father_loc = (slot + self.tree.len()) / 2; + while father_loc > 0 { + if let Some(father) = self.tree[father_loc] + && winner.1 < Some(father.1) + { + self.tree[father_loc] = winner.1.map(|value| (winner.0, value)); + winner = (father.0, Some(father.1)); + } + father_loc /= 2; + } + self.tree[0] = winner.1.map(|value| (winner.0, value)); + } } #[cfg(test)] @@ -124,55 +126,69 @@ mod test { #[test] fn basic() { - let data = vec![ - Some(4), - Some(6), - Some(5), - Some(9), - Some(8), - Some(2), - Some(3), - Some(7), - ]; - let mut tree = LoserTree::from(data); - tree.rebuild(); - - for i in 2..=9 { - assert_eq!(*tree.peek(), Some(11 - i)); - assert_eq!(*tree.peek(), Some(11 - i)); - if i == 9 { - assert_eq!(*tree.peek_top2(), None); - assert_eq!(*tree.peek_top2(), None); - } else { - assert_eq!(*tree.peek_top2(), Some(10 - i)); - assert_eq!(*tree.peek_top2(), Some(10 - i)); - } - let win = tree.winner(); - tree.update(win, None); + let data = vec![4, 6, 5, 9, 8, 2, 3, 7]; + let mut tree = LoserTree::with_capacity(data.len()); + tree.rebuild_from(data); + + for expected in (2..=9).rev() { + assert_eq!(tree.peek(), Some(&expected)); + assert_eq!(tree.promote_with(None), expected); } - assert_eq!(*tree.peek(), None); - assert_eq!(*tree.peek_top2(), None); + assert_eq!(tree.peek(), None); } #[test] fn in_place_update() { - let data = vec![9, 6, 5, 7, 8]; - let mut tree = LoserTree::from(data); - tree.rebuild(); - - *tree.peek_mut() = 4; - tree.adjust_top(); - assert_eq!(tree.winner(), 4); - assert_eq!(*tree.peek(), 8); - - *tree.peek_mut() = 3; - tree.adjust_top(); - assert_eq!(tree.winner(), 3); - assert_eq!(*tree.peek(), 7); - - *tree.peek_mut() = 2; - tree.adjust_top(); - assert_eq!(tree.winner(), 1); - assert_eq!(*tree.peek(), 6); + let mut tree = LoserTree::with_capacity(5); + tree.rebuild_from(vec![9, 6, 5, 7, 8]); + + assert_eq!(tree.promote_with(Some(4)), 9); + assert_eq!(tree.peek(), Some(&8)); + + assert_eq!(tree.promote_with(Some(3)), 8); + assert_eq!(tree.peek(), Some(&7)); + + assert_eq!(tree.promote_with(Some(2)), 7); + assert_eq!(tree.peek(), Some(&6)); + } + + #[test] + fn replace_winner_on_same_path() { + let mut tree = LoserTree::with_capacity(2); + tree.rebuild_from(vec![2, 1]); + + assert_eq!(tree.promote_with(Some(3)), 2); + assert_eq!(tree.promote_with(None), 3); + assert_eq!(tree.promote_with(None), 1); + assert_eq!(tree.peek(), None); + } + + #[test] + fn replace_taken_top() { + let mut tree = LoserTree::with_capacity(3); + tree.rebuild_from(vec![3, 2, 1]); + + let (slot, value) = tree.take_top(); + assert_eq!(value, 3); + assert_eq!(tree.peek(), None); + + tree.replace_top(slot, Some(4)); + assert_eq!(tree.peek(), Some(&4)); + assert_eq!(tree.promote_with(None), 4); + assert_eq!(tree.promote_with(None), 2); + assert_eq!(tree.promote_with(None), 1); + } + + #[test] + fn remove_taken_top() { + let mut tree = LoserTree::with_capacity(3); + tree.rebuild_from(vec![3, 2, 1]); + + let (slot, value) = tree.take_top(); + assert_eq!(value, 3); + tree.replace_top(slot, None); + assert_eq!(tree.promote_with(None), 2); + assert_eq!(tree.promote_with(None), 1); + assert_eq!(tree.peek(), None); } } diff --git a/src/query/pipeline/transforms/src/processors/transforms/sorts/core/merger.rs b/src/query/pipeline/transforms/src/processors/transforms/sorts/core/merger.rs index 63fddfc392473..4100498aa278b 100644 --- a/src/query/pipeline/transforms/src/processors/transforms/sorts/core/merger.rs +++ b/src/query/pipeline/transforms/src/processors/transforms/sorts/core/merger.rs @@ -12,14 +12,17 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::cmp::Reverse; +use std::cmp::Ordering; use std::collections::VecDeque; +use std::future::Future; use databend_common_exception::Result; use databend_common_expression::ChunkIndex; use databend_common_expression::Column; use databend_common_expression::DataBlock; use databend_common_expression::DataBlockVec; +use databend_common_expression::types::DataType; +use futures::future::Either; use super::Rows; use super::algorithm::*; @@ -116,6 +119,106 @@ impl BufferState { } } +/// Owns both the sort cursors and the rows their items borrow from. +/// +/// `cursors` must be declared before `rows` so it is dropped first. +struct CursorStorage { + cursors: A, + rows: Box<[Option]>, +} + +impl CursorStorage { + fn with_capacity(capacity: usize) -> Self { + Self { + cursors: A::with_capacity(capacity), + rows: (0..capacity) + .map(|_| None) + .collect::>() + .into_boxed_slice(), + } + } + + #[inline(always)] + fn is_empty(&self) -> bool { + self.cursors.is_empty() + } + + #[inline(always)] + fn len(&self) -> usize { + self.cursors.len() + } + + #[inline(always)] + fn peek(&self) -> Option<&Cursor> { + self.cursors.peek() + } + + #[inline(always)] + fn peek_top2(&self) -> &Cursor { + self.cursors.peek_top2() + } + + fn rows(&self, stream_index: usize) -> &A::Rows { + self.rows[stream_index] + .as_ref() + .expect("cursor must have originating rows") + } + + fn item_cmp( + &self, + stream_index: usize, + row_index: usize, + item: ::Item<'static>, + ) -> Ordering { + // Safety: the item is used only while its originating Rows is borrowed. + unsafe { self.rows(stream_index).row_stable(row_index).cmp(&item) } + } + + fn push(&mut self, stream_index: usize, rows: A::Rows) { + debug_assert!(self.rows[stream_index].is_none()); + self.rows[stream_index] = Some(rows); + + let rows = self.rows[stream_index].as_ref().unwrap(); + let num_rows = rows.len(); + debug_assert!(num_rows > 0); + // Safety: Rows guarantees its items survive moving the Rows wrapper. + // CursorStorage keeps the originating Rows alive until this cursor is + // removed, and its field order drops all cursors before any rows. + let (current, last) = unsafe { (rows.row_stable(0), rows.row_stable(num_rows - 1)) }; + let cursor = Cursor::new(stream_index, num_rows, current, last); + self.cursors.push(stream_index, cursor); + } + + fn rebuild(&mut self) { + self.cursors.rebuild(); + } + + /// Advances the top cursor and returns its stream index when it is exhausted. + fn advance_top(&mut self, count: usize) -> Option { + let cursor = self.cursors.peek().unwrap(); + let stream_index = cursor.input_index; + let row_index = cursor.row_index + count; + let num_rows = self.rows(stream_index).len(); + debug_assert!(row_index <= num_rows); + // Safety: the originating Rows remains in this storage until after the + // cursor has been removed below. + let current = (row_index < num_rows) + .then(|| unsafe { self.rows(stream_index).row_stable(row_index) }); + + let mut peek_mut = self.cursors.peek_mut(); + let cursor = &mut peek_mut; + cursor.advance(count, current); + + if !cursor.is_finished() { + return None; + } + + A::pop_mut(peek_mut); + self.rows[stream_index] = None; + Some(stream_index) + } +} + /// A merge sort operator to merge multiple sorted streams and output one sorted stream. pub struct Merger where A: SortAlgorithm @@ -125,7 +228,7 @@ where A: SortAlgorithm unsorted_streams: Vec, pending_streams: VecDeque, - sorted_cursors: A, + cursor_storage: CursorStorage, buffers: BufferState, } @@ -136,13 +239,13 @@ where A: SortAlgorithm // We only create a merger when there are at least two streams. debug_assert!(streams.len() > 1, "streams.len() = {}", streams.len()); - let sorted_cursors = A::with_capacity(streams.len()); + let cursor_storage = CursorStorage::with_capacity(streams.len()); let pending_streams = (0..streams.len()).collect(); let buffers = BufferState::new(streams.len()); Self { unsorted_streams: streams, - sorted_cursors, + cursor_storage, batch_rows, limit, pending_streams, @@ -152,7 +255,7 @@ where A: SortAlgorithm #[inline(always)] pub fn is_finished(&self) -> bool { - (self.sorted_cursors.is_empty() && !self.has_pending_stream() && !self.buffers.has_output()) + (self.cursor_storage.is_empty() && !self.has_pending_stream() && !self.buffers.has_output()) || self.limit == Some(0) } @@ -161,14 +264,14 @@ where A: SortAlgorithm !self.pending_streams.is_empty() } - /// To evaluate the current cursor, and update the top of the sorted_cursors if necessary. - /// This method can only be called when iterating the sorted_cursors. + /// To evaluate the current cursor, and update the top of the cursor storage if necessary. + /// This method can only be called when iterating the cursor storage. /// /// Return `true` if the batch is full (need to output). #[inline(always)] fn evaluate_cursor(&mut self) -> bool { - let cursor = if let Some(Reverse(cursor)) = self.sorted_cursors.peek() { - cursor + let cursor = if let Some(cursor) = self.cursor_storage.peek() { + *cursor } else { return false; }; @@ -177,19 +280,11 @@ where A: SortAlgorithm let buffer_index = self.buffers.stream_to_buffer[stream_index] .expect("cursor must point to active stream buffer"); let start = cursor.row_index; - let count = self.evaluate_cursor_count(cursor); + let count = self.evaluate_cursor_count(&cursor); self.buffers.record_output_range(buffer_index, start, count); - // `self.sorted_cursors.peek_mut` will return a `PeekMut` object which allows us to modify the top element of the sorted_cursors. - // The sorted_cursors will adjust itself automatically when the `PeekMut` object is dropped (RAII). - let mut peek_mut = self.sorted_cursors.peek_mut(); - let cursor = &mut peek_mut.0; - cursor.row_index += count; - - if cursor.is_finished() { - // Pop the current `cursor`. - A::pop_mut(peek_mut); + if let Some(stream_index) = self.cursor_storage.advance_top(count) { self.buffers.detach(buffer_index, stream_index); self.pending_streams.push_back(stream_index); } @@ -208,21 +303,25 @@ where A: SortAlgorithm .num_rows() .min(start + max_rows - self.buffers.output_len()); - if self.sorted_cursors.len() == 1 || cursor.current() == cursor.last() { + if self.cursor_storage.len() == 1 || cursor.current() == cursor.last() { return row_index_limit - start; } if !A::SHOULD_PEEK_TOP2 { - let mut p = cursor.cursor_mut(); - p.advance(); - let item = &cursor.current(); - while p.row_index < row_index_limit && p.current() == *item { - p.advance(); + let mut row_index = start + 1; + let item = cursor.current(); + while row_index < row_index_limit + && self + .cursor_storage + .item_cmp(cursor.input_index, row_index, item) + == Ordering::Equal + { + row_index += 1; } - return p.row_index - start; + return row_index - start; } - let next_cursor = &self.sorted_cursors.peek_top2().0; + let next_cursor = self.cursor_storage.peek_top2(); if cursor.last() <= next_cursor.current() { // Short Path: // If the last row of current block is smaller than the next cursor, @@ -230,14 +329,18 @@ where A: SortAlgorithm return row_index_limit - start; } - let mut p = cursor.cursor_mut(); - p.advance(); - let item = &next_cursor.current(); - while p.row_index < row_index_limit && p.current() <= *item { + let mut row_index = start + 1; + let item = next_cursor.current(); + while row_index < row_index_limit + && self + .cursor_storage + .item_cmp(cursor.input_index, row_index, item) + != Ordering::Greater + { // If the cursor is equals or smaller than the next cursor, continue advance. - p.advance(); + row_index += 1; } - p.row_index - start + row_index - start } fn build_output(&mut self) -> Result { @@ -271,12 +374,15 @@ where if let Some((block, col)) = input { let rows = A::Rows::from_column(&col)?; self.buffers.attach_stream_block(i, block)?; - let cursor = Cursor::new(i, rows); - self.sorted_cursors.push(i, Reverse(cursor)); + self.cursor_storage.push(i, rows); } } - self.sorted_cursors.rebuild(); self.pending_streams.extend(continue_pendings); + // `rebuild` is the mutation/read barrier: no cursor is observed while + // an input can still contribute an unresolved next block. + if self.pending_streams.is_empty() { + self.cursor_storage.rebuild(); + } Ok(()) } @@ -296,7 +402,7 @@ where } // No pending streams now. - if self.sorted_cursors.is_empty() { + if self.cursor_storage.is_empty() { return if self.buffers.has_output() { Ok(Some(self.build_output()?)) } else { @@ -322,7 +428,7 @@ where A: SortAlgorithm, S: AsyncSortedStream + Send, { - // This method can only be called when there is no data of the stream in the sorted_cursors. + // This method can only be called when there is no data of the stream in the cursor storage. pub async fn async_poll_pending_stream(&mut self) -> Result<()> { let mut continue_pendings = Vec::new(); while let Some(i) = self.pending_streams.pop_front() { @@ -335,12 +441,14 @@ where if let Some((block, col)) = input { let rows = A::Rows::from_column(&col)?; self.buffers.attach_stream_block(i, block)?; - let cursor = Cursor::new(i, rows); - self.sorted_cursors.push(i, Reverse(cursor)); + self.cursor_storage.push(i, rows); } } - self.sorted_cursors.rebuild(); self.pending_streams.extend(continue_pendings); + // Keep the same read barrier as the synchronous path. + if self.pending_streams.is_empty() { + self.cursor_storage.rebuild(); + } Ok(()) } @@ -358,7 +466,7 @@ where } // No pending streams now. - if self.sorted_cursors.is_empty() { + if self.cursor_storage.is_empty() { return if self.buffers.has_output() { Ok(Some(self.build_output()?)) } else { @@ -379,6 +487,139 @@ where } } -pub type HeapMerger = Merger, S>; +pub type LoserTreeMerger = Merger, S>; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum MergeAlgorithm { + Heap, + LoserTree, + LoserTreeTop2, +} + +/// Runtime selection around fully monomorphized mergers. +/// +/// Dispatch happens once per public merger operation. The cursor evaluation and +/// loser-tree replay loops remain inside a concrete [`Merger`], so they do +/// not pay an enum branch per row or per replay. +pub enum SelectedMerger { + Heap(Merger, S>), + LoserTree(Merger, S>), + LoserTreeTop2(Merger, S>), +} + +impl SelectedMerger { + pub fn new( + algorithm: MergeAlgorithm, + streams: Vec, + batch_rows: usize, + limit: Option, + ) -> Self { + match algorithm { + MergeAlgorithm::Heap => Self::Heap(Merger::new(streams, batch_rows, limit)), + MergeAlgorithm::LoserTree => Self::LoserTree(Merger::new(streams, batch_rows, limit)), + MergeAlgorithm::LoserTreeTop2 => { + Self::LoserTreeTop2(Merger::new(streams, batch_rows, limit)) + } + } + } + + pub fn new_auto( + streams: Vec, + batch_rows: usize, + limit: Option, + enable_loser_tree: bool, + ) -> Self { + let algorithm = if !enable_loser_tree { + MergeAlgorithm::Heap + } else if streams.len() >= 16 { + MergeAlgorithm::LoserTree + } else if std::matches!( + R::data_type(), + DataType::Boolean + | DataType::Number(_) + | DataType::Decimal(_) + | DataType::Timestamp + | DataType::TimestampTz + | DataType::Date + | DataType::Interval + | DataType::Opaque(_) + ) { + MergeAlgorithm::Heap + } else { + MergeAlgorithm::LoserTreeTop2 + }; + Self::new(algorithm, streams, batch_rows, limit) + } + + pub fn is_finished(&self) -> bool { + match self { + Self::Heap(merger) => merger.is_finished(), + Self::LoserTreeTop2(merger) => merger.is_finished(), + Self::LoserTree(merger) => merger.is_finished(), + } + } + + pub fn has_pending_stream(&self) -> bool { + match self { + Self::Heap(merger) => merger.has_pending_stream(), + Self::LoserTreeTop2(merger) => merger.has_pending_stream(), + Self::LoserTree(merger) => merger.has_pending_stream(), + } + } + + pub fn streams(self) -> Vec { + match self { + Self::Heap(merger) => merger.streams(), + Self::LoserTreeTop2(merger) => merger.streams(), + Self::LoserTree(merger) => merger.streams(), + } + } +} + +impl SelectedMerger +where + R: Rows, + S: SortedStream + Send, +{ + pub fn poll_pending_stream(&mut self) -> Result<()> { + match self { + Self::Heap(merger) => merger.poll_pending_stream(), + Self::LoserTreeTop2(merger) => merger.poll_pending_stream(), + Self::LoserTree(merger) => merger.poll_pending_stream(), + } + } + + pub fn next_block(&mut self) -> Result> { + match self { + Self::Heap(merger) => merger.next_block(), + Self::LoserTreeTop2(merger) => merger.next_block(), + Self::LoserTree(merger) => merger.next_block(), + } + } +} -pub type LoserTreeMerger = Merger, S>; +impl SelectedMerger +where + R: Rows, + S: AsyncSortedStream + Send, +{ + pub fn async_poll_pending_stream(&mut self) -> impl Future> + Send + '_ { + match self { + Self::Heap(merger) => Either::Left(Either::Left(merger.async_poll_pending_stream())), + Self::LoserTreeTop2(merger) => { + Either::Left(Either::Right(merger.async_poll_pending_stream())) + } + Self::LoserTree(merger) => Either::Right(merger.async_poll_pending_stream()), + } + } + + pub fn async_next_block( + &mut self, + ) -> impl Future>> + Send + '_ { + match self { + Self::Heap(merger) => Either::Left(Either::Left(merger.async_next_block())), + Self::LoserTreeTop2(merger) => Either::Left(Either::Right(merger.async_next_block())), + Self::LoserTree(merger) => Either::Right(merger.async_next_block()), + } + } +} diff --git a/src/query/pipeline/transforms/src/processors/transforms/sorts/core/row_convert/fixed.rs b/src/query/pipeline/transforms/src/processors/transforms/sorts/core/row_convert/fixed.rs index 7f45ceb8b6bad..57a69aa3aa11e 100644 --- a/src/query/pipeline/transforms/src/processors/transforms/sorts/core/row_convert/fixed.rs +++ b/src/query/pipeline/transforms/src/processors/transforms/sorts/core/row_convert/fixed.rs @@ -80,7 +80,9 @@ pub(super) fn choose_encode_method(fields: &[RowSortField]) -> Option { #[derive(Debug, Clone)] pub struct FixedRows(Buffer<[u64; N]>); -impl Rows for FixedRows { +// SAFETY: Buffer stores its elements in an Arc-backed allocation, so moving the +// Buffer wrapper does not relocate referenced elements. +unsafe impl Rows for FixedRows { const IS_ASC_COLUMN: bool = true; type Item<'a> diff --git a/src/query/pipeline/transforms/src/processors/transforms/sorts/core/row_convert/simple.rs b/src/query/pipeline/transforms/src/processors/transforms/sorts/core/row_convert/simple.rs index 4f2156f353165..b449ff7307f30 100644 --- a/src/query/pipeline/transforms/src/processors/transforms/sorts/core/row_convert/simple.rs +++ b/src/query/pipeline/transforms/src/processors/transforms/sorts/core/row_convert/simple.rs @@ -36,10 +36,12 @@ pub struct SimpleRowsAsc { inner: T::Column, } -impl Rows for SimpleRowsAsc +// SAFETY: scalar references either contain copied values or point into the +// column's heap-backed buffers, which are not relocated with the column wrapper. +unsafe impl Rows for SimpleRowsAsc where T: ArgType, - for<'a> T::ScalarRef<'a>: Ord + Send, + for<'a> T::ScalarRef<'a>: Ord + Copy + Send, { const IS_ASC_COLUMN: bool = true; type Item<'a> @@ -90,10 +92,12 @@ pub struct SimpleRowsDesc { inner: T::Column, } -impl Rows for SimpleRowsDesc +// SAFETY: Reverse only wraps the same move-stable scalar references as the +// ascending implementation. +unsafe impl Rows for SimpleRowsDesc where T: ArgType, - for<'a> T::ScalarRef<'a>: Ord + Send, + for<'a> T::ScalarRef<'a>: Ord + Copy + Send, { const IS_ASC_COLUMN: bool = false; type Item<'a> @@ -150,7 +154,7 @@ pub struct SimpleRowConverter { impl RowConverter> for SimpleRowConverter where T: ArgType, - for<'a> T::ScalarRef<'a>: Ord + Send, + for<'a> T::ScalarRef<'a>: Ord + Copy + Send, { fn new(desc: SortKeyDescription) -> Result { let sort_offset = desc.into_single_sort_offset(true); @@ -172,7 +176,7 @@ where impl RowConverter> for SimpleRowConverter where T: ArgType, - for<'a> T::ScalarRef<'a>: Ord + Send, + for<'a> T::ScalarRef<'a>: Ord + Copy + Send, { fn new(desc: SortKeyDescription) -> Result { let sort_offset = desc.into_single_sort_offset(false); diff --git a/src/query/pipeline/transforms/src/processors/transforms/sorts/core/row_convert/variable.rs b/src/query/pipeline/transforms/src/processors/transforms/sorts/core/row_convert/variable.rs index a9cebe58184a1..d30412f0ceeef 100644 --- a/src/query/pipeline/transforms/src/processors/transforms/sorts/core/row_convert/variable.rs +++ b/src/query/pipeline/transforms/src/processors/transforms/sorts/core/row_convert/variable.rs @@ -63,7 +63,9 @@ use super::variable_encode::var_encode; pub type VariableRows = BinaryColumn; -impl Rows for VariableRows { +// SAFETY: BinaryColumn stores data and offsets in Arc-backed buffers, so moving +// the column wrapper does not relocate referenced byte slices. +unsafe impl Rows for VariableRows { const IS_ASC_COLUMN: bool = true; type Item<'a> = &'a [u8]; type Type = BinaryType; diff --git a/src/query/pipeline/transforms/src/processors/transforms/sorts/core/rows.rs b/src/query/pipeline/transforms/src/processors/transforms/sorts/core/rows.rs index f1bb2e0eacac2..2848b221d22f9 100644 --- a/src/query/pipeline/transforms/src/processors/transforms/sorts/core/rows.rs +++ b/src/query/pipeline/transforms/src/processors/transforms/sorts/core/rows.rs @@ -24,17 +24,31 @@ use databend_common_expression::types::DataType; use super::RowConverter; /// Rows can be compared. -pub trait Rows -where Self: Sized + Clone + Debug + Send +/// +/// # Safety +/// Items returned by [`Rows::row`] must remain valid when `Self` is moved. They +/// may only be invalidated when the originating `Rows` value is dropped. +pub unsafe trait Rows +where Self: Sized + Clone + Debug + Send + 'static { const IS_ASC_COLUMN: bool; - type Item<'a>: Ord + Debug + type Item<'a>: Ord + Debug + Copy + Send where Self: 'a; type Type: ArgType; type Converter: RowConverter; fn len(&self) -> usize; fn row(&self, index: usize) -> Self::Item<'_>; + + /// Returns an item whose lifetime is managed by the caller. + /// + /// # Safety + /// The caller must keep this `Rows` value alive for the returned lifetime. + unsafe fn row_stable<'a>(&self, index: usize) -> Self::Item<'a> { + let item = self.row(index); + unsafe { std::mem::transmute_copy(&item) } + } + fn to_column(&self) -> Column; fn from_column(col: &Column) -> Result; diff --git a/src/query/pipeline/transforms/src/processors/transforms/sorts/mod.rs b/src/query/pipeline/transforms/src/processors/transforms/sorts/mod.rs index fd9b983d39ebc..60792706b0eb8 100644 --- a/src/query/pipeline/transforms/src/processors/transforms/sorts/mod.rs +++ b/src/query/pipeline/transforms/src/processors/transforms/sorts/mod.rs @@ -57,6 +57,7 @@ pub struct Base { pub spiller: S, pub sort_row_offset: usize, pub limit: Option, + pub enable_loser_tree: bool, } #[derive(Debug)] diff --git a/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_collect.rs b/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_collect.rs index 47d14605583f2..3518644a91861 100644 --- a/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_collect.rs +++ b/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_collect.rs @@ -31,18 +31,17 @@ use super::SortSpillParams; use super::TransformSortMergeLimit; use super::core::RowConverter; use super::core::Rows; -use super::core::algorithm::SortAlgorithm; use crate::traits::SortSpiller; #[allow(clippy::large_enum_variant)] -enum Inner { +enum Inner { Collect(Vec), - Limit(TransformSortMergeLimit), - Spill(Vec, SortSpill), + Limit(TransformSortMergeLimit), + Spill(Vec, SortSpill), None, } -pub struct TransformSortCollect { +pub struct TransformSortCollect { name: &'static str, input: Arc, output: Arc, @@ -50,18 +49,18 @@ pub struct TransformSortCollect { max_block_size: usize, default_num_merge: usize, - order_col_converter: Option<::Converter>, + order_col_converter: Option, base: Base, - inner: Inner, + inner: Inner, enable_restore_prefetch: bool, enable_sort_spill_stream_regroup: bool, } -impl TransformSortCollect +impl TransformSortCollect where - A: SortAlgorithm, + R: Rows, S: SortSpiller, { pub fn new( @@ -71,7 +70,7 @@ where max_block_size: usize, default_num_merge: usize, sort_limit: bool, - order_col_converter: Option<::Converter>, + order_col_converter: Option, enable_restore_prefetch: bool, enable_sort_spill_stream_regroup: bool, ) -> Result { @@ -183,7 +182,7 @@ where } None => match &mut self.inner { Inner::Limit(limit_sort) => { - let rows = A::Rows::from_column( + let rows = R::from_column( &block.get_by_offset(self.base.sort_row_offset).to_column(), )?; limit_sort.add_block(block, rows) @@ -234,11 +233,10 @@ where } #[async_trait::async_trait] -impl Processor for TransformSortCollect +impl Processor for TransformSortCollect where - A: SortAlgorithm + 'static, - A::Rows: 'static, - ::Converter: Send + 'static, + R: Rows + 'static, + R::Converter: Send + 'static, S: SortSpiller, { fn name(&self) -> String { diff --git a/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_k_way_merge.rs b/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_k_way_merge.rs index 4dbd37a20ada4..97998c9d07a55 100644 --- a/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_k_way_merge.rs +++ b/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_k_way_merge.rs @@ -42,7 +42,7 @@ use super::core::RowsTypeVisitor; use super::core::SortKeyDescription; use super::core::SortedStream; use super::core::algorithm::HeapSort; -use super::core::algorithm::LoserTreeSort; +use super::core::algorithm::LoserTreeTop2Sort; use super::core::algorithm::SortAlgorithm; use super::core::select_row_type; @@ -103,7 +103,7 @@ impl RowsTypeVisitor for Builder<'_> { R::Converter: Send + 'static, { if self.enable_loser_tree { - self.build::>() + self.build::>() } else { self.build::>() } diff --git a/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_local_merge.rs b/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_local_merge.rs index 51891cacf841d..8d87fcbec237b 100644 --- a/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_local_merge.rs +++ b/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_local_merge.rs @@ -26,16 +26,15 @@ use databend_common_pipeline::core::OutputPort; use databend_common_pipeline::core::Processor; use super::Base; -use super::MemoryMerger; use super::MergeSort; use super::OutputData; use super::RowsStat; +use super::SelectedMemoryMerger; use super::SortSpill; use super::SortSpillParams; use super::TransformSortMergeLimit; use super::core::RowConverter; use super::core::Rows; -use super::core::algorithm::SortAlgorithm; use super::create_memory_merger; use crate::traits::SortSpiller; @@ -50,21 +49,21 @@ enum State { } #[allow(clippy::large_enum_variant)] -enum Inner { +enum Inner { Collect(Vec), - Limit(TransformSortMergeLimit), - Memory(MemoryMerger), - Spill(Vec, SortSpill), + Limit(TransformSortMergeLimit), + Memory(SelectedMemoryMerger), + Spill(Vec, SortSpill), } -pub struct TransformSort { +pub struct TransformSort { name: &'static str, input: Arc, output: Arc, output_data: VecDeque, state: State, - row_converter: ::Converter, + row_converter: R::Converter, /// If the next transform of current transform is [`super::transform_multi_sort_merge::MultiSortMergeProcessor`], /// we can generate and output the order column to avoid the extra converting in the next transform. remove_order_col: bool, @@ -75,16 +74,16 @@ pub struct TransformSort { input_has_order_col: bool, base: Base, - inner: Inner, + inner: Inner, max_block_size: usize, enable_restore_prefetch: bool, enable_sort_spill_stream_regroup: bool, } -impl TransformSort +impl TransformSort where - A: SortAlgorithm, + R: Rows, S: SortSpiller, { #[allow(clippy::too_many_arguments)] @@ -92,12 +91,13 @@ where input: Arc, output: Arc, sort_row_offset: usize, - row_converter: ::Converter, + row_converter: R::Converter, max_block_size: usize, limit: Option<(usize, bool)>, spiller: S, remove_order_col: bool, input_has_order_col: bool, + enable_loser_tree: bool, enable_restore_prefetch: bool, enable_sort_spill_stream_regroup: bool, ) -> Result { @@ -124,6 +124,7 @@ where spiller, sort_row_offset, limit, + enable_loser_tree, }, inner, max_block_size, @@ -186,7 +187,7 @@ where if self.input_has_order_col { match &mut self.inner { Inner::Limit(limit_sort) => { - let rows = A::Rows::from_column( + let rows = R::from_column( &block.get_by_offset(self.base.sort_row_offset).to_column(), )?; limit_sort.add_block(block, rows) @@ -225,11 +226,12 @@ where self.state = State::Finish; return Ok(()); } - let mut merger = create_memory_merger::( + let mut merger = create_memory_merger::( input_data, self.base.sort_row_offset, self.base.limit, self.max_block_size, + self.base.enable_loser_tree, ); if let Some(block) = merger.next_block()? { @@ -284,11 +286,10 @@ where } #[async_trait::async_trait] -impl Processor for TransformSort +impl Processor for TransformSort where - A: SortAlgorithm + 'static, - A::Rows: 'static, - ::Converter: Send + 'static, + R: Rows + 'static, + R::Converter: Send + 'static, S: SortSpiller, { fn name(&self) -> String { diff --git a/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_merge.rs b/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_merge.rs index 76d0b5d26bcae..ff2e167445d89 100644 --- a/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_merge.rs +++ b/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_merge.rs @@ -31,7 +31,7 @@ use super::core::RowsTypeVisitor; use super::core::SortKeyDescription; use super::core::SortedStream; use super::core::algorithm::HeapSort; -use super::core::algorithm::LoserTreeSort; +use super::core::algorithm::LoserTreeTop2Sort; use super::core::algorithm::SortAlgorithm; use super::core::select_row_type; use super::sort_merge_base::MergeSort; @@ -138,7 +138,7 @@ impl TransformSortMerge { } if self.enable_loser_tree { - self.merge_sort_algo::>(batch_size, size_hint) + self.merge_sort_algo::>(batch_size, size_hint) } else { self.merge_sort_algo::>(batch_size, size_hint) } diff --git a/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_merge_limit.rs b/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_merge_limit.rs index 4ae8f0b4475da..e7fb8b5ff4303 100644 --- a/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_merge_limit.rs +++ b/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_merge_limit.rs @@ -31,7 +31,10 @@ use super::sort_merge_base::MergeSort; /// This is a specific version of [`super::transform_sort_merge::TransformSortMerge`] which sort blocks with limit. pub struct TransformSortMergeLimit { - heap: FixedHeap>>, + // `heap` must be dropped before `rows`, because its cursors borrow items + // from the originating rows. + heap: FixedHeap>>, + rows: HashMap, buffer: HashMap, /// Record current memory usage. @@ -56,30 +59,56 @@ impl MergeSort for TransformSortMergeLimit { let input_index = self.next_index; self.next_index += 1; - let mut cursor = Cursor::new(input_index, init_rows); let block_num_bytes = block.memory_size() as u64; self.num_bytes += block_num_bytes; self.num_rows += block.num_rows(); let cur_index = input_index; self.buffer.insert(cur_index, block); - - while !cursor.is_finished() { - if let Some(Reverse(evict)) = self.heap.push(Reverse(cursor.clone())) { - if evict.row_index == 0 { - // Evict the first row of the block, - // which means the block must not appear in the Top-N result. - if let Some(block) = self.buffer.remove(&evict.input_index) { - self.num_bytes -= block.memory_size() as u64; - self.num_rows -= block.num_rows(); + let old = self.rows.insert(cur_index, init_rows); + debug_assert!(old.is_none()); + + { + let rows = self.rows.get(&cur_index).unwrap(); + let num_rows = rows.len(); + debug_assert!(num_rows > 0); + // Safety: Rows guarantees its items survive moving the Rows + // wrapper. `self.rows` keeps the originating Rows alive while any + // cursor for this input remains in `heap`. + let (current, last) = unsafe { (rows.row_stable(0), rows.row_stable(num_rows - 1)) }; + let mut cursor = Cursor::new(input_index, num_rows, current, last); + + while !cursor.is_finished() { + if let Some(Reverse(evict)) = self.heap.push(Reverse(cursor)) { + if evict.row_index == 0 { + // Evict the first row of the block, + // which means the block must not appear in the Top-N result. + if let Some(block) = self.buffer.remove(&evict.input_index) { + self.num_bytes -= block.memory_size() as u64; + self.num_rows -= block.num_rows(); + } + if evict.input_index != cur_index { + let rows = self.rows.remove(&evict.input_index); + debug_assert!(rows.is_some()); + } } - } - if evict.input_index == cur_index { - // The Top-N heap is full, and later rows in current block cannot be put into the heap. - break; + if evict.input_index == cur_index { + // The Top-N heap is full, and later rows in current block cannot be put into the heap. + break; + } } + let row_index = cursor.row_index + 1; + // Safety: the originating Rows remains in `self.rows` for the + // entire loop. + let current = (row_index < num_rows) + .then(|| unsafe { self.rows.get(&cur_index).unwrap().row_stable(row_index) }); + cursor.advance(1, current); } - cursor.advance(); + } + + if !self.buffer.contains_key(&cur_index) { + let rows = self.rows.remove(&cur_index); + debug_assert!(rows.is_some()); } // String views may keep source buffers alive after filtering or slicing. Compact only @@ -126,66 +155,15 @@ impl MergeSort for TransformSortMergeLimit { } } -#[cfg(test)] -mod tests { - use databend_common_exception::Result; - use databend_common_expression::Column; - use databend_common_expression::DataBlock; - use databend_common_expression::FromData; - use databend_common_expression::types::Int32Type; - use databend_common_expression::types::StringType; - - use super::MergeSort; - use super::TransformSortMergeLimit; - use crate::sorts::core::Rows; - use crate::sorts::core::SimpleRowsAsc; - - #[test] - fn test_top_n_compacts_retained_string_views() -> Result<()> { - const SOURCE_ROWS: i32 = 2_000; - const LIMIT: usize = 10; - - let payload_suffix = "x".repeat(256); - let keys = (0..SOURCE_ROWS).collect::>(); - let payloads = keys - .iter() - .map(|key| format!("{key:08}-{payload_suffix}")) - .collect::>(); - let block = DataBlock::new_from_columns(vec![ - Int32Type::from_data(keys), - StringType::from_data(payloads), - ]) - .slice(0..LIMIT); - let rows = SimpleRowsAsc::::from_column(&block.get_by_offset(0).to_column())?; - - let mut sort = TransformSortMergeLimit::create(4_096, LIMIT); - sort.add_block(block, rows)?; - - let retained = sort.buffer.values().next().unwrap(); - let Column::String(payloads) = retained.get_by_offset(1).to_column() else { - unreachable!("expected string payload column") - }; - assert_eq!(payloads.total_bytes_len(), LIMIT * (8 + 1 + 256)); - assert!( - payloads.total_buffer_len() < 16 * 1024, - "Top-N retained {} bytes of source string buffers", - payloads.total_buffer_len(), - ); - assert_eq!(sort.num_bytes().0, retained.memory_size() as u64); - - Ok(()) - } -} - #[derive(Clone, Copy)] struct LocalCursorOrder; impl CursorOrder for LocalCursorOrder { - fn eq(a: &Cursor, b: &Cursor) -> bool { + fn eq<'a>(a: &Cursor<'a, R, Self>, b: &Cursor<'a, R, Self>) -> bool { (a.input_index == b.input_index && a.row_index == b.row_index) || a.current() == b.current() } - fn cmp(a: &Cursor, b: &Cursor) -> Ordering { + fn cmp<'a>(a: &Cursor<'a, R, Self>, b: &Cursor<'a, R, Self>) -> Ordering { if a.input_index == b.input_index { return a.row_index.cmp(&b.row_index); } @@ -199,6 +177,7 @@ impl TransformSortMergeLimit { pub fn create(block_size: usize, limit: usize) -> Self { TransformSortMergeLimit { heap: FixedHeap::new(limit), + rows: HashMap::with_capacity(limit), buffer: HashMap::with_capacity(limit), block_size, next_index: 0, @@ -237,9 +216,92 @@ impl TransformSortMergeLimit { } self.buffer.clear(); + self.rows.clear(); self.num_bytes = ByteSize(0); self.num_rows = 0; output_blocks } } + +#[cfg(test)] +mod tests { + use databend_common_exception::Result; + use databend_common_expression::Column; + use databend_common_expression::DataBlock; + use databend_common_expression::FromData; + use databend_common_expression::types::AccessType; + use databend_common_expression::types::BinaryType; + use databend_common_expression::types::Int32Type; + use databend_common_expression::types::StringType; + + use super::MergeSort; + use super::TransformSortMergeLimit; + use crate::sorts::core::Rows; + use crate::sorts::core::SimpleRowsAsc; + use crate::sorts::core::VariableRows; + + #[test] + fn test_top_n_core_cursor_rows_lifetime() -> Result<()> { + let mut sort = TransformSortMergeLimit::::create(4_096, 3); + for (keys, values) in [ + ( + vec![b"10".as_slice(), b"20".as_slice(), b"30".as_slice()], + vec![10, 20, 30], + ), + ( + vec![b"01".as_slice(), b"02".as_slice(), b"03".as_slice()], + vec![1, 2, 3], + ), + ] { + let block = DataBlock::new_from_columns(vec![Int32Type::from_data(values)]); + let rows = VariableRows::from_column(&BinaryType::from_data(keys))?; + sort.add_block(block, rows)?; + } + + assert_eq!(sort.rows.len(), 1); + let output = sort.on_finish(true)?; + assert!(sort.rows.is_empty()); + assert_eq!(output.len(), 1); + let values = Int32Type::try_downcast_column(&output[0].get_by_offset(0).to_column())?; + assert_eq!(values.as_slice(), &[1, 2, 3]); + + Ok(()) + } + + #[test] + fn test_top_n_compacts_retained_string_views() -> Result<()> { + const SOURCE_ROWS: i32 = 2_000; + const LIMIT: usize = 10; + + let payload_suffix = "x".repeat(256); + let keys = (0..SOURCE_ROWS).collect::>(); + let payloads = keys + .iter() + .map(|key| format!("{key:08}-{payload_suffix}")) + .collect::>(); + let block = DataBlock::new_from_columns(vec![ + Int32Type::from_data(keys), + StringType::from_data(payloads), + ]) + .slice(0..LIMIT); + let rows = SimpleRowsAsc::::from_column(&block.get_by_offset(0).to_column())?; + + let mut sort = TransformSortMergeLimit::create(4_096, LIMIT); + sort.add_block(block, rows)?; + + let retained = sort.buffer.values().next().unwrap(); + let Column::String(payloads) = retained.get_by_offset(1).to_column() else { + unreachable!("expected string payload column") + }; + assert_eq!(payloads.total_bytes_len(), LIMIT * (8 + 1 + 256)); + assert!( + payloads.total_buffer_len() < 16 * 1024, + "Top-N retained {} bytes of source string buffers", + payloads.total_buffer_len(), + ); + assert_eq!(sort.num_bytes().0, retained.memory_size() as u64); + + Ok(()) + } +} diff --git a/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_multi_merge.rs b/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_multi_merge.rs index 4814c9d025f31..82252fc4167a8 100644 --- a/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_multi_merge.rs +++ b/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_multi_merge.rs @@ -34,7 +34,7 @@ use super::core::RowsTypeVisitor; use super::core::SortKeyDescription; use super::core::SortedStream; use super::core::algorithm::HeapSort; -use super::core::algorithm::LoserTreeSort; +use super::core::algorithm::LoserTreeTop2Sort; use super::core::algorithm::SortAlgorithm; use super::core::select_row_type; @@ -99,7 +99,7 @@ impl RowsTypeVisitor for MultiSortMergeBuilder { R::Converter: Send + 'static, { if self.enable_loser_tree { - self.create_processor::>() + self.create_processor::>() } else { self.create_processor::>() } diff --git a/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_restore.rs b/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_restore.rs index 3ddf94c9dfac0..fb1ee67c423cd 100644 --- a/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_restore.rs +++ b/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_restore.rs @@ -27,14 +27,14 @@ use super::Base; use super::SortBound; use super::SortBoundNext; use super::SortCollectedMeta; -use super::core::algorithm::SortAlgorithm; +use super::core::Rows; use super::sort_spill::OutputData; use super::sort_spill::SortSpill; use crate::HookTransform; use crate::HookTransformer; use crate::traits::SortSpiller; -pub struct TransformSortRestore { +pub struct TransformSortRestore { input: Vec, output: Option, @@ -43,12 +43,12 @@ pub struct TransformSortRestore { remove_order_col: bool, base: Base, - inner: Option>, + inner: Option>, } -impl TransformSortRestore +impl TransformSortRestore where - A: SortAlgorithm + Send + 'static, + R: Rows + Send + 'static, S: SortSpiller, { pub fn new( @@ -68,10 +68,9 @@ where } #[async_trait::async_trait] -impl HookTransform for TransformSortRestore +impl HookTransform for TransformSortRestore where - A: SortAlgorithm + 'static, - A::Rows: 'static, + R: Rows + 'static, S: SortSpiller, { const NAME: &'static str = "TransformSortRestore"; diff --git a/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_spill.rs b/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_spill.rs index b939d5748e68e..5f4ccae733c36 100644 --- a/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_spill.rs +++ b/src/query/pipeline/transforms/src/processors/transforms/sorts/sort_spill.rs @@ -42,31 +42,31 @@ use super::core::AsyncSortedStream; use super::core::Bounds; use super::core::Merger; use super::core::Rows; +use super::core::SelectedMerger; use super::core::SortedStream; -use super::core::algorithm::SortAlgorithm; use super::sort_spill_regroup::IntervalGroupingPayload; use super::sort_spill_regroup::regroup_min_interval_groups; use crate::MemorySettings; use crate::traits::Location; use crate::traits::SortSpiller; -pub struct SortSpill { +pub struct SortSpill { base: Base, - step: Step, + step: Step, } -enum Step { - Collect(StepCollect), - Sort(StepSort), +enum Step { + Collect(StepCollect), + Sort(StepSort), } -struct StepCollect { +struct StepCollect { params: SortSpillParams, sampler: FixedRateSampler, - streams: Vec>, + streams: Vec>, } -struct StepSort { +struct StepSort { params: SortSpillParams, /// Partition boundaries for restoring and sorting blocks. /// Each boundary represents a cutoff point where data less than or equal to it belongs to one partition. @@ -74,15 +74,15 @@ struct StepSort { cur_bound: Option, bound_index: i32, - subsequent: Vec>, - current: Vec>, + subsequent: Vec>, + current: Vec>, - output_merger: Option>>, + output_merger: Option>>, } -impl SortSpill +impl SortSpill where - A: SortAlgorithm, + R: Rows, S: SortSpiller, { pub fn new(base: Base, params: SortSpillParams) -> Self { @@ -201,7 +201,7 @@ where } } - pub fn format_memory_usage(&self) -> FmtMemoryUsage<'_, A, S> { + pub fn format_memory_usage(&self) -> FmtMemoryUsage<'_, R, S> { FmtMemoryUsage(self) } @@ -237,7 +237,7 @@ where } } -impl StepCollect { +impl StepCollect { #[fastrace::trace(name = "StepCollect::sort_input_data")] async fn sort_input_data( &mut self, @@ -260,8 +260,13 @@ impl StepCollect { } vec![block].into() } else { - let mut merger = - create_memory_merger::(input_data, base.sort_row_offset, base.limit, batch_rows); + let mut merger = create_memory_merger::( + input_data, + base.sort_row_offset, + base.limit, + batch_rows, + base.enable_loser_tree, + ); let mut sorted = VecDeque::new(); while let Some(data) = merger.next_block()? { @@ -282,10 +287,10 @@ impl StepCollect { Ok(()) } - fn next_step(&mut self, base: &Base) -> Result> { + fn next_step(&mut self, base: &Base) -> Result> { self.sampler.compact_blocks(true); let sampled_rows = std::mem::take(&mut self.sampler.dense_blocks); - let bounds = base.determine_bounds::(sampled_rows, self.params.batch_rows)?; + let bounds = base.determine_bounds::(sampled_rows, self.params.batch_rows)?; Ok(StepSort { bounds, @@ -305,7 +310,7 @@ pub struct OutputData { pub finish: bool, } -impl StepSort { +impl StepSort { fn next_bound(&mut self) { match self.bounds.next_bound() { Some(bound) => self.cur_bound = Some(bound), @@ -325,7 +330,12 @@ impl StepSort { .drain(self.current.len() - num_merge..) .collect(); - let mut merger = Merger::::new(streams, self.params.batch_rows, None); + let mut merger = SelectedMerger::::new_auto( + streams, + self.params.batch_rows, + None, + base.enable_loser_tree, + ); let mut sorted = VecDeque::new(); while let Some(data) = merger.async_next_block().await? { @@ -381,8 +391,12 @@ impl StepSort { }); } - let merger = - Merger::::new(mem::take(&mut self.current), self.params.batch_rows, None); + let merger = SelectedMerger::::new_auto( + mem::take(&mut self.current), + self.params.batch_rows, + None, + base.enable_loser_tree, + ); self.output_merger.insert(merger) } }; @@ -564,23 +578,23 @@ impl Base { SpillableBlock::new(data, self.sort_row_offset) } - fn determine_bounds( + fn determine_bounds( &self, sampled_rows: Vec, batch_rows: usize, ) -> Result { match sampled_rows.len() { 0 => Ok(Bounds::default()), - 1 => Bounds::from_column::(sampled_rows[0].get_last_column().clone()), + 1 => Bounds::from_column::(sampled_rows[0].get_last_column().clone()), _ => { let ls = sampled_rows .into_iter() .map(|data| { let col = data.get_last_column().clone(); - Bounds::from_column::(col) + Bounds::from_column::(col) }) .collect::>>()?; - Bounds::merge::(ls, batch_rows) + Bounds::merge::(ls, batch_rows) } } } @@ -617,9 +631,9 @@ impl RowsStat for Vec> { } } -pub struct FmtMemoryUsage<'a, A: SortAlgorithm, S: SortSpiller>(&'a SortSpill); +pub struct FmtMemoryUsage<'a, R: Rows, S: SortSpiller>(&'a SortSpill); -impl fmt::Debug for FmtMemoryUsage<'_, A, S> { +impl fmt::Debug for FmtMemoryUsage<'_, R, S> { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { let debug = &mut f.debug_struct("SortSpill"); match &self.0.step { @@ -1058,17 +1072,20 @@ impl DataBlockStream { pub type MemoryMerger = Merger; -pub fn create_memory_merger( +pub type SelectedMemoryMerger = SelectedMerger; + +pub fn create_memory_merger( blocks: Vec, sort_row_offset: usize, limit: Option, batch_rows: usize, -) -> MemoryMerger { + enable_loser_tree: bool, +) -> SelectedMemoryMerger { let streams = blocks .into_iter() .map(|data| DataBlockStream::new(data, sort_row_offset)) .collect(); - Merger::::new(streams, batch_rows, limit) + SelectedMerger::new_auto(streams, batch_rows, limit, enable_loser_tree) } fn get_domain(entry: &BlockEntry) -> Column { @@ -1117,7 +1134,6 @@ mod tests { use super::*; use crate::sorts::core::SimpleRowsAsc; use crate::sorts::core::SimpleRowsDesc; - use crate::sorts::core::algorithm::HeapSort; use crate::sorts::core::convert_rows; fn test_data() -> (DataSchemaRef, DataBlock) { @@ -1703,7 +1719,7 @@ mod tests { let existed_subsequent = create_int_stream_asc(spiller.clone(), vec![vec![10, 11]], None, true).await?; - let mut sort = StepSort::>, _> { + let mut sort = StepSort::, _> { params: SortSpillParams { batch_rows: 2, num_merge: 2, diff --git a/src/query/pipeline/transforms/tests/it/merger.rs b/src/query/pipeline/transforms/tests/it/merger.rs index 6655e0086a0d6..9bb7fdc53176f 100644 --- a/src/query/pipeline/transforms/tests/it/merger.rs +++ b/src/query/pipeline/transforms/tests/it/merger.rs @@ -21,11 +21,14 @@ use databend_common_expression::FromData; use databend_common_expression::block_debug::pretty_format_blocks; use databend_common_expression::types::Int32Type; use databend_common_pipeline_transforms::sorts::core::AsyncSortedStream; +use databend_common_pipeline_transforms::sorts::core::MergeAlgorithm; use databend_common_pipeline_transforms::sorts::core::Merger; +use databend_common_pipeline_transforms::sorts::core::SelectedMerger; use databend_common_pipeline_transforms::sorts::core::SimpleRowsAsc; use databend_common_pipeline_transforms::sorts::core::SortedStream; use databend_common_pipeline_transforms::sorts::core::algorithm::HeapSort; use databend_common_pipeline_transforms::sorts::core::algorithm::LoserTreeSort; +use databend_common_pipeline_transforms::sorts::core::algorithm::LoserTreeTop2Sort; use databend_common_pipeline_transforms::sorts::core::algorithm::SortAlgorithm; use itertools::Itertools; use rand::Rng; @@ -74,7 +77,9 @@ impl AsyncSortedStream for TestStream { type TestMerger = Merger; type TestHeapSort = HeapSort>; +type TestLoserTreeTop2Sort = LoserTreeTop2Sort>; type TestLoserTreeSort = LoserTreeSort>; +type TestSelectedMerger = SelectedMerger, TestStream>; fn prepare_input_and_result( data: Vec>>, @@ -151,6 +156,19 @@ fn create_test_merger( TestMerger::::new(streams, 4, limit) } +fn create_selected_test_merger( + algorithm: MergeAlgorithm, + input: Vec>, + limit: Option, +) -> TestSelectedMerger { + let streams = input + .into_iter() + .map(|v| TestStream::new(v.into_iter().collect::>())) + .collect::>(); + + TestSelectedMerger::new(algorithm, streams, 4, limit) +} + fn check_result(result: Vec, expected: DataBlock) { if expected.is_empty() { if !result.is_empty() && !DataBlock::concat(&result).unwrap().is_empty() { @@ -203,13 +221,55 @@ async fn async_test( Ok(()) } +fn test_selected(mut merger: TestSelectedMerger, expected: DataBlock) -> Result<()> { + let mut result = Vec::new(); + while !merger.is_finished() { + if let Some(block) = merger.next_block()? { + result.push(block); + } + } + check_result(result, expected); + Ok(()) +} + +async fn async_test_selected(mut merger: TestSelectedMerger, expected: DataBlock) -> Result<()> { + let mut result = Vec::new(); + while !merger.is_finished() { + if let Some(block) = merger.async_next_block().await? { + result.push(block); + } + } + check_result(result, expected); + Ok(()) +} + +fn test_move_after_initial_poll( + mut merger: TestMerger, + expected: DataBlock, +) -> Result<()> { + while merger.has_pending_stream() { + merger.poll_pending_stream()?; + } + + // Move the Merger after its cursors have borrowed rows from CursorStorage. + let mut merger = Box::new(merger); + let mut result = Vec::new(); + while !merger.is_finished() { + if let Some(block) = merger.next_block()? { + result.push(block); + } + } + check_result(result, expected); + Ok(()) +} + fn test_basic(limit: Option) -> Result<()> { let (input, expected) = basic_test_data(limit); let merger = create_test_merger::(input, limit); test(merger, expected)?; let (input, expected) = basic_test_data(limit); - let merger = create_test_merger::(input, limit); + let merger = create_test_merger::(input, limit); test(merger, expected) } @@ -219,7 +279,7 @@ async fn async_test_basic(limit: Option) -> Result<()> { async_test(merger, expected).await?; let (input, expected) = basic_test_data(limit); - let merger = create_test_merger::>>(input, limit); + let merger = create_test_merger::>>(input, limit); async_test(merger, expected).await } @@ -235,6 +295,51 @@ fn test_basic_with_limit() -> anyhow::Result<()> { Ok(()) } +#[test] +fn test_selected_merger() -> anyhow::Result<()> { + for algorithm in [ + MergeAlgorithm::Heap, + MergeAlgorithm::LoserTreeTop2, + MergeAlgorithm::LoserTree, + ] { + let (input, expected) = basic_test_data(None); + test_selected( + create_selected_test_merger(algorithm, input, None), + expected, + )?; + } + Ok(()) +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_selected_merger_async() -> anyhow::Result<()> { + for algorithm in [ + MergeAlgorithm::Heap, + MergeAlgorithm::LoserTreeTop2, + MergeAlgorithm::LoserTree, + ] { + let (input, expected) = basic_test_data(None); + async_test_selected( + create_selected_test_merger(algorithm, input, None), + expected, + ) + .await?; + } + Ok(()) +} + +#[test] +fn test_move_merger_with_active_cursors() -> anyhow::Result<()> { + let (input, expected) = basic_test_data(None); + let merger = create_test_merger::(input, None); + test_move_after_initial_poll(merger, expected)?; + + let (input, expected) = basic_test_data(None); + let merger = create_test_merger::(input, None); + test_move_after_initial_poll(merger, expected)?; + Ok(()) +} + #[tokio::test(flavor = "multi_thread")] async fn async_test_basic_with_limit() -> anyhow::Result<()> { async_test_basic(None).await?; @@ -257,6 +362,12 @@ fn test_fuzz() -> anyhow::Result<()> { test(merger, expected)?; } + for _ in 0..10 { + let (input, expected, limit) = random_test_data(&mut rng); + let merger = create_test_merger::(input, limit); + test(merger, expected)?; + } + for _ in 0..10 { let (input, expected, limit) = random_test_data(&mut rng); let merger = create_test_merger::(input, limit); @@ -278,7 +389,7 @@ async fn test_fuzz_async() -> anyhow::Result<()> { for _ in 0..10 { let (input, expected, limit) = random_test_data(&mut rng); - let merger = create_test_merger::(input, limit); + let merger = create_test_merger::(input, limit); async_test(merger, expected).await?; } diff --git a/src/query/service/src/pipelines/processors/transforms/sort/sort_builder.rs b/src/query/service/src/pipelines/processors/transforms/sort/sort_builder.rs index b75920fb76de7..0798b10096e95 100644 --- a/src/query/service/src/pipelines/processors/transforms/sort/sort_builder.rs +++ b/src/query/service/src/pipelines/processors/transforms/sort/sort_builder.rs @@ -39,7 +39,7 @@ use databend_common_pipeline_transforms::sorts::core::Rows; use databend_common_pipeline_transforms::sorts::core::RowsTypeVisitor; use databend_common_pipeline_transforms::sorts::core::SortKeyDescription; use databend_common_pipeline_transforms::sorts::core::algorithm::HeapSort; -use databend_common_pipeline_transforms::sorts::core::algorithm::LoserTreeSort; +use databend_common_pipeline_transforms::sorts::core::algorithm::LoserTreeTop2Sort; use databend_common_pipeline_transforms::sorts::core::algorithm::SortAlgorithm; use databend_common_pipeline_transforms::sorts::core::select_row_type; use databend_common_pipeline_transforms::traits::SortSpiller; @@ -220,6 +220,7 @@ impl TransformSortBuilder { sort_row_offset: self.key_desc.sort_row_offset(), spiller: self.spiller.clone().unwrap(), limit: self.limit, + enable_loser_tree: self.enable_loser_tree, } } @@ -278,20 +279,20 @@ struct Build<'a, S: SortSpiller> { } impl Build<'_, S> { - fn build_sort( + fn build_sort( &mut self, sort_limit: bool, input: Arc, ) -> Result> where - A: SortAlgorithm + 'static, - ::Converter: Send + 'static, + R: Rows + 'static, + R::Converter: Send + 'static, { let key_desc = self.params.key_desc.clone(); let uses_source_sort_col = key_desc.uses_source_sort_col(); let sort_row_offset = key_desc.sort_row_offset(); - let row_converter = ::Converter::new(key_desc)?; - Ok(Box::new(TransformSort::::new( + let row_converter = R::Converter::new(key_desc)?; + Ok(Box::new(TransformSort::::new( input, self.output.clone(), sort_row_offset, @@ -303,24 +304,25 @@ impl Build<'_, S> { // and cannot be dropped by a trailing `pop_columns(1)`. !self.params.keep_order_col && !uses_source_sort_col, self.params.input_has_order_col || uses_source_sort_col, + self.params.enable_loser_tree, self.params.enable_restore_prefetch, self.params.enable_sort_spill_stream_regroup, )?)) } - fn build_sort_collect( + fn build_sort_collect( &mut self, input: Arc, sort_limit: bool, default_num_merge: usize, ) -> Result> where - A: SortAlgorithm + 'static, - ::Converter: Send + 'static, + R: Rows + 'static, + R::Converter: Send + 'static, { assert!(!self.params.input_has_order_col); - let row_converter = ::Converter::new(self.params.key_desc.clone())?; - Ok(Box::new(TransformSortCollect::::new( + let row_converter = R::Converter::new(self.params.key_desc.clone())?; + Ok(Box::new(TransformSortCollect::::new( input, self.output.clone(), self.params.new_base(), @@ -337,9 +339,9 @@ impl Build<'_, S> { )?)) } - fn build_sort_restore(&mut self, input: Arc) -> Result> - where A: SortAlgorithm + 'static { - Ok(Box::new(TransformSortRestore::::new( + fn build_sort_restore(&mut self, input: Arc) -> Result> + where R: Rows + 'static { + Ok(Box::new(TransformSortRestore::::new( input, self.output.clone(), self.params.new_base(), @@ -392,34 +394,19 @@ impl RowsTypeVisitor for Build<'_, S> { { let sort_limit = self.params.should_use_sort_limit(); match self.typ.take().unwrap() { - SortType::Sort(input) => match self.params.enable_loser_tree { - true => self.build_sort::>(sort_limit, input), - false => self.build_sort::>(sort_limit, input), - }, + SortType::Sort(input) => self.build_sort::(sort_limit, input), SortType::Collect { input, default_num_merge, - } => match self.params.enable_loser_tree { - true => self.build_sort_collect::>( - input, - sort_limit, - default_num_merge, - ), - false => { - self.build_sort_collect::>(input, sort_limit, default_num_merge) - } - }, + } => self.build_sort_collect::(input, sort_limit, default_num_merge), SortType::BoundBroadcast { input, state } => { self.build_bound_broadcast::(input, state) } - SortType::Restore(input) => match self.params.enable_loser_tree { - true => self.build_sort_restore::>(input), - false => self.build_sort_restore::>(input), - }, + SortType::Restore(input) => self.build_sort_restore::(input), SortType::BoundedMergeSort(inputs) => match self.params.enable_loser_tree { - true => self.build_bounded_merge_sort::>(inputs), + true => self.build_bounded_merge_sort::>(inputs), false => self.build_bounded_merge_sort::>(inputs), }, } diff --git a/src/query/storages/paimon/Cargo.toml b/src/query/storages/paimon/Cargo.toml index 0471c340e05f0..4556383a7aeaa 100644 --- a/src/query/storages/paimon/Cargo.toml +++ b/src/query/storages/paimon/Cargo.toml @@ -21,13 +21,11 @@ databend-common-expression = { workspace = true } databend-common-meta-app = { workspace = true } databend-common-pipeline = { workspace = true } databend-common-pipeline-transforms = { workspace = true } -databend-common-storage = { workspace = true } databend-common-users = { workspace = true } databend-meta-client = { workspace = true } databend-storages-common-table-meta = { workspace = true } educe = { workspace = true } futures = { workspace = true } -log = { workspace = true } paimon = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } diff --git a/tests/sqllogictests/suites/query/join/merge_join.test b/tests/sqllogictests/suites/query/join/merge_join.test index 32ead0072f741..0e8f746b7cf07 100644 --- a/tests/sqllogictests/suites/query/join/merge_join.test +++ b/tests/sqllogictests/suites/query/join/merge_join.test @@ -1,23 +1,17 @@ statement ok -drop table if exists t1; - -statement ok -drop table if exists t2; - -statement ok -create table t1(a int null); +create or replace table t1(a int null); statement ok insert into t1 values(0), (1), (NULL), (2); statement ok -create table t2(b int null); +create or replace table t2(b int null); statement ok insert into t2 values(1), (NULL), (3); query II -select * from t1 join t2 on t1.a < t2.b order by t1.a; +select * from t1 join t2 on t1.a < t2.b order by t1.a, b; ---- 0 1 0 3