You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@lucene.apache.org by vi...@apache.org on 2024/04/11 03:43:06 UTC
(lucene) branch main updated: Fix failing BaseKnnVectorQueryTestCase#testTimeout (#13283)
This is an automated email from the ASF dual-hosted git repository.
vigyasharma pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/lucene.git
The following commit(s) were added to refs/heads/main by this push:
new 13cf882677c Fix failing BaseKnnVectorQueryTestCase#testTimeout (#13283)
13cf882677c is described below
commit 13cf882677c14fee4bb299da05ebbdedaf5fe4fc
Author: Kaival Parikh <46...@users.noreply.github.com>
AuthorDate: Thu Apr 11 09:13:01 2024 +0530
Fix failing BaseKnnVectorQueryTestCase#testTimeout (#13283)
* Fix failing BaseKnnVectorQueryTestCase#testTimeout
* Also fix ParentBlockJoinKnnVectorQueryTestCase#testTimeout
---------
Co-authored-by: Kaival Parikh <ka...@gmail.com>
---
.../lucene/search/BaseKnnVectorQueryTestCase.java | 52 ++++++++++++++++++++--
.../ParentBlockJoinKnnVectorQueryTestCase.java | 7 ++-
2 files changed, 51 insertions(+), 8 deletions(-)
diff --git a/lucene/core/src/test/org/apache/lucene/search/BaseKnnVectorQueryTestCase.java b/lucene/core/src/test/org/apache/lucene/search/BaseKnnVectorQueryTestCase.java
index acb50aa69e2..2ae0ae14a29 100644
--- a/lucene/core/src/test/org/apache/lucene/search/BaseKnnVectorQueryTestCase.java
+++ b/lucene/core/src/test/org/apache/lucene/search/BaseKnnVectorQueryTestCase.java
@@ -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
diff --git a/lucene/join/src/test/org/apache/lucene/search/join/ParentBlockJoinKnnVectorQueryTestCase.java b/lucene/join/src/test/org/apache/lucene/search/join/ParentBlockJoinKnnVectorQueryTestCase.java
index 35addcff3f7..7cec3b71592 100644
--- a/lucene/join/src/test/org/apache/lucene/search/join/ParentBlockJoinKnnVectorQueryTestCase.java
+++ b/lucene/join/src/test/org/apache/lucene/search/join/ParentBlockJoinKnnVectorQueryTestCase.java
@@ -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