Added debug logging, measure execution times of 2nd search.

Signed-off-by: Dooyong Kim <[email protected]>

Author: 0ctopus13prime

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

@@ -24,7 +24,6 @@
 import org.apache.lucene.util.BitSetIterator;
 import org.apache.lucene.util.Bits;
 import org.apache.lucene.util.FixedBitSet;
-import org.opensearch.common.Nullable;
 import org.opensearch.common.StopWatch;
 import org.opensearch.common.lucene.Lucene;
 import org.opensearch.knn.common.FieldInfoExtractor;
@@ -52,6 +51,9 @@
 import static org.opensearch.knn.common.KNNConstants.SPACE_TYPE;
 import static org.opensearch.knn.common.KNNConstants.VECTOR_DATA_TYPE_FIELD;
 
+import static org.opensearch.knn.profile.StopWatchUtils.startStopWatch;
+import static org.opensearch.knn.profile.StopWatchUtils.stopStopWatchAndLog;
+
 /**
  * {@link KNNWeight} serves as a template for implementing approximate nearest neighbor (ANN)
  * and radius search over a native index type, such as Faiss.
@@ -298,9 +300,9 @@ public PerLeafResult searchLeaf(LeafReaderContext context, int k) throws IOExcep
         final SegmentReader reader = Lucene.segmentReader(context.reader());
         final String segmentName = reader.getSegmentName();
 
-        StopWatch stopWatch = startStopWatch();
+        final StopWatch stopWatch = startStopWatch(log);
         final BitSet filterBitSet = getFilteredDocsBitSet(context);
-        stopStopWatchAndLog(stopWatch, "FilterBitSet creation", segmentName);
+        stopStopWatchAndLog(log, stopWatch, "FilterBitSet creation", knnQuery.getShardId(), segmentName, knnQuery.getField());
 
         // Save its cardinality, as the cardinality calculation is expensive.
         final int filterCardinality = filterBitSet.cardinality();
@@ -330,9 +332,10 @@ public PerLeafResult searchLeaf(LeafReaderContext context, int k) throws IOExcep
             );
         }
 
-        StopWatch annStopWatch = startStopWatch();
+        final StopWatch annStopWatch = startStopWatch(log);
         final TopDocs topDocs = approximateSearch(context, filterBitSet, filterCardinality, k);
-        stopStopWatchAndLog(annStopWatch, "ANN search", segmentName);
+        stopStopWatchAndLog(log, stopWatch, "ANN search", knnQuery.getShardId(), segmentName, knnQuery.getField());
+
         if (knnQuery.isExplain()) {
             knnExplanation.addLeafResult(context.id(), topDocs.scoreDocs.length);
         }
@@ -358,14 +361,6 @@ public PerLeafResult searchLeaf(LeafReaderContext context, int k) throws IOExcep
         );
     }
 
-    private void stopStopWatchAndLog(@Nullable final StopWatch stopWatch, final String prefixMessage, String segmentName) {
-        if (log.isDebugEnabled() && stopWatch != null) {
-            stopWatch.stop();
-            final String logMessage = prefixMessage + " shard: [{}], segment: [{}], field: [{}], time in nanos:[{}] ";
-            log.debug(logMessage, knnQuery.getShardId(), segmentName, knnQuery.getField(), stopWatch.totalTime().nanos());
-        }
-    }
-
     protected BitSet getFilteredDocsBitSet(final LeafReaderContext ctx) throws IOException {
         if (this.filterWeight == null) {
             return new FixedBitSet(0);
@@ -593,10 +588,10 @@ protected void addExplainIfRequired(final TopDocs results, final KNNEngine knnEn
      */
     public TopDocs exactSearch(final LeafReaderContext leafReaderContext, final ExactSearcher.ExactSearcherContext exactSearcherContext)
         throws IOException {
-        StopWatch stopWatch = startStopWatch();
+        final StopWatch stopWatch = startStopWatch(log);
         TopDocs exactSearchResults = exactSearcher.searchLeaf(leafReaderContext, exactSearcherContext);
         final SegmentReader reader = Lucene.segmentReader(leafReaderContext.reader());
-        stopStopWatchAndLog(stopWatch, "Exact search", reader.getSegmentName());
+        stopStopWatchAndLog(log, stopWatch, "Exact search", knnQuery.getShardId(), reader.getSegmentName(), knnQuery.getField());
         return exactSearchResults;
     }
 
@@ -713,13 +708,6 @@ private boolean isMissingNativeEngineFiles(LeafReaderContext context) {
         return engineFiles.isEmpty();
     }
 
-    private StopWatch startStopWatch() {
-        if (log.isDebugEnabled()) {
-            return new StopWatch().start();
-        }
-        return null;
-    }
-
     protected int[] getParentIdsArray(final LeafReaderContext context) throws IOException {
         if (knnQuery.getParentsFilter() == null) {
             return null;

src/main/java/org/opensearch/knn/index/query/memoryoptsearch/MemoryOptimizedKNNWeight.java

@@ -31,6 +31,7 @@
 import org.opensearch.knn.index.query.KNNWeight;
 import org.opensearch.knn.index.query.MemoryOptimizedSearchScoreConverter;
 import org.opensearch.lucene.OptimisticKnnCollectorManager;
+import org.opensearch.lucene.ReentrantKnnCollectorManager;
 
 import java.io.IOException;
 
@@ -50,7 +51,7 @@ public class MemoryOptimizedKNNWeight extends KNNWeight {
 
     private final KnnCollectorManager knnCollectorManager;
     @Setter
-    private KnnCollectorManager optimistic2ndKnnCollectorManager;
+    private ReentrantKnnCollectorManager reentrantKNNCollectorManager;
 
     public MemoryOptimizedKNNWeight(KNNQuery query, float boost, final Weight filterWeight, IndexSearcher searcher, Integer k) {
         super(query, boost, filterWeight);
@@ -188,8 +189,8 @@ private TopDocs queryIndex(
         }
 
         // Create a collector + bitset
-        final KnnCollectorManager collectorManager = optimistic2ndKnnCollectorManager != null
-            ? optimistic2ndKnnCollectorManager
+        final KnnCollectorManager collectorManager = reentrantKNNCollectorManager != null
+            ? reentrantKNNCollectorManager
             : knnCollectorManager;
         final KnnCollector knnCollector = collectorManager.newCollector(visitedLimit, DEFAULT_HNSW_SEARCH_STRATEGY, context);
         final AcceptDocs acceptDocs = getAcceptedDocs(reader, cardinality, filterIdsBitSet);

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

@@ -53,6 +53,9 @@
 import java.util.concurrent.Callable;
 import java.util.stream.Collectors;
 
+import static org.opensearch.knn.profile.StopWatchUtils.startStopWatch;
+import static org.opensearch.knn.profile.StopWatchUtils.stopStopWatchAndLog;
+
 /**
  * {@link KNNQuery} executes approximate nearest neighbor search (ANN) on a segment level.
  * {@link NativeEngineKnnVectorQuery} executes approximate nearest neighbor search but gives
@@ -274,13 +277,19 @@ private List<PerLeafResult> doSearch(
 
         // For memory optimized search, it should kick off 2nd search if optimistic
         if (knnQuery.isMemoryOptimizedSearch() && perLeafResults.size() > 1) {
-            run2ndOptimisticSearch(perLeafResults, knnWeight, leafReaderContexts, k, indexSearcher);
+            log.debug(
+                "Running second deep dive search in optimistic while memory optimized search is enabled. perLeafResults.size()={}",
+                perLeafResults.size()
+            );
+            final StopWatch stopWatch = startStopWatch(log);
+            reentrantSearch(perLeafResults, knnWeight, leafReaderContexts, k, indexSearcher);
+            stopStopWatchAndLog(log, stopWatch, "2ndOptimisticSearch", knnQuery.getShardId(), "All Shards", knnQuery.getField());
         }
 
         return perLeafResults;
     }
 
-    private void run2ndOptimisticSearch(
+    private void reentrantSearch(
         final List<PerLeafResult> perLeafResults,
         final KNNWeight knnWeight,
         final List<LeafReaderContext> leafReaderContexts,
@@ -299,7 +308,7 @@ private void run2ndOptimisticSearch(
 
         assert (perLeafResults.size() == leafReaderContexts.size());
 
-        // Get collector manager first
+        // Get memory optimized knn weight first, it's safe get it, we checked it already.
         final MemoryOptimizedKNNWeight memoryOptKNNWeight = (MemoryOptimizedKNNWeight) knnWeight;
 
         // How many results have we collected?
@@ -313,25 +322,24 @@ private void run2ndOptimisticSearch(
             return;
         }
 
-        // Build segment to results table
-        final Map<Integer, TopDocs> segmentOrdToResults = new HashMap<>(leafReaderContexts.size());
-        for (int i = 0; i < perLeafResults.size(); i++) {
-            segmentOrdToResults.put(leafReaderContexts.get(i).ord, perLeafResults.get(i).getResult());
-        }
-
         // Start 2nd deep dive, and get the minimum bar.
         final float minTopKScore = OptimisticSearchStrategyUtils.findKthLargestScore(perLeafResults, knnQuery.getK(), totalResults);
 
         // Select candidate segments for 2nd search. Pick whatever segment returned all vectors whose score values are greater than `kth`
         // value in the merged results.
         final List<Callable<TopDocs>> secondDeepDiveTasks = new ArrayList<>();
         final List<Integer> contextIndices = new ArrayList<>();
+        final Map<Integer, TopDocs> segmentOrdToResults = new HashMap<>();
+
         for (int i = 0; i < leafReaderContexts.size(); ++i) {
             final LeafReaderContext leafReaderContext = leafReaderContexts.get(i);
             final PerLeafResult perLeafResult = perLeafResults.get(i);
-            final TopDocs perLeaf = segmentOrdToResults.get(leafReaderContext.ord);
+            final TopDocs perLeaf = perLeafResults.get(i).getResult();
             if (perLeaf.scoreDocs.length > 0 && perLeafResult.getSearchMode() == PerLeafResult.SearchMode.APPROXIMATE_SEARCH) {
                 if (FORCE_REENTER_TESTING || perLeaf.scoreDocs[perLeaf.scoreDocs.length - 1].score >= minTopKScore) {
+                    // For the target segment, save top results. Which will be used as seeds.
+                    segmentOrdToResults.put(leafReaderContext.ord, perLeaf);
+
                     // All this leaf's hits are at or above the global topK min score; explore it further
                     secondDeepDiveTasks.add(
                         () -> knnWeight.approximateSearch(
@@ -348,24 +356,24 @@ private void run2ndOptimisticSearch(
 
         // Kick off 2nd search tasks
         if (secondDeepDiveTasks.isEmpty() == false) {
-            final ReentrantKnnCollectorManager knnCollectorManagerPhase2 = new ReentrantKnnCollectorManager(
+            final ReentrantKnnCollectorManager reentrantCollectorManager = new ReentrantKnnCollectorManager(
                 new TopKnnCollectorManager(k, indexSearcher),
                 segmentOrdToResults,
                 knnQuery.getQueryVector(),
                 knnQuery.getField()
             );
 
             // Make weight use reentrant collector manager
-            memoryOptKNNWeight.setOptimistic2ndKnnCollectorManager(knnCollectorManagerPhase2);
+            memoryOptKNNWeight.setReentrantKNNCollectorManager(reentrantCollectorManager);
 
             final List<TopDocs> deepDiveTopDocs = indexSearcher.getTaskExecutor().invokeAll(secondDeepDiveTasks);
 
             // Override results for target context
             for (int i = 0; i < deepDiveTopDocs.size(); ++i) {
                 // Override with the new results
-                final TopDocs resultsFrom2ncDeepDive = deepDiveTopDocs.get(i);
+                final TopDocs resultsFromDeepDive = deepDiveTopDocs.get(i);
                 final PerLeafResult perLeafResult = perLeafResults.get(contextIndices.get(i));
-                perLeafResult.setResult(resultsFrom2ncDeepDive);
+                perLeafResult.setResult(resultsFromDeepDive);
             }
         }
     }