You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by GitBox <gi...@apache.org> on 2022/11/21 08:21:30 UTC

[GitHub] [arrow-datafusion] richox opened a new pull request, #4301: Use tournament loser tree for k-way sort-merging

richox opened a new pull request, #4301:
URL: https://github.com/apache/arrow-datafusion/pull/4301

   # Which issue does this PR close?
   
   <!--
   We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123.
   -->
   
   Closes #4300.
   
   # Rationale for this change
   
   <!--
    Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
    Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.  
   -->
   
   # What changes are included in this PR?
   
   <!--
   There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
   -->
   
   # Are these changes tested?
   
   <!--
   We typically require tests for all PRs in order to:
   1. Prevent the code from being accidentally broken by subsequent changes
   2. Serve as another way to document the expected behavior of the code
   
   If tests are not included in your PR, please explain why (for example, are they covered by existing tests)?
   -->
   
   running benchmarks with `cargo bench --bench merge`:
   
   ```
   merge i64               time:   [8.0959 ms 8.2001 ms 8.3260 ms]
                           change: [-58.629% -57.906% -57.138%] (p = 0.00 < 0.05)
                           Performance has improved.
   Found 11 outliers among 100 measurements (11.00%)
     4 (4.00%) high mild
     7 (7.00%) high severe
   
   merge f64               time:   [8.2829 ms 8.3378 ms 8.4028 ms]
                           change: [-56.127% -55.475% -54.829%] (p = 0.00 < 0.05)
                           Performance has improved.
   Found 10 outliers among 100 measurements (10.00%)
     5 (5.00%) high mild
     5 (5.00%) high severe
   
   merge utf8 low cardinality
                           time:   [7.5418 ms 7.6425 ms 7.7538 ms]
                           change: [-65.897% -65.335% -64.703%] (p = 0.00 < 0.05)
                           Performance has improved.
   Found 13 outliers among 100 measurements (13.00%)
     7 (7.00%) high mild
     6 (6.00%) high severe
   
   merge utf8 high cardinality
                           time:   [9.9438 ms 10.094 ms 10.257 ms]
                           change: [-52.043% -51.155% -50.249%] (p = 0.00 < 0.05)
                           Performance has improved.
   Found 8 outliers among 100 measurements (8.00%)
     4 (4.00%) high mild
     4 (4.00%) high severe
   
   merge utf8 tuple        time:   [20.197 ms 20.493 ms 20.826 ms]
                           change: [-40.663% -39.505% -38.279%] (p = 0.00 < 0.05)
                           Performance has improved.
   Found 10 outliers among 100 measurements (10.00%)
     6 (6.00%) high mild
     4 (4.00%) high severe
   
   merge utf8 dictionary   time:   [5.1537 ms 5.2135 ms 5.2788 ms]
                           change: [-72.062% -71.616% -71.137%] (p = 0.00 < 0.05)
                           Performance has improved.
   Found 8 outliers among 100 measurements (8.00%)
     7 (7.00%) high mild
     1 (1.00%) high severe
   
   merge utf8 dictionary tuple
                           time:   [7.9585 ms 8.0592 ms 8.1797 ms]
                           change: [-63.981% -63.324% -62.678%] (p = 0.00 < 0.05)
                           Performance has improved.
   Found 11 outliers among 100 measurements (11.00%)
     6 (6.00%) high mild
     5 (5.00%) high severe
   
   merge mixed utf8 dictionary tuple
                           time:   [14.528 ms 14.669 ms 14.829 ms]
                           change: [-43.886% -43.201% -42.489%] (p = 0.00 < 0.05)
                           Performance has improved.
   Found 10 outliers among 100 measurements (10.00%)
     5 (5.00%) high mild
     5 (5.00%) high severe
   
   merge mixed tuple       time:   [17.824 ms 18.068 ms 18.333 ms]
                           change: [-39.957% -38.888% -37.748%] (p = 0.00 < 0.05)
                           Performance has improved.
   Found 5 outliers among 100 measurements (5.00%)
     4 (4.00%) high mild
     1 (1.00%) high severe
   
   ```
   # Are there any user-facing changes?
   
   <!--
   If there are user-facing changes then we may require documentation to be updated before approving the PR.
   -->
   
   <!--
   If there are any breaking changes to public APIs, please add the `api change` label.
   -->


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [arrow-datafusion] alamb merged pull request #4301: Use tournament loser tree for k-way sort-merging, increase merge speed by 50%

Posted by GitBox <gi...@apache.org>.
alamb merged PR #4301:
URL: https://github.com/apache/arrow-datafusion/pull/4301


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [arrow-datafusion] alamb commented on pull request #4301: Use tournament loser tree for k-way sort-merging, increase merge speed by 50%

Posted by GitBox <gi...@apache.org>.
alamb commented on PR #4301:
URL: https://github.com/apache/arrow-datafusion/pull/4301#issuecomment-1329744186

   Follow on PR: https://github.com/apache/arrow-datafusion/pull/4407


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [arrow-datafusion] alamb commented on pull request #4301: Use tournament loser tree for k-way sort-merging

Posted by GitBox <gi...@apache.org>.
alamb commented on PR #4301:
URL: https://github.com/apache/arrow-datafusion/pull/4301#issuecomment-1322210468

   cc @tustvold 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [arrow-datafusion] viirya commented on a diff in pull request #4301: Use tournament loser tree for k-way sort-merging

Posted by GitBox <gi...@apache.org>.
viirya commented on code in PR #4301:
URL: https://github.com/apache/arrow-datafusion/pull/4301#discussion_r1028605200


##########
datafusion/core/src/physical_plan/sorts/sort_preserving_merge.rs:
##########
@@ -551,17 +558,46 @@ impl SortPreservingMergeStream {
         if self.aborted {
             return Poll::Ready(None);
         }
+        let num_streams = self.streams.num_streams();
+
+        // Init all cursors and the loser tree in the first poll
+        if self.loser_tree.is_empty() {
+            // Ensure all non-exhausted streams have a cursor from which
+            // rows can be pulled
+            for i in 0..num_streams {
+                match futures::ready!(self.maybe_poll_stream(cx, i)) {
+                    Ok(_) => {}
+                    Err(e) => {
+                        self.aborted = true;
+                        return Poll::Ready(Some(Err(e)));
+                    }
+                }
+            }
 
-        // Ensure all non-exhausted streams have a cursor from which
-        // rows can be pulled
-        for i in 0..self.streams.num_streams() {
-            match futures::ready!(self.maybe_poll_stream(cx, i)) {
-                Ok(_) => {}
-                Err(e) => {
-                    self.aborted = true;
-                    return Poll::Ready(Some(Err(e)));
+            // Init loser tree
+            self.loser_tree.resize(num_streams, usize::MAX);
+            for i in 0..num_streams {
+                let mut winner = i;
+                let mut cmp_node = (num_streams + i) / 2;
+                while cmp_node != 0 && self.loser_tree[cmp_node] != usize::MAX {
+                    let challenger = self.loser_tree[cmp_node];
+                    let challenger_win =
+                        match (&self.cursors[winner], &self.cursors[challenger]) {
+                            (None, _) => true,
+                            (_, None) => false,
+                            (Some(winner), Some(challenger)) => challenger < winner,
+                        };
+                    if challenger_win {
+                        self.loser_tree[cmp_node] = winner;
+                        winner = challenger;
+                    } else {
+                        self.loser_tree[cmp_node] = challenger;
+                    }

Review Comment:
   ```suggestion
                       if challenger_win {
                           self.loser_tree[cmp_node] = winner;
                           winner = challenger;
                       }
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [arrow-datafusion] ursabot commented on pull request #4301: Use tournament loser tree for k-way sort-merging, increase merge speed by 50%

Posted by GitBox <gi...@apache.org>.
ursabot commented on PR #4301:
URL: https://github.com/apache/arrow-datafusion/pull/4301#issuecomment-1329708051

   Benchmark runs are scheduled for baseline = 52e198ea456c588cd57d44e9d98b082d0d2bd163 and contender = 0d334cf7d67ecdf2bd095987242085e0d59484f3. 0d334cf7d67ecdf2bd095987242085e0d59484f3 is a master commit associated with this PR. Results will be available as each benchmark for each run completes.
   Conbench compare runs links:
   [Skipped :warning: Benchmarking of arrow-datafusion-commits is not supported on ec2-t3-xlarge-us-east-2] [ec2-t3-xlarge-us-east-2](https://conbench.ursa.dev/compare/runs/8a1d6a9e8144471ea160fbf595429b98...e87c2696f67f401ca5b2010f755161db/)
   [Skipped :warning: Benchmarking of arrow-datafusion-commits is not supported on test-mac-arm] [test-mac-arm](https://conbench.ursa.dev/compare/runs/1f452b0b9e094beeaf7cb42adf4764b1...1682458985cc45edafdc8130921044e5/)
   [Skipped :warning: Benchmarking of arrow-datafusion-commits is not supported on ursa-i9-9960x] [ursa-i9-9960x](https://conbench.ursa.dev/compare/runs/e1dcc9c6029849b0804021114e9c24e3...f6fee4f8665843639b7a63f2317ccd6b/)
   [Skipped :warning: Benchmarking of arrow-datafusion-commits is not supported on ursa-thinkcentre-m75q] [ursa-thinkcentre-m75q](https://conbench.ursa.dev/compare/runs/694098fac7f7436ebe9f500c72b861aa...6724d4d02f8645ad88909256d6eaaa0d/)
   Buildkite builds:
   Supported benchmarks:
   ec2-t3-xlarge-us-east-2: Supported benchmark langs: Python, R. Runs only benchmarks with cloud = True
   test-mac-arm: Supported benchmark langs: C++, Python, R
   ursa-i9-9960x: Supported benchmark langs: Python, R, JavaScript
   ursa-thinkcentre-m75q: Supported benchmark langs: C++, Java
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [arrow-datafusion] alamb commented on pull request #4301: Use tournament loser tree for k-way sort-merging

Posted by GitBox <gi...@apache.org>.
alamb commented on PR #4301:
URL: https://github.com/apache/arrow-datafusion/pull/4301#issuecomment-1328040325

   If  we don't hear back from @richox  by Monday I plan to merge this PR and then create a follow on PR with the proposed changes from @tustvold  and @viirya 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [arrow-datafusion] tustvold commented on a diff in pull request #4301: Use tournament loser tree for k-way sort-merging

Posted by GitBox <gi...@apache.org>.
tustvold commented on code in PR #4301:
URL: https://github.com/apache/arrow-datafusion/pull/4301#discussion_r1028358218


##########
datafusion/core/src/physical_plan/sorts/sort_preserving_merge.rs:
##########
@@ -570,45 +606,57 @@ impl SortPreservingMergeStream {
         let _timer = elapsed_compute.timer();
 
         loop {
-            match self.heap.pop() {
-                Some(Reverse(mut cursor)) => {
-                    let stream_idx = cursor.stream_idx();
-                    let batch_idx = self.batches[stream_idx].len() - 1;
-                    let row_idx = cursor.advance();
-
-                    let mut cursor_finished = false;
-                    // insert the cursor back to heap if the record batch is not exhausted
-                    if !cursor.is_finished() {
-                        self.heap.push(Reverse(cursor));
-                    } else {
-                        cursor_finished = true;
-                        self.cursor_finished[stream_idx] = true;
+            // Adjust the loser tree if necessary
+            if !self.loser_tree_adjusted {
+                let mut winner = self.loser_tree[0];

Review Comment:
   It might be easier to follow if this was moved into a method called `replace_loser_tree_winner`, perhaps with a link to this GIF - https://en.wikipedia.org/wiki/K-way_merge_algorithm#/media/File:Loser_tree_replacement_selection.gif



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [arrow-datafusion] alamb commented on pull request #4301: Use tournament loser tree for k-way sort-merging, increase merge speed by 50%

Posted by GitBox <gi...@apache.org>.
alamb commented on PR #4301:
URL: https://github.com/apache/arrow-datafusion/pull/4301#issuecomment-1329703852

   I will have a follow on PR up shortly


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [arrow-datafusion] alamb commented on a diff in pull request #4301: Use tournament loser tree for k-way sort-merging, increase merge speed by 50%

Posted by GitBox <gi...@apache.org>.
alamb commented on code in PR #4301:
URL: https://github.com/apache/arrow-datafusion/pull/4301#discussion_r1034040105


##########
datafusion/core/src/physical_plan/sorts/sort_preserving_merge.rs:
##########
@@ -570,45 +606,57 @@ impl SortPreservingMergeStream {
         let _timer = elapsed_compute.timer();
 
         loop {
-            match self.heap.pop() {
-                Some(Reverse(mut cursor)) => {
-                    let stream_idx = cursor.stream_idx();
-                    let batch_idx = self.batches[stream_idx].len() - 1;
-                    let row_idx = cursor.advance();
-
-                    let mut cursor_finished = false;
-                    // insert the cursor back to heap if the record batch is not exhausted
-                    if !cursor.is_finished() {
-                        self.heap.push(Reverse(cursor));
-                    } else {
-                        cursor_finished = true;
-                        self.cursor_finished[stream_idx] = true;
+            // Adjust the loser tree if necessary
+            if !self.loser_tree_adjusted {
+                let mut winner = self.loser_tree[0];
+                match futures::ready!(self.maybe_poll_stream(cx, winner)) {
+                    Ok(_) => {}
+                    Err(e) => {
+                        self.aborted = true;
+                        return Poll::Ready(Some(Err(e)));
                     }
+                }
 
-                    self.in_progress.push(RowIndex {
-                        stream_idx,
-                        batch_idx,
-                        row_idx,
-                    });
-
-                    if self.in_progress.len() == self.batch_size {
-                        return Poll::Ready(Some(self.build_record_batch()));
+                let mut cmp_node = (num_streams + winner) / 2;
+                while cmp_node != 0 {
+                    let challenger = self.loser_tree[cmp_node];
+                    let challenger_win =
+                        match (&self.cursors[winner], &self.cursors[challenger]) {
+                            (None, _) => true,
+                            (_, None) => false,
+                            (Some(winner), Some(challenger)) => challenger < winner,
+                        };
+                    if challenger_win {
+                        self.loser_tree[cmp_node] = winner;
+                        winner = challenger;
                     }
+                    cmp_node /= 2;
+                }
+                self.loser_tree[0] = winner;
+                self.loser_tree_adjusted = true;
+            }
 
-                    // If removed the last row from the cursor, need to fetch a new record
-                    // batch if possible, before looping round again
-                    if cursor_finished {
-                        match futures::ready!(self.maybe_poll_stream(cx, stream_idx)) {
-                            Ok(_) => {}
-                            Err(e) => {
-                                self.aborted = true;
-                                return Poll::Ready(Some(Err(e)));
-                            }
-                        }
-                    }
+            let min_cursor_idx = self.loser_tree[0];

Review Comment:
   I couldn't make this work in https://github.com/apache/arrow-datafusion/pull/4407



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [arrow-datafusion] alamb commented on pull request #4301: Use tournament loser tree for k-way sort-merging, increase merge speed by 50%

Posted by GitBox <gi...@apache.org>.
alamb commented on PR #4301:
URL: https://github.com/apache/arrow-datafusion/pull/4301#issuecomment-1329704231

   Thanks again @richox @tustvold and @viirya 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [arrow-datafusion] tustvold commented on a diff in pull request #4301: Use tournament loser tree for k-way sort-merging

Posted by GitBox <gi...@apache.org>.
tustvold commented on code in PR #4301:
URL: https://github.com/apache/arrow-datafusion/pull/4301#discussion_r1028338480


##########
datafusion/core/src/physical_plan/sorts/sort_preserving_merge.rs:
##########
@@ -570,45 +606,57 @@ impl SortPreservingMergeStream {
         let _timer = elapsed_compute.timer();
 
         loop {
-            match self.heap.pop() {
-                Some(Reverse(mut cursor)) => {
-                    let stream_idx = cursor.stream_idx();
-                    let batch_idx = self.batches[stream_idx].len() - 1;
-                    let row_idx = cursor.advance();
-
-                    let mut cursor_finished = false;
-                    // insert the cursor back to heap if the record batch is not exhausted
-                    if !cursor.is_finished() {
-                        self.heap.push(Reverse(cursor));
-                    } else {
-                        cursor_finished = true;
-                        self.cursor_finished[stream_idx] = true;
+            // Adjust the loser tree if necessary
+            if !self.loser_tree_adjusted {
+                let mut winner = self.loser_tree[0];
+                match futures::ready!(self.maybe_poll_stream(cx, winner)) {
+                    Ok(_) => {}
+                    Err(e) => {
+                        self.aborted = true;
+                        return Poll::Ready(Some(Err(e)));
                     }
+                }
 
-                    self.in_progress.push(RowIndex {
-                        stream_idx,
-                        batch_idx,
-                        row_idx,
-                    });
-
-                    if self.in_progress.len() == self.batch_size {
-                        return Poll::Ready(Some(self.build_record_batch()));
+                let mut cmp_node = (num_streams + winner) / 2;

Review Comment:
   ```suggestion
                   // Replace overall winner by walking tree of losers
                   let mut cmp_node = (num_streams + winner) / 2;
   ```



##########
datafusion/core/src/physical_plan/sorts/sort_preserving_merge.rs:
##########
@@ -323,8 +318,17 @@ pub(crate) struct SortPreservingMergeStream {
     /// An id to uniquely identify the input stream batch
     next_batch_id: usize,
 
-    /// Heap that yields [`SortKeyCursor`] in increasing order
-    heap: BinaryHeap<Reverse<SortKeyCursor>>,
+    /// Vector that holds all [`SortKeyCursor`]s
+    cursors: Vec<Option<SortKeyCursor>>,
+
+    /// The loser tree that always produces the minimum cursor
+    ///
+    /// Node 0 stores the top winner, Nodes 1..num_streams store
+    /// the loser nodes
+    loser_tree: Vec<usize>,
+
+    /// Identify whether the loser tree is adjusted

Review Comment:
   ```suggestion
       /// Identify whether the most recently yielded overall winner has been replaced
       /// within the loser tree, a value of `false` indicates that they overall winner
       /// has been yielded but the loser tree has not been updated
   ```
   
   Or something to make it clearer what adjusted actually means.
   
   FWIW a boolean of `should_replace_winner` or something might be clearer



##########
datafusion/core/src/physical_plan/sorts/sort_preserving_merge.rs:
##########
@@ -551,17 +558,46 @@ impl SortPreservingMergeStream {
         if self.aborted {
             return Poll::Ready(None);
         }
+        let num_streams = self.streams.num_streams();
+
+        // Init all cursors and the loser tree in the first poll
+        if self.loser_tree.is_empty() {
+            // Ensure all non-exhausted streams have a cursor from which

Review Comment:
   It might be easier to follow if this method were split into a method called `init_loser_tree` with a doc comment explaining what it does



##########
datafusion/core/src/physical_plan/sorts/sort_preserving_merge.rs:
##########
@@ -570,45 +606,57 @@ impl SortPreservingMergeStream {
         let _timer = elapsed_compute.timer();
 
         loop {
-            match self.heap.pop() {
-                Some(Reverse(mut cursor)) => {
-                    let stream_idx = cursor.stream_idx();
-                    let batch_idx = self.batches[stream_idx].len() - 1;
-                    let row_idx = cursor.advance();
-
-                    let mut cursor_finished = false;
-                    // insert the cursor back to heap if the record batch is not exhausted
-                    if !cursor.is_finished() {
-                        self.heap.push(Reverse(cursor));
-                    } else {
-                        cursor_finished = true;
-                        self.cursor_finished[stream_idx] = true;
+            // Adjust the loser tree if necessary
+            if !self.loser_tree_adjusted {
+                let mut winner = self.loser_tree[0];

Review Comment:
   It might be easier to follow if this was moved into a method called `replace_loser_tree_winner`



##########
datafusion/core/src/physical_plan/sorts/cursor.rs:
##########
@@ -109,8 +109,14 @@ impl PartialOrd for SortKeyCursor {
 
 impl Ord for SortKeyCursor {
     fn cmp(&self, other: &Self) -> Ordering {
-        self.current()
-            .cmp(&other.current())
-            .then_with(|| self.stream_idx.cmp(&other.stream_idx))
+        match (self.is_finished(), other.is_finished()) {

Review Comment:
   ```suggestion
           // Order finished cursors last
           match (self.is_finished(), other.is_finished()) {
   ```



##########
datafusion/core/src/physical_plan/sorts/sort_preserving_merge.rs:
##########
@@ -570,45 +606,57 @@ impl SortPreservingMergeStream {
         let _timer = elapsed_compute.timer();
 
         loop {
-            match self.heap.pop() {
-                Some(Reverse(mut cursor)) => {
-                    let stream_idx = cursor.stream_idx();
-                    let batch_idx = self.batches[stream_idx].len() - 1;
-                    let row_idx = cursor.advance();
-
-                    let mut cursor_finished = false;
-                    // insert the cursor back to heap if the record batch is not exhausted
-                    if !cursor.is_finished() {
-                        self.heap.push(Reverse(cursor));
-                    } else {
-                        cursor_finished = true;
-                        self.cursor_finished[stream_idx] = true;
+            // Adjust the loser tree if necessary
+            if !self.loser_tree_adjusted {
+                let mut winner = self.loser_tree[0];
+                match futures::ready!(self.maybe_poll_stream(cx, winner)) {
+                    Ok(_) => {}
+                    Err(e) => {
+                        self.aborted = true;
+                        return Poll::Ready(Some(Err(e)));
                     }
+                }
 
-                    self.in_progress.push(RowIndex {
-                        stream_idx,
-                        batch_idx,
-                        row_idx,
-                    });
-
-                    if self.in_progress.len() == self.batch_size {
-                        return Poll::Ready(Some(self.build_record_batch()));
+                let mut cmp_node = (num_streams + winner) / 2;
+                while cmp_node != 0 {
+                    let challenger = self.loser_tree[cmp_node];
+                    let challenger_win =
+                        match (&self.cursors[winner], &self.cursors[challenger]) {
+                            (None, _) => true,
+                            (_, None) => false,
+                            (Some(winner), Some(challenger)) => challenger < winner,
+                        };
+                    if challenger_win {
+                        self.loser_tree[cmp_node] = winner;
+                        winner = challenger;
                     }
+                    cmp_node /= 2;
+                }
+                self.loser_tree[0] = winner;
+                self.loser_tree_adjusted = true;
+            }
 
-                    // If removed the last row from the cursor, need to fetch a new record
-                    // batch if possible, before looping round again
-                    if cursor_finished {
-                        match futures::ready!(self.maybe_poll_stream(cx, stream_idx)) {
-                            Ok(_) => {}
-                            Err(e) => {
-                                self.aborted = true;
-                                return Poll::Ready(Some(Err(e)));
-                            }
-                        }
-                    }
+            let min_cursor_idx = self.loser_tree[0];

Review Comment:
   I think this could be made easier to follow if it were written along the lines of
   
   
   ```
   let min_cursor = self.cursors[min_cursor_idx];
   if min_cursor.is_finished() {
       // All streams are exhausted
       return Poll::Ready((!self.in_progress.is_empty()).then(|| self.build_record_batch()))
   }
   
   self.loser_tree_adjusted = false;
   self.in_progress.push(...)
   if self.in_progress.len() == self.batch_size {
       return Poll::Ready(Some(self.build_record_batch()));
   }
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [arrow-datafusion] alamb commented on pull request #4301: Use tournament loser tree for k-way sort-merging

Posted by GitBox <gi...@apache.org>.
alamb commented on PR #4301:
URL: https://github.com/apache/arrow-datafusion/pull/4301#issuecomment-1323839682

   Thanks @richox @tustvold   and @viirya  -- this looks awesome. Let's wait to see if @richox  would like to respond to comments.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org