Add optimistic collection to DiversifyingNearestChildren* vector queries (#15063)

Optimistic knn search addresses a major issue where we return inconsistent results due to race conditions in the shared queue previously used over multi-segment search.

It not only returns consistent results, but is generally better overall when it comes to recall & latency. 

This PR refactors the optimistic querying in AbstractKnnVectorQuery so that more sub-queries can take advantage of this logic. 

Additionally, this refactors PatienceKnnVectorQuery to better fit this API and utilize the underlying multi-segment search logic.


Here are the results for nested vs. baseline. Note how baseline recall fluctuates, where it is consistent in candidate.

closes: #15059

Author: benwtrent

lucene/CHANGES.txt

@@ -48,9 +48,7 @@ Improvements
   This avoids clashes with other ASM versions on classpath because it removes the dependency.
   (Uwe Schindler)
 
-* GITHUB#14226: Use optimistic-with-checking KNN Query execution strategy in place of cross-thread
-  global queue min-score checking. Improves performance and consistency. (Mike Sokolov, Dzung Bui,
-  Ben Trent)
+
 
 * GITHUB#15002: Remove synchronized WeakHashMap from IndexReader. This was added to help prevent
   SIGSEGV with the old unsafe "mmap hack", which has been replaced by MemorySegments.
@@ -171,7 +169,14 @@ Improvements
 
 * GITHUB#14737: Clean up redundant code and comments in query node classes. (Stefan Vodita)
 
-* GITHUB#14816: Expose search strategy in KNN query
+* GITHUB#14816: Expose search strategy in KNN query (Tommaso Teofili)
+
+* GITHUB#14226: Use optimistic-with-checking KNN Query execution strategy in place of cross-thread
+  global queue min-score checking. Improves performance and consistency. (Mike Sokolov, Dzung Bui,
+  Ben Trent)
+
+* GITHUB#15063: Use optimistic-with-checking KNN Query strategy for `DiversifyingChildren*` queries.
+  (Ben Trent)
 
 Optimizations
 ---------------------

lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java

@@ -103,29 +103,31 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException {
       filterWeight = null;
     }
 
-    KnnCollectorManager knnCollectorManagerInner = getKnnCollectorManager(k, indexSearcher);
-    TimeLimitingKnnCollectorManager knnCollectorManager =
-        new TimeLimitingKnnCollectorManager(knnCollectorManagerInner, indexSearcher.getTimeout());
+    KnnCollectorManager knnCollectorManager = getKnnCollectorManager(k, indexSearcher);
+    OptimisticKnnCollectorManager optimisticCollectorManager =
+        new OptimisticKnnCollectorManager(k, knnCollectorManager);
+    TimeLimitingKnnCollectorManager timeLimitingKnnCollectorManager =
+        new TimeLimitingKnnCollectorManager(optimisticCollectorManager, indexSearcher.getTimeout());
     TaskExecutor taskExecutor = indexSearcher.getTaskExecutor();
     List<LeafReaderContext> leafReaderContexts = new ArrayList<>(reader.leaves());
     List<Callable<TopDocs>> tasks = new ArrayList<>(leafReaderContexts.size());
     for (LeafReaderContext context : leafReaderContexts) {
-      tasks.add(() -> searchLeaf(context, filterWeight, knnCollectorManager));
+      tasks.add(() -> searchLeaf(context, filterWeight, timeLimitingKnnCollectorManager));
     }
     Map<Integer, TopDocs> perLeafResults = new HashMap<>();
     TopDocs topK = runSearchTasks(tasks, taskExecutor, perLeafResults, leafReaderContexts);
     int reentryCount = 0;
     if (topK.scoreDocs.length > 0
         && perLeafResults.size() > 1
         // only re-enter if we used the optimistic collection
-        && knnCollectorManagerInner instanceof OptimisticKnnCollectorManager
+        && knnCollectorManager.isOptimistic()
         // don't re-enter the search if we early terminated
         && topK.totalHits.relation() == TotalHits.Relation.EQUAL_TO) {
       float minTopKScore = topK.scoreDocs[topK.scoreDocs.length - 1].score;
       TimeLimitingKnnCollectorManager knnCollectorManagerPhase2 =
           new TimeLimitingKnnCollectorManager(
               new ReentrantKnnCollectorManager(
-                  new TopKnnCollectorManager(k, indexSearcher), perLeafResults),
+                  getKnnCollectorManager(k, indexSearcher), perLeafResults),
               indexSearcher.getTimeout());
       Iterator<LeafReaderContext> ctxIter = leafReaderContexts.iterator();
       while (ctxIter.hasNext()) {
@@ -236,26 +238,33 @@ private TopDocs getLeafResults(
   }
 
   protected KnnCollectorManager getKnnCollectorManager(int k, IndexSearcher searcher) {
-    return new OptimisticKnnCollectorManager(k);
+    return new TopKnnCollectorManager(k, searcher);
   }
 
   static class OptimisticKnnCollectorManager implements KnnCollectorManager {
     private final int k;
+    private final KnnCollectorManager delegate;
 
-    OptimisticKnnCollectorManager(int k) {
+    OptimisticKnnCollectorManager(int k, KnnCollectorManager delegate) {
       this.k = k;
+      this.delegate = delegate;
     }
 
     @Override
     public KnnCollector newCollector(
         int visitedLimit, KnnSearchStrategy searchStrategy, LeafReaderContext context)
         throws IOException {
-      @SuppressWarnings("resource")
-      float leafProportion = context.reader().maxDoc() / (float) context.parent.reader().maxDoc();
-      int perLeafTopK = perLeafTopKCalculation(k, leafProportion);
-      // if we divided by zero above, leafProportion can be NaN and then this would be 0
-      assert perLeafTopK > 0;
-      return new TopKnnCollector(perLeafTopK, visitedLimit, searchStrategy);
+      // The delegate supports optimistic collection
+      if (delegate.isOptimistic()) {
+        @SuppressWarnings("resource")
+        float leafProportion = context.reader().maxDoc() / (float) context.parent.reader().maxDoc();
+        int perLeafTopK = perLeafTopKCalculation(k, leafProportion);
+        // if we divided by zero above, leafProportion can be NaN and then this would be 0
+        assert perLeafTopK > 0;
+        return delegate.newOptimisticCollector(visitedLimit, searchStrategy, context, perLeafTopK);
+      }
+      // We don't support optimistic collection, so just do regular execution path
+      return delegate.newCollector(visitedLimit, searchStrategy, context);
     }
   }
 

lucene/core/src/java/org/apache/lucene/search/HnswQueueSaturationCollector.java

@@ -99,7 +99,13 @@ public KnnSearchStrategy getSearchStrategy() {
       return new KnnSearchStrategy.Patience(this, hnswStrategy.filteredSearchThreshold());
     } else if (delegateStrategy instanceof KnnSearchStrategy.Seeded seededStrategy) {
       if (seededStrategy.originalStrategy() instanceof KnnSearchStrategy.Hnsw hnswStrategy) {
-        return new KnnSearchStrategy.Patience(this, hnswStrategy.filteredSearchThreshold());
+        // rewrap the underlying HNSW strategy with patience
+        // this way we still use the seeded entry points, filter threshold,
+        // and can utilize patience thresholds
+        KnnSearchStrategy.Patience patienceStrategy =
+            new KnnSearchStrategy.Patience(this, hnswStrategy.filteredSearchThreshold());
+        return new KnnSearchStrategy.Seeded(
+            seededStrategy.entryPoints(), seededStrategy.numberOfEntryPoints(), patienceStrategy);
       }
     }
     return delegateStrategy;

lucene/core/src/java/org/apache/lucene/search/PatienceKnnVectorQuery.java

@@ -190,7 +190,7 @@ public String toString(String field) {
 
   @Override
   protected KnnCollectorManager getKnnCollectorManager(int k, IndexSearcher searcher) {
-    return delegate.getKnnCollectorManager(k, searcher);
+    return new PatienceCollectorManager(delegate.getKnnCollectorManager(k, searcher));
   }
 
   @Override
@@ -200,8 +200,7 @@ protected TopDocs approximateSearch(
       int visitedLimit,
       KnnCollectorManager knnCollectorManager)
       throws IOException {
-    return delegate.approximateSearch(
-        context, acceptDocs, visitedLimit, new PatienceCollectorManager(knnCollectorManager));
+    return delegate.approximateSearch(context, acceptDocs, visitedLimit, knnCollectorManager);
   }
 
   @Override
@@ -273,6 +272,25 @@ public KnnCollector newCollector(
           saturationThreshold,
           patience);
     }
+
+    @Override
+    public KnnCollector newOptimisticCollector(
+        int visitLimit, KnnSearchStrategy searchStrategy, LeafReaderContext ctx, int k)
+        throws IOException {
+      if (knnCollectorManager.isOptimistic()) {
+        return new HnswQueueSaturationCollector(
+            knnCollectorManager.newOptimisticCollector(visitLimit, searchStrategy, ctx, k),
+            saturationThreshold,
+            patience);
+      } else {
+        return null;
+      }
+    }
+
+    @Override
+    public boolean isOptimistic() {
+      return knnCollectorManager.isOptimistic();
+    }
   }
 
   @Override

lucene/core/src/java/org/apache/lucene/search/knn/KnnCollectorManager.java

@@ -38,4 +38,25 @@ public interface KnnCollectorManager {
   KnnCollector newCollector(
       int visitedLimit, KnnSearchStrategy searchStrategy, LeafReaderContext context)
       throws IOException;
+
+  /**
+   * Return a new {@link KnnCollector} instance, generally with a specific k value, scaled per leaf
+   * statistics
+   *
+   * @param visitedLimit the maximum number of nodes that the search is allowed to visit
+   * @param searchStrategy the optional search strategy configuration
+   * @param context the leaf reader context
+   * @param k the number of neighbors to collect, this is the expected number of results
+   * @return a new KnnCollector instance
+   * @throws IOException if there is an error creating the collector
+   */
+  default KnnCollector newOptimisticCollector(
+      int visitedLimit, KnnSearchStrategy searchStrategy, LeafReaderContext context, int k)
+      throws IOException {
+    return null;
+  }
+
+  default boolean isOptimistic() {
+    return false;
+  }
 }

lucene/core/src/java/org/apache/lucene/search/knn/KnnSearchStrategy.java

@@ -165,7 +165,9 @@ public int hashCode() {
     }
 
     @Override
-    public void nextVectorsBlock() {}
+    public void nextVectorsBlock() {
+      originalStrategy.nextVectorsBlock();
+    }
   }
 
   /**

lucene/core/src/java/org/apache/lucene/search/knn/TopKnnCollectorManager.java

@@ -22,24 +22,15 @@
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.KnnCollector;
 import org.apache.lucene.search.TopKnnCollector;
-import org.apache.lucene.util.hnsw.BlockingFloatHeap;
 
-/**
- * TopKnnCollectorManager responsible for creating {@link TopKnnCollector} instances. When
- * concurrency is supported, the {@link BlockingFloatHeap} is used to track the global top scores
- * collected across all leaves.
- */
+/** TopKnnCollectorManager responsible for creating {@link TopKnnCollector} instances. */
 public class TopKnnCollectorManager implements KnnCollectorManager {
 
   // the number of docs to collect
   private final int k;
-  // the global score queue used to track the top scores collected across all leaves
-  private final BlockingFloatHeap globalScoreQueue;
 
   public TopKnnCollectorManager(int k, IndexSearcher indexSearcher) {
-    boolean isMultiSegments = indexSearcher.getIndexReader().leaves().size() > 1;
     this.k = k;
-    this.globalScoreQueue = isMultiSegments ? new BlockingFloatHeap(k) : null;
   }
 
   /**
@@ -52,11 +43,17 @@ public TopKnnCollectorManager(int k, IndexSearcher indexSearcher) {
   public KnnCollector newCollector(
       int visitedLimit, KnnSearchStrategy searchStrategy, LeafReaderContext context)
       throws IOException {
-    if (globalScoreQueue == null) {
-      return new TopKnnCollector(k, visitedLimit, searchStrategy);
-    } else {
-      return new MultiLeafKnnCollector(
-          k, globalScoreQueue, new TopKnnCollector(k, visitedLimit, searchStrategy));
-    }
+    return new TopKnnCollector(k, visitedLimit, searchStrategy);
+  }
+
+  @Override
+  public KnnCollector newOptimisticCollector(
+      int visitedLimit, KnnSearchStrategy searchStrategy, LeafReaderContext context, int k) {
+    return new TopKnnCollector(k, visitedLimit, searchStrategy);
+  }
+
+  @Override
+  public boolean isOptimistic() {
+    return true;
   }
 }

lucene/join/src/java/org/apache/lucene/search/join/DiversifyingNearestChildrenKnnCollectorManager.java

@@ -23,9 +23,7 @@
 import org.apache.lucene.search.KnnCollector;
 import org.apache.lucene.search.knn.KnnCollectorManager;
 import org.apache.lucene.search.knn.KnnSearchStrategy;
-import org.apache.lucene.search.knn.MultiLeafKnnCollector;
 import org.apache.lucene.util.BitSet;
-import org.apache.lucene.util.hnsw.BlockingFloatHeap;
 
 /**
  * DiversifyingNearestChildrenKnnCollectorManager responsible for creating {@link
@@ -37,7 +35,6 @@ public class DiversifyingNearestChildrenKnnCollectorManager implements KnnCollec
   private final int k;
   // filter identifying the parent documents.
   private final BitSetProducer parentsFilter;
-  private final BlockingFloatHeap globalScoreQueue;
 
   /**
    * Constructor
@@ -49,8 +46,6 @@ public DiversifyingNearestChildrenKnnCollectorManager(
       int k, BitSetProducer parentsFilter, IndexSearcher indexSearcher) {
     this.k = k;
     this.parentsFilter = parentsFilter;
-    this.globalScoreQueue =
-        indexSearcher.getIndexReader().leaves().size() > 1 ? new BlockingFloatHeap(k) : null;
   }
 
   /**
@@ -67,15 +62,24 @@ public KnnCollector newCollector(
     if (parentBitSet == null) {
       return null;
     }
-    if (globalScoreQueue == null) {
-      return new DiversifyingNearestChildrenKnnCollector(
-          k, visitedLimit, searchStrategy, parentBitSet);
-    } else {
-      return new MultiLeafKnnCollector(
-          k,
-          globalScoreQueue,
-          new DiversifyingNearestChildrenKnnCollector(
-              k, visitedLimit, searchStrategy, parentBitSet));
+    return new DiversifyingNearestChildrenKnnCollector(
+        k, visitedLimit, searchStrategy, parentBitSet);
+  }
+
+  @Override
+  public KnnCollector newOptimisticCollector(
+      int visitedLimit, KnnSearchStrategy searchStrategy, LeafReaderContext context, int k)
+      throws IOException {
+    BitSet parentBitSet = parentsFilter.getBitSet(context);
+    if (parentBitSet == null) {
+      return null;
     }
+    return new DiversifyingNearestChildrenKnnCollector(
+        k, visitedLimit, searchStrategy, parentBitSet);
+  }
+
+  @Override
+  public boolean isOptimistic() {
+    return true;
   }
 }