You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by al...@apache.org on 2023/05/09 15:30:08 UTC
[arrow-datafusion] branch main updated: Add more documentation to SortPreservingMergeStream (#6260)
This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 894beb1d3b Add more documentation to SortPreservingMergeStream (#6260)
894beb1d3b is described below
commit 894beb1d3b2b42ecc58c582f98e17800ac2b9527
Author: Armin Primadi <ap...@gmail.com>
AuthorDate: Tue May 9 22:30:01 2023 +0700
Add more documentation to SortPreservingMergeStream (#6260)
* Add more documentation to SortPreservingMergeStream
* Update documentation
* Fix cargo test doc error
* Add merging 5 partitions test
* More documentation
* Update datafusion/core/src/physical_plan/sorts/sort_preserving_merge.rs
Co-authored-by: Armin Primadi <ap...@gmail.com>
---------
Co-authored-by: Andrew Lamb <an...@nerdnetworks.org>
---
datafusion/core/src/physical_plan/sorts/merge.rs | 65 ++++++++++++++++++++++--
1 file changed, 61 insertions(+), 4 deletions(-)
diff --git a/datafusion/core/src/physical_plan/sorts/merge.rs b/datafusion/core/src/physical_plan/sorts/merge.rs
index 7e2d986e9d..e6533deaba 100644
--- a/datafusion/core/src/physical_plan/sorts/merge.rs
+++ b/datafusion/core/src/physical_plan/sorts/merge.rs
@@ -108,6 +108,24 @@ struct SortPreservingMergeStream<C> {
/// the path from bottom to top is visited, keeping the number of
/// comparisons close to the theoretical limit of `log(S)`.
///
+ /// The current implementation uses a vector to store the tree.
+ /// Conceptually, it looks like this (assuming 8 streams):
+ ///
+ /// ```text
+ /// 0 (winner)
+ ///
+ /// 1
+ /// / \
+ /// 2 3
+ /// / \ / \
+ /// 4 5 6 7
+ /// ```
+ ///
+ /// Where element at index 0 in the vector is the current winner. Element
+ /// at index 1 is the root of the loser tree, element at index 2 is the
+ /// left child of the root, and element at index 3 is the right child of
+ /// the root and so on.
+ ///
/// reference: <https://en.wikipedia.org/wiki/K-way_merge_algorithm#Tournament_Tree>
loser_tree: Vec<usize>,
@@ -242,6 +260,45 @@ impl<C: Cursor> SortPreservingMergeStream<C> {
}
}
+ /// Find the leaf node index in the loser tree for the given cursor index
+ ///
+ /// Note that this is not necessarily a leaf node in the tree, but it can
+ /// also be a half-node (a node with only one child). This happens when the
+ /// number of cursors/streams is not a power of two. Thus, the loser tree
+ /// will be unbalanced, but it will still work correctly.
+ ///
+ /// For example, with 5 streams, the loser tree will look like this:
+ ///
+ /// ```text
+ /// 0 (winner)
+ ///
+ /// 1
+ /// / \
+ /// 2 3
+ /// / \ / \
+ /// 4 | | |
+ /// / \ | | |
+ /// -+---+--+---+---+---- Below is not a part of loser tree
+ /// S3 S4 S0 S1 S2
+ /// ```
+ ///
+ /// S0, S1, ... S4 are the streams (read: stream at index 0, stream at
+ /// index 1, etc.)
+ ///
+ /// Zooming in at node 2 in the loser tree as an example, we can see that
+ /// it takes as input the next item at (S0) and the loser of (S3, S4).
+ ///
+ #[inline]
+ fn lt_leaf_node_index(&self, cursor_index: usize) -> usize {
+ (self.cursors.len() + cursor_index) / 2
+ }
+
+ /// Find the parent node index for the given node index
+ #[inline]
+ fn lt_parent_node_index(&self, node_idx: usize) -> usize {
+ node_idx / 2
+ }
+
/// Attempts to initialize the loser tree with one value from each
/// non exhausted input, if possible
fn init_loser_tree(&mut self) {
@@ -249,7 +306,7 @@ impl<C: Cursor> SortPreservingMergeStream<C> {
self.loser_tree = vec![usize::MAX; self.cursors.len()];
for i in 0..self.cursors.len() {
let mut winner = i;
- let mut cmp_node = (self.cursors.len() + i) / 2;
+ let mut cmp_node = self.lt_leaf_node_index(i);
while cmp_node != 0 && self.loser_tree[cmp_node] != usize::MAX {
let challenger = self.loser_tree[cmp_node];
if self.is_gt(winner, challenger) {
@@ -257,7 +314,7 @@ impl<C: Cursor> SortPreservingMergeStream<C> {
winner = challenger;
}
- cmp_node /= 2;
+ cmp_node = self.lt_parent_node_index(cmp_node);
}
self.loser_tree[cmp_node] = winner;
}
@@ -268,14 +325,14 @@ impl<C: Cursor> SortPreservingMergeStream<C> {
fn update_loser_tree(&mut self) {
let mut winner = self.loser_tree[0];
// Replace overall winner by walking tree of losers
- let mut cmp_node = (self.cursors.len() + winner) / 2;
+ let mut cmp_node = self.lt_leaf_node_index(winner);
while cmp_node != 0 {
let challenger = self.loser_tree[cmp_node];
if self.is_gt(winner, challenger) {
self.loser_tree[cmp_node] = winner;
winner = challenger;
}
- cmp_node /= 2;
+ cmp_node = self.lt_parent_node_index(cmp_node);
}
self.loser_tree[0] = winner;
self.loser_tree_adjusted = true;