Remove explicit K from exact search context

CHANGELOG.md

@@ -38,3 +38,4 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Refactoring
 * Does not create additional KNNVectorValues in NativeEngines990KNNVectorWriter when quantization is not needed [#2133](https://github.com/opensearch-project/k-NN/pull/2133)
 * Minor refactoring and refactored some unit test [#2167](https://github.com/opensearch-project/k-NN/pull/2167)
+* Remove explicit K from exact search context  [#2212](https://github.com/opensearch-project/k-NN/pull/2212)

src/main/java/org/opensearch/knn/index/query/ExactSearcher.java

@@ -63,11 +63,11 @@ public Map<Integer, Float> searchLeaf(final LeafReaderContext leafReaderContext,
         if (exactSearcherContext.getKnnQuery().getRadius() != null) {
             return doRadialSearch(leafReaderContext, exactSearcherContext, iterator);
         }
-        if (exactSearcherContext.getMatchedDocs() != null
-            && exactSearcherContext.getMatchedDocs().cardinality() <= exactSearcherContext.getK()) {
+        final int k = exactSearcherContext.getKnnQuery().getK();
+        if (exactSearcherContext.getMatchedDocs() != null && exactSearcherContext.getMatchedDocs().cardinality() <= k) {
             return scoreAllDocs(iterator);
         }
-        return searchTopCandidates(iterator, exactSearcherContext.getK(), Predicates.alwaysTrue());
+        return searchTopCandidates(iterator, k, Predicates.alwaysTrue());
     }
 
     /**
@@ -232,7 +232,6 @@ public static class ExactSearcherContext {
          * re-scoring we need to re-score using full precision vectors and not quantized vectors.
          */
         boolean useQuantizedVectorsForSearch;
-        int k;
         BitSet matchedDocs;
         KNNQuery knnQuery;
         /**

src/main/java/org/opensearch/knn/index/query/KNNWeight.java

@@ -143,15 +143,15 @@ public Map<Integer, Float> searchLeaf(LeafReaderContext context, int k) throws I
          * This improves the recall.
          */
         if (isFilteredExactSearchPreferred(cardinality)) {
-            return doExactSearch(context, filterBitSet, k);
+            return doExactSearch(context, filterBitSet);
         }
         Map<Integer, Float> docIdsToScoreMap = doANNSearch(context, filterBitSet, cardinality, k);
         // See whether we have to perform exact search based on approx search results
         // This is required if there are no native engine files or if approximate search returned
         // results less than K, though we have more than k filtered docs
         if (isExactSearchRequire(context, cardinality, docIdsToScoreMap.size())) {
             final BitSet docs = filterWeight != null ? filterBitSet : null;
-            return doExactSearch(context, docs, k);
+            return doExactSearch(context, docs);
         }
         return docIdsToScoreMap;
     }
@@ -208,10 +208,9 @@ private int[] bitSetToIntArray(final BitSet bitSet) {
         return intArray;
     }
 
-    private Map<Integer, Float> doExactSearch(final LeafReaderContext context, final BitSet acceptedDocs, int k) throws IOException {
+    private Map<Integer, Float> doExactSearch(final LeafReaderContext context, final BitSet acceptedDocs) throws IOException {
         final ExactSearcherContextBuilder exactSearcherContextBuilder = ExactSearcher.ExactSearcherContext.builder()
             .isParentHits(true)
-            .k(k)
             // setting to true, so that if quantization details are present we want to do search on the quantized
             // vectors as this flow is used in first pass of search.
             .useQuantizedVectorsForSearch(true)

src/main/java/org/opensearch/knn/index/query/nativelib/NativeEngineKnnVectorQuery.java

@@ -70,7 +70,7 @@ public Weight createWeight(IndexSearcher indexSearcher, ScoreMode scoreMode, flo
             }
 
             StopWatch stopWatch = new StopWatch().start();
-            perLeafResults = doRescore(indexSearcher, leafReaderContexts, knnWeight, perLeafResults, finalK);
+            perLeafResults = doRescore(indexSearcher, leafReaderContexts, knnWeight, perLeafResults);
             long rescoreTime = stopWatch.stop().totalTime().millis();
             log.debug("Rescoring results took {} ms. oversampled k:{}, segments:{}", rescoreTime, firstPassK, leafReaderContexts.size());
         }
@@ -104,8 +104,7 @@ private List<Map<Integer, Float>> doRescore(
         final IndexSearcher indexSearcher,
         List<LeafReaderContext> leafReaderContexts,
         KNNWeight knnWeight,
-        List<Map<Integer, Float>> perLeafResults,
-        int k
+        List<Map<Integer, Float>> perLeafResults
     ) throws IOException {
         List<Callable<Map<Integer, Float>>> rescoreTasks = new ArrayList<>(leafReaderContexts.size());
         for (int i = 0; i < perLeafResults.size(); i++) {
@@ -117,7 +116,6 @@ private List<Map<Integer, Float>> doRescore(
                     .matchedDocs(convertedBitSet)
                     // setting to false because in re-scoring we want to do exact search on full precision vectors
                     .useQuantizedVectorsForSearch(false)
-                    .k(k)
                     .isParentHits(false)
                     .knnQuery(knnQuery)
                     .build();