Make the HitQueue size more appropriate for KNN exact search (#13184)

Currently, when performing KNN exact search, we consistently set the HitQueue size to k. However, there may be instances where the number of candidates is actually lower than k.
apache · Mar 19, 2024 · 7a08eea · 7a08eea
1 parent d393b9d
commit 7a08eea
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 3 deletions.
diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
@@ -222,6 +222,7 @@ Optimizations
 
 * GITHUB#13121: Speedup multi-segment HNSW graph search for diversifying child kNN queries. Builds on GITHUB#12962.
   (Ben Trent)
+* GITHUB#13184: Make the HitQueue size more appropriate for KNN exact search (Pan Guixin)
 
 Bug Fixes
 ---------------------

diff --git a/lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java b/lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java
@@ -190,7 +190,8 @@ protected TopDocs exactSearch(LeafReaderContext context, DocIdSetIterator accept
     if (vectorScorer == null) {
       return NO_RESULTS;
     }
-    HitQueue queue = new HitQueue(k, true);
+    final int queueSize = Math.min(k, Math.toIntExact(acceptIterator.cost()));
+    HitQueue queue = new HitQueue(queueSize, true);
     ScoreDoc topDoc = queue.top();
     int doc;
     while ((doc = acceptIterator.nextDoc()) != DocIdSetIterator.NO_MORE_DOCS) {

diff --git a/...e/join/src/java/org/apache/lucene/search/join/DiversifyingChildrenByteKnnVectorQuery.java b/...e/join/src/java/org/apache/lucene/search/join/DiversifyingChildrenByteKnnVectorQuery.java
@@ -98,7 +98,8 @@ protected TopDocs exactSearch(LeafReaderContext context, DocIdSetIterator accept
             parentBitSet,
             query,
             fi.getVectorSimilarityFunction());
-    HitQueue queue = new HitQueue(k, true);
+    final int queueSize = Math.min(k, Math.toIntExact(acceptIterator.cost()));
+    HitQueue queue = new HitQueue(queueSize, true);
     ScoreDoc topDoc = queue.top();
     while (vectorScorer.nextParent() != DocIdSetIterator.NO_MORE_DOCS) {
       float score = vectorScorer.score();

diff --git a/.../join/src/java/org/apache/lucene/search/join/DiversifyingChildrenFloatKnnVectorQuery.java b/.../join/src/java/org/apache/lucene/search/join/DiversifyingChildrenFloatKnnVectorQuery.java
@@ -98,7 +98,8 @@ protected TopDocs exactSearch(LeafReaderContext context, DocIdSetIterator accept
             parentBitSet,
             query,
             fi.getVectorSimilarityFunction());
-    HitQueue queue = new HitQueue(k, true);
+    final int queueSize = Math.min(k, Math.toIntExact(acceptIterator.cost()));
+    HitQueue queue = new HitQueue(queueSize, true);
     ScoreDoc topDoc = queue.top();
     while (vectorScorer.nextParent() != DocIdSetIterator.NO_MORE_DOCS) {
       float score = vectorScorer.score();