Make sure KnnVectorQuery applies search boost (#956)

Before, the rewritten query DocAndScoreQuery ignored the boost.
This commit is contained in:
Julie Tibshirani 2022-06-14 08:34:02 -07:00 committed by GitHub
parent 83461601ad
commit 89dbe651cc
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 31 additions and 6 deletions

View File

@ -103,6 +103,8 @@ Bug Fixes
* LUCENE-10605: Fix error in 32bit jvm object alignment gap calculation (Sun Wuqiang)
* GITHUB#956: Make sure KnnVectorQuery applies search boost. (Julie Tibshirani)
Other
---------------------

View File

@ -346,7 +346,7 @@ public class KnnVectorQuery extends Query {
if (found < 0) {
return Explanation.noMatch("not in top " + k);
}
return Explanation.match(scores[found], "within top " + k);
return Explanation.match(scores[found] * boost, "within top " + k);
}
@Override
@ -388,18 +388,18 @@ public class KnnVectorQuery extends Query {
}
@Override
public float getMaxScore(int docid) {
docid += context.docBase;
public float getMaxScore(int docId) {
docId += context.docBase;
float maxScore = 0;
for (int idx = Math.max(0, upTo); idx < upper && docs[idx] <= docid; idx++) {
for (int idx = Math.max(0, upTo); idx < upper && docs[idx] <= docId; idx++) {
maxScore = Math.max(maxScore, scores[idx]);
}
return maxScore;
return maxScore * boost;
}
@Override
public float score() {
return scores[upTo];
return scores[upTo] * boost;
}
@Override

View File

@ -112,6 +112,29 @@ public class TestKnnVectorQuery extends LuceneTestCase {
}
}
public void testSearchBoost() throws IOException {
try (Directory indexStore =
getIndexStore("field", new float[] {0, 1}, new float[] {1, 2}, new float[] {0, 0});
IndexReader reader = DirectoryReader.open(indexStore)) {
IndexSearcher searcher = newSearcher(reader);
Query vectorQuery = new KnnVectorQuery("field", new float[] {0, 0}, 10);
ScoreDoc[] scoreDocs = searcher.search(vectorQuery, 3).scoreDocs;
Query boostQuery = new BoostQuery(vectorQuery, 3.0f);
ScoreDoc[] boostScoreDocs = searcher.search(boostQuery, 3).scoreDocs;
assertEquals(scoreDocs.length, boostScoreDocs.length);
for (int i = 0; i < scoreDocs.length; i++) {
ScoreDoc scoreDoc = scoreDocs[i];
ScoreDoc boostScoreDoc = boostScoreDocs[i];
assertEquals(scoreDoc.doc, boostScoreDoc.doc);
assertEquals(scoreDoc.score * 3.0f, boostScoreDoc.score, 0.001f);
}
}
}
/** Tests that a KnnVectorQuery applies the filter query */
public void testSimpleFilter() throws IOException {
try (Directory indexStore =