Apache
/
lucene
mirror of https://github.com/apache/lucene.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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.
This commit is contained in:
panguixin 2024-03-20 00:24:36 +08:00 committed by GitHub
parent d393b9d039
commit 7a08eeab47
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 7 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
lucene/CHANGES.txt
View File

@ -222,6 +222,7 @@ Optimizations
* GITHUB#13121: Speedup multi-segment HNSW graph search for diversifying child kNN queries. Builds on GITHUB#12962. * GITHUB#13121: Speedup multi-segment HNSW graph search for diversifying child kNN queries. Builds on GITHUB#12962.
(Ben Trent) (Ben Trent)
* GITHUB#13184: Make the HitQueue size more appropriate for KNN exact search (Pan Guixin)
Bug Fixes Bug Fixes
--------------------- ---------------------

3
lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java
View File

@ -190,7 +190,8 @@ abstract class AbstractKnnVectorQuery extends Query {
if (vectorScorer == null) { if (vectorScorer == null) {
return NO_RESULTS; 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(); ScoreDoc topDoc = queue.top();
int doc; int doc;
while ((doc = acceptIterator.nextDoc()) != DocIdSetIterator.NO_MORE_DOCS) { while ((doc = acceptIterator.nextDoc()) != DocIdSetIterator.NO_MORE_DOCS) {

3
lucene/join/src/java/org/apache/lucene/search/join/DiversifyingChildrenByteKnnVectorQuery.java
View File

@ -98,7 +98,8 @@ public class DiversifyingChildrenByteKnnVectorQuery extends KnnByteVectorQuery {
parentBitSet, parentBitSet,
query, query,
fi.getVectorSimilarityFunction()); 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(); ScoreDoc topDoc = queue.top();
while (vectorScorer.nextParent() != DocIdSetIterator.NO_MORE_DOCS) { while (vectorScorer.nextParent() != DocIdSetIterator.NO_MORE_DOCS) {
float score = vectorScorer.score(); float score = vectorScorer.score();

3
lucene/join/src/java/org/apache/lucene/search/join/DiversifyingChildrenFloatKnnVectorQuery.java
View File

@ -98,7 +98,8 @@ public class DiversifyingChildrenFloatKnnVectorQuery extends KnnFloatVectorQuery
parentBitSet, parentBitSet,
query, query,
fi.getVectorSimilarityFunction()); 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(); ScoreDoc topDoc = queue.top();
while (vectorScorer.nextParent() != DocIdSetIterator.NO_MORE_DOCS) { while (vectorScorer.nextParent() != DocIdSetIterator.NO_MORE_DOCS) {
float score = vectorScorer.score(); float score = vectorScorer.score();