diff --git a/CHANGELOG.md b/CHANGELOG.md index 1adcb23033..089d2366b0 100644 --- 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 index 5b6029766c..9b78315f48 100644 --- 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 scoreAllDocs(KNNIterator iterator) throws IOExceptio } private Map 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 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) {