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

Fix failing BaseKnnVectorQueryTestCase#testTimeout (#13283) 

* Fix failing BaseKnnVectorQueryTestCase#testTimeout

* Also fix ParentBlockJoinKnnVectorQueryTestCase#testTimeout

---------

Co-authored-by: Kaival Parikh <kaivalp2000@gmail.com>
This commit is contained in:
Kaival Parikh 2024-04-11 09:13:01 +05:30 committed by GitHub
parent 203cf7b553
commit 13cf882677
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 51 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

52
lucene/core/src/test/org/apache/lucene/search/BaseKnnVectorQueryTestCase.java
View File

@ -43,6 +43,8 @@ import org.apache.lucene.index.QueryTimeout;
import org.apache.lucene.index.StoredFields;
import org.apache.lucene.index.Term;
import org.apache.lucene.index.VectorSimilarityFunction;
import org.apache.lucene.search.knn.KnnCollectorManager;
import org.apache.lucene.search.knn.TopKnnCollectorManager;
import org.apache.lucene.store.Directory;
import org.apache.lucene.tests.analysis.MockAnalyzer;
import org.apache.lucene.tests.codecs.asserting.AssertingCodec;
@ -766,8 +768,50 @@ abstract class BaseKnnVectorQueryTestCase extends LuceneTestCase {
}
}
/** Test functionality of {@link TimeLimitingKnnCollectorManager}. */
public void testTimeLimitingKnnCollectorManager() 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);
KnnCollectorManager delegate = new TopKnnCollectorManager(3, searcher);
// A collector manager with no timeout
TimeLimitingKnnCollectorManager noTimeoutManager =
new TimeLimitingKnnCollectorManager(delegate, null);
KnnCollector noTimeoutCollector =
noTimeoutManager.newCollector(Integer.MAX_VALUE, searcher.leafContexts.getFirst());
// Check that a normal collector is created without timeout
assertTrue(noTimeoutCollector instanceof TopKnnCollector);
noTimeoutCollector.collect(0, 0);
assertFalse(noTimeoutCollector.earlyTerminated());
// Check that results are complete
TopDocs noTimeoutTopDocs = noTimeoutCollector.topDocs();
assertEquals(TotalHits.Relation.EQUAL_TO, noTimeoutTopDocs.totalHits.relation);
assertEquals(1, noTimeoutTopDocs.scoreDocs.length);
// A collector manager that immediately times out
TimeLimitingKnnCollectorManager timeoutManager =
new TimeLimitingKnnCollectorManager(delegate, () -> true);
KnnCollector timeoutCollector =
timeoutManager.newCollector(Integer.MAX_VALUE, searcher.leafContexts.getFirst());
// Check that a time limiting collector is created, which returns partial results
assertFalse(timeoutCollector instanceof TopKnnCollector);
timeoutCollector.collect(0, 0);
assertTrue(timeoutCollector.earlyTerminated());
// Check that partial results are returned
TopDocs timeoutTopDocs = timeoutCollector.topDocs();
assertEquals(TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO, timeoutTopDocs.totalHits.relation);
assertEquals(1, timeoutTopDocs.scoreDocs.length);
}
}
/** Test that the query times out correctly. */
@AwaitsFix(bugUrl = "https://github.com/apache/lucene/issues/13272")
public void testTimeout() throws IOException {
try (Directory indexStore =
getIndexStore("field", new float[] {0, 1}, new float[] {1, 2}, new float[] {0, 0});
@ -786,9 +830,9 @@ abstract class BaseKnnVectorQueryTestCase extends LuceneTestCase {
assertEquals(0, searcher.count(exactQuery)); // Same for exact search
searcher.setTimeout(new CountingQueryTimeout(1)); // Only score 1 doc
// Note: This depends on the HNSW graph having just one layer,
// would be 0 in case of multiple layers
assertEquals(1, searcher.count(query)); // Expect only 1 result
// Note: We get partial results when the HNSW graph has 1 layer, but no results for > 1 layer
// because the timeout is exhausted while finding the best entry node for the last level
assertTrue(searcher.count(query) <= 1); // Expect at most 1 result
searcher.setTimeout(new CountingQueryTimeout(1)); // Only score 1 doc
assertEquals(1, searcher.count(exactQuery)); // Expect only 1 result

7
lucene/join/src/test/org/apache/lucene/search/join/ParentBlockJoinKnnVectorQueryTestCase.java
View File

@ -285,7 +285,6 @@ abstract class ParentBlockJoinKnnVectorQueryTestCase extends LuceneTestCase {
}
/** Test that the query times out correctly. */
@AwaitsFix(bugUrl = "https://github.com/apache/lucene/issues/13272")
public void testTimeout() throws IOException {
try (Directory indexStore =
getIndexStore("field", new float[] {0, 1}, new float[] {1, 2}, new float[] {0, 0});
@ -306,9 +305,9 @@ abstract class ParentBlockJoinKnnVectorQueryTestCase extends LuceneTestCase {
assertEquals(0, searcher.count(exactQuery)); // Same for exact search
searcher.setTimeout(new CountingQueryTimeout(1)); // Only score 1 parent
// Note: This depends on the HNSW graph having just one layer,
// would be 0 in case of multiple layers
assertEquals(1, searcher.count(query)); // Expect only 1 result
// Note: We get partial results when the HNSW graph has 1 layer, but no results for > 1 layer
// because the timeout is exhausted while finding the best entry node for the last level
assertTrue(searcher.count(query) <= 1); // Expect at most 1 result
searcher.setTimeout(new CountingQueryTimeout(1)); // Only score 1 parent
assertEquals(1, searcher.count(exactQuery)); // Expect only 1 result