Skip to content

Commit

Permalink
Optimize reduceToTopK in ResultUtil by removing pre-filling and reduc…
Browse files Browse the repository at this point in the history
…ing peek calls (#2146)

Signed-off-by: Junqiu Lei <junqiu@amazon.com>
  • Loading branch information
junqiu-lei authored Sep 27, 2024
1 parent eba9d98 commit e0c3afe
Show file tree
Hide file tree
Showing 2 changed files with 4 additions and 4 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
* Add AVX512 support to k-NN for FAISS library [#2069](https://github.com/opensearch-project/k-NN/pull/2069)
### Enhancements
* Add short circuit if no live docs are in segments [#2059](https://github.com/opensearch-project/k-NN/pull/2059)
* Optimize reduceToTopK in ResultUtil by removing pre-filling and reducing peek calls [#2146](https://github.com/opensearch-project/k-NN/pull/2146)
* Update Default Rescore Context based on Dimension [#2149](https://github.com/opensearch-project/k-NN/pull/2149)
### Bug Fixes
* KNN80DocValues should only be considered for BinaryDocValues fields [#2147](https://github.com/opensearch-project/k-NN/pull/2147)
Expand Down
7 changes: 3 additions & 4 deletions src/main/java/org/opensearch/knn/index/query/ResultUtil.java
Original file line number Diff line number Diff line change
Expand Up @@ -33,15 +33,14 @@ public final class ResultUtil {
public static void reduceToTopK(List<Map<Integer, Float>> perLeafResults, int k) {
// Iterate over all scores to get min competitive score
PriorityQueue<Float> topKMinQueue = new PriorityQueue<>(k);
for (int i = 0; i < k; i++) {
topKMinQueue.add(-Float.MAX_VALUE);
}

int count = 0;
for (Map<Integer, Float> perLeafResult : perLeafResults) {
count += perLeafResult.size();
for (Float score : perLeafResult.values()) {
if (topKMinQueue.peek() != null && score > topKMinQueue.peek()) {
if (topKMinQueue.size() < k) {
topKMinQueue.add(score);
} else if (topKMinQueue.peek() != null && score > topKMinQueue.peek()) {
topKMinQueue.poll();
topKMinQueue.add(score);
}
Expand Down

0 comments on commit e0c3afe

Please sign in to comment.