[C++] Improve select_k_unstable performance

[asfimport]

Recently sort_indices was optimized and refactored in ARROW-14165. One of those optimizations could also be applied to select_k_unstable's implementation.

Reporter: David Li / @lidavidm

Related issues:

• [C++][Compute] Implement TopK/BottomK (relates to)
• [C++] Improve table sort performance #2 (relates to)

Original Issue Attachments:

• image-2022-02-16-02-35-42-818.png

PRs and other links:

• GitHub Pull Request #12164
• GitHub Pull Request #12582

_{Note: This issue was originally created as ARROW-14183. Please see the migration documentation for further details.}

[asfimport]

Alvin Chunga Mamani / @AlvinJ15:
The same optimization used in sort of tables was applied. The old implementation that used a fixed heap of size K was changed by doing a sort by batches, merges and selecting the first K indices, but there is no change in terms of improvement.

What could be optimized for an implementation for small K using only a merge of the first K elements per block, and by having a K=10 for example at the time of merging, only the first 10 elements of each block would be merged. And for a slightly larger K, apply the normal merge of the entire block until the total size of the block is less than K, then it would proceed to only mege K elements.

[asfimport]

Antoine Pitrou / @pitrou:
So, first, I think you can assume that K will be small compared to N (for example K=100 vs. N several millions). You probably should confirm this by asking the kernels team.

Maintaining a K-element max-heap is O(N log K). You're unlikely to beat this by much if you switch to a sorting approach with complexity O(N log N), even if the sort is more carefully optimized than the heap maintenance.

So I think you should try to explore this approach:

• for each batch, create a K-element max-heap of the top (or bottom) elements in the batch; since the batch is contiguous, you will not be paying the cost of chunked index access
• compute the final result from the K-element heaps: naively, we can create a new K-element max-heap and feed it from the existing onces (I don't know if there are better algorithms to exploit the fact that the input is already a bunch of max-heaps; but that's a second order concern IMHO)

[asfimport]

Alvin Chunga Mamani / @AlvinJ15:
In #12582 was implemented a RadixSort per batches(taking advantage of contiguous elements in memory as same of table-sort) + Merge K first Elements, and the results seems to be equal to the current implementation of use a Heap (just a ~5% more fast in some cases). With this implementation the RadixSort takes most of the time because it sort in batches the whole elements, instead of that I can implement just a Merge K first elements and for small values of K against thousands/millions of elements, but maybe it will be same complexity of the Heap, because it will sort N logK until K < length(current_block)
CC: @lidavidm @pitrou  

[asfimport]

Antoine Pitrou / @pitrou:
Well, K should most of the time be a lot smaller than the block length (if K is typically 10 or 100, for example).

Another idea for the RecordBatch K-top (the first step of a Table K-top, before the K-merge) is to avoid the cost of multi-column compares by doing a sort of Radix Heap.

Do a (K+1)-top on the current (first) batch column. There are two possible outcomes:

• if top[K-1] != top[K], then you don't need to handle other columns since there is no tie
• if top[K-1] == top[K], we need to recurse into the next column by:
  1. finding the first J such that top[J]==top[K]
  2. finding all indices in the current column such that the value is equal to top[K]
  3. running a (K-J+1)-top on the next batch column with the row indices found in step 2) above

[asfimport]

Apache Arrow JIRA Bot:
This issue was last updated over 90 days ago, which may be an indication it is no longer being actively worked. To better reflect the current state, the issue is being unassigned per project policy. Please feel free to re-take assignment of the issue if it is being actively worked, or if you plan to start that work soon.