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

short-circuit for case there is no oversample

This commit is contained in:
Anh Dung Bui 2024-11-26 15:18:03 +09:00
parent feda6af1a7
commit 3178bbc22e
2 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
lucene/core/src/java/org/apache/lucene/search/RerankKnnFloatVectorQuery.java
View File

@ -55,6 +55,10 @@ public class RerankKnnFloatVectorQuery extends Query {
public Query rewrite(IndexSearcher indexSearcher) throws IOException {
IndexReader reader = indexSearcher.getIndexReader();
Query rewritten = indexSearcher.rewrite(query);
// short-circuit: don't re-rank if we already got all possible results
if (query.getK() <= k) {
return rewritten;
}
Weight weight = indexSearcher.createWeight(rewritten, ScoreMode.COMPLETE_NO_SCORES, 1.0f);
HitQueue queue = new HitQueue(k, false);
for (var leaf : reader.leaves()) {

2
lucene/core/src/test/org/apache/lucene/search/TestRerankKnnFloatVectorQuery.java
View File

@ -91,7 +91,7 @@ public class TestRerankKnnFloatVectorQuery extends LuceneTestCase {
IndexSearcher searcher = new IndexSearcher(reader);
float[] targetVector = randomFloatVector(VECTOR_DIMENSION, random);
int k = 10;
double oversample = 1.0;
double oversample = random.nextFloat(1.5f, 3.0f);
KnnFloatVectorQuery knnQuery =
new KnnFloatVectorQuery(FIELD, targetVector, k + (int) (k * oversample));