Segregate segments based on search type

For exact search, it is not required to perform
qunatization during rescore with oversamples.
However, to avoid normalization between segments from
approx search and exact search, we will first identify
segments that needs approxsearch and will perform oversamples
and, at end, after rescore, we will add scores from segments that
will perform exact search.

Signed-off-by: Vijayan Balasubramanian <balasvij@amazon.com>

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

@@ -153,6 +153,26 @@ public Map<Integer, Float> searchLeaf(LeafReaderContext context, int k) throws I
         return docIdsToScoreMap;
     }
 
+    /**
+     * For given {@link LeafReaderContext}, this api will return will KNNWeight perform exact search or not
+     * always. This decision is based on two properties, 1) if there are no native engine files in segments,
+     * exact search will always be performed, 2) if number of docs after filter is less than 'k'
+     * @param context
+     * @return
+     * @throws IOException
+     */
+    public boolean isExactSearchPreferred(LeafReaderContext context) throws IOException {
+        final BitSet filterBitSet = getFilteredDocsBitSet(context);
+        int cardinality = filterBitSet.cardinality();
+        if (isFilteredExactSearchPreferred(cardinality)) {
+            return true;
+        }
+        if (isMissingNativeEngineFiles(context)) {
+            return true;
+        }
+        return false;
+    }
+
     private BitSet getFilteredDocsBitSet(final LeafReaderContext ctx) throws IOException {
         if (this.filterWeight == null) {
             return new FixedBitSet(0);

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

@@ -65,15 +65,27 @@ public Weight createWeight(IndexSearcher indexSearcher, ScoreMode scoreMode, flo
             boolean isShardLevelRescoringEnabled = KNNSettings.isShardLevelRescoringEnabledForDiskBasedVector(knnQuery.getIndexName());
             int dimension = knnQuery.getQueryVector().length;
             int firstPassK = rescoreContext.getFirstPassK(finalK, isShardLevelRescoringEnabled, dimension);
-            perLeafResults = doSearch(indexSearcher, leafReaderContexts, knnWeight, firstPassK);
+            // split segments into whether exact search will be performed or not
+            List<LeafReaderContext> exactSearchSegments = new ArrayList<>();
+            List<LeafReaderContext> approxSearchSegments = new ArrayList<>();
+            for (LeafReaderContext leafReaderContext : leafReaderContexts) {
+                if (knnWeight.isExactSearchPreferred(leafReaderContext)) {
+                    exactSearchSegments.add(leafReaderContext);
+                } else {
+                    approxSearchSegments.add(leafReaderContext);
+                }
+            }
+            perLeafResults = doSearch(indexSearcher, approxSearchSegments, knnWeight, firstPassK);
             if (isShardLevelRescoringEnabled == true) {
                 ResultUtil.reduceToTopK(perLeafResults, firstPassK);
             }
 
             StopWatch stopWatch = new StopWatch().start();
             perLeafResults = doRescore(indexSearcher, knnWeight, perLeafResults, finalK);
             long rescoreTime = stopWatch.stop().totalTime().millis();
-            log.debug("Rescoring results took {} ms. oversampled k:{}, segments:{}", rescoreTime, firstPassK, leafReaderContexts.size());
+            log.debug("Rescoring results took {} ms. oversampled k:{}, segments:{}", rescoreTime, firstPassK, perLeafResults.size());
+            // do exact search on rest of segments and append to result lists
+            perLeafResults.addAll(doExactSearch(indexSearcher, knnWeight, exactSearchSegments));
         }
         ResultUtil.reduceToTopK(perLeafResults, finalK);
         TopDocs[] topDocs = new TopDocs[perLeafResults.size()];
@@ -127,6 +139,28 @@ private List<Map.Entry<LeafReaderContext, Map<Integer, Float>>> doRescore(
         return indexSearcher.getTaskExecutor().invokeAll(rescoreTasks);
     }
 
+    private List<Map.Entry<LeafReaderContext, Map<Integer, Float>>> doExactSearch(
+        final IndexSearcher indexSearcher,
+        KNNWeight knnWeight,
+        List<LeafReaderContext> leafReaderContexts
+    ) throws IOException {
+        List<Callable<Map.Entry<LeafReaderContext, Map<Integer, Float>>>> exactSearchTasks = new ArrayList<>(leafReaderContexts.size());
+        for (LeafReaderContext context : leafReaderContexts) {
+            exactSearchTasks.add(() -> {
+                final ExactSearcher.ExactSearcherContext exactSearcherContext = ExactSearcher.ExactSearcherContext.builder()
+                    // setting to false because we want to do exact search on full precision vectors
+                    .useQuantizedVectorsForSearch(false)
+                    .k(knnQuery.getK())
+                    .knnQuery(knnQuery)
+                    .isParentHits(true)
+                    .build();
+                final Map<Integer, Float> searchResults = knnWeight.exactSearch(context, exactSearcherContext);
+                return new AbstractMap.SimpleEntry<>(context, searchResults);
+            });
+        }
+        return indexSearcher.getTaskExecutor().invokeAll(exactSearchTasks);
+    }
+
     private Query createDocAndScoreQuery(IndexReader reader, TopDocs topK) {
         int len = topK.scoreDocs.length;
         Arrays.sort(topK.scoreDocs, Comparator.comparingInt(a -> a.doc));