Optimize searchTopK method in ExactSearcher

Signed-off-by: Junqiu Lei <[email protected]>
junqiu-lei · Sep 24, 2024 · e988ccf · e988ccf
1 parent 189562e
commit e988ccf
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 14 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -22,6 +22,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 searchTopK method in ExactSearcher [#2146](https://github.com/opensearch-project/k-NN/pull/2146)
 ### Bug Fixes
 ### Infrastructure
 ### Documentation

diff --git a/src/main/java/org/opensearch/knn/index/query/ExactSearcher.java b/src/main/java/org/opensearch/knn/index/query/ExactSearcher.java
@@ -68,26 +68,25 @@ private Map<Integer, Float> scoreAllDocs(KNNIterator iterator) throws IOExceptio
     }
 
     private Map<Integer, Float> searchTopK(KNNIterator iterator, int k) throws IOException {
-        // Creating min heap and init with MAX DocID and Score as -INF.
         final HitQueue queue = new HitQueue(k, true);
-        ScoreDoc topDoc = queue.top();
         final Map<Integer, Float> docToScore = new HashMap<>();
         int docId;
+
         while ((docId = iterator.nextDoc()) != DocIdSetIterator.NO_MORE_DOCS) {
-            if (iterator.score() > topDoc.score) {
-                topDoc.score = iterator.score();
-                topDoc.doc = docId;
-                // As the HitQueue is min heap, updating top will bring the doc with -INF score or worst score we
-                // have seen till now on top.
-                topDoc = queue.updateTop();
+            float score = iterator.score();
+
+            if (score < 0) {
+                continue;
             }
-        }
 
-        // If scores are negative we will remove them.
-        // This is done, because there can be negative values in the Heap as we init the heap with Score as -INF.
-        // If filterIds < k, the some values in heap can have a negative score.
-        while (queue.size() > 0 && queue.top().score < 0) {
-            queue.pop();
+            if (queue.size() < k) {
+                queue.add(new ScoreDoc(docId, score));
+            } else if (score > queue.top().score) {
+                ScoreDoc topDoc = queue.top();
+                topDoc.score = score;
+                topDoc.doc = docId;
+                queue.updateTop();
+            }
         }
 
         while (queue.size() > 0) {