Fix code review comments

Signed-off-by: Vijayan Balasubramanian <[email protected]>
VijayanB · Oct 2, 2024 · 145c6a4 · 145c6a4
1 parent 9637fb7
commit 145c6a4
Show file tree

Hide file tree

Showing 2 changed files with 194 additions and 28 deletions.
diff --git a/src/main/java/org/opensearch/knn/index/query/ExactSearcher.java b/src/main/java/org/opensearch/knn/index/query/ExactSearcher.java
@@ -5,8 +5,10 @@
 
 package org.opensearch.knn.index.query;
 
+import com.google.common.base.Predicates;
 import lombok.AllArgsConstructor;
 import lombok.Builder;
+import lombok.NonNull;
 import lombok.Value;
 import lombok.extern.log4j.Log4j2;
 import org.apache.lucene.index.FieldInfo;
@@ -65,7 +67,7 @@ public Map<Integer, Float> searchLeaf(final LeafReaderContext leafReaderContext,
             && exactSearcherContext.getMatchedDocs().cardinality() <= exactSearcherContext.getK()) {
             return scoreAllDocs(iterator);
         }
-        return searchTopK(iterator, exactSearcherContext.getK());
+        return searchTopCandidates(iterator, exactSearcherContext.getK(), Predicates.alwaysTrue());
     }
 
     /**
@@ -104,18 +106,16 @@ private Map<Integer, Float> scoreAllDocs(KNNIterator iterator) throws IOExceptio
         return docToScore;
     }
 
-    private Map<Integer, Float> searchTopCandidates(KNNIterator iterator, int limit, Predicate<Float> filterScore) throws IOException {
+    private Map<Integer, Float> searchTopCandidates(KNNIterator iterator, int limit, @NonNull Predicate<Float> filterScore)
+        throws IOException {
         // Creating min heap and init with MAX DocID and Score as -INF.
         final HitQueue queue = new HitQueue(limit, true);
         ScoreDoc topDoc = queue.top();
         final Map<Integer, Float> docToScore = new HashMap<>();
         int docId;
         while ((docId = iterator.nextDoc()) != DocIdSetIterator.NO_MORE_DOCS) {
             final float currentScore = iterator.score();
-            if (filterScore != null && Predicate.not(filterScore).test(currentScore)) {
-                continue;
-            }
-            if (currentScore > topDoc.score) {
+            if (filterScore.test(currentScore) && currentScore > topDoc.score) {
                 topDoc.score = currentScore;
                 topDoc.doc = docId;
                 // As the HitQueue is min heap, updating top will bring the doc with -INF score or worst score we
@@ -138,10 +138,6 @@ private Map<Integer, Float> searchTopCandidates(KNNIterator iterator, int limit,
         return docToScore;
     }
 
-    private Map<Integer, Float> searchTopK(KNNIterator iterator, int k) throws IOException {
-        return searchTopCandidates(iterator, k, null);
-    }
-
     private Map<Integer, Float> filterDocsByMinScore(ExactSearcherContext context, KNNIterator iterator, float minScore)
         throws IOException {
         int maxResultWindow = context.getKnnQuery().getContext().getMaxResultWindow();