You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@lucene.apache.org by so...@apache.org on 2025/08/01 15:03:04 UTC
(lucene) branch branch_10x updated: Don't revisit knn query results unless optimistic collection was used; fixes #15005
This is an automated email from the ASF dual-hosted git repository.
sokolov pushed a commit to branch branch_10x
in repository https://gitbox.apache.org/repos/asf/lucene.git
The following commit(s) were added to refs/heads/branch_10x by this push:
new 7a97c6357cb Don't revisit knn query results unless optimistic collection was used; fixes #15005
7a97c6357cb is described below
commit 7a97c6357cb8a2c4b2d597a201f7954fc71f56ff
Author: Michael Sokolov <so...@amazon.com>
AuthorDate: Thu Jul 31 12:50:43 2025 -0400
Don't revisit knn query results unless optimistic collection was used; fixes #15005
---
.../lucene/search/AbstractKnnVectorQuery.java | 42 ++++++++++++--------
.../ParentBlockJoinKnnVectorQueryTestCase.java | 45 ++++++++++++++++++++++
.../TestParentBlockJoinByteKnnVectorQuery.java | 2 +-
3 files changed, 73 insertions(+), 16 deletions(-)
diff --git a/lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java b/lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java
index f5a73172b3e..87ae6a30f05 100644
--- a/lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java
+++ b/lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java
@@ -106,9 +106,9 @@ abstract class AbstractKnnVectorQuery extends Query {
filterWeight = null;
}
+ KnnCollectorManager knnCollectorManagerInner = getKnnCollectorManager(k, indexSearcher);
TimeLimitingKnnCollectorManager knnCollectorManager =
- new TimeLimitingKnnCollectorManager(
- getKnnCollectorManager(k, indexSearcher), indexSearcher.getTimeout());
+ new TimeLimitingKnnCollectorManager(knnCollectorManagerInner, indexSearcher.getTimeout());
TaskExecutor taskExecutor = indexSearcher.getTaskExecutor();
List<LeafReaderContext> leafReaderContexts = new ArrayList<>(reader.leaves());
List<Callable<TopDocs>> tasks = new ArrayList<>(leafReaderContexts.size());
@@ -119,10 +119,12 @@ abstract class AbstractKnnVectorQuery extends Query {
TopDocs topK = runSearchTasks(tasks, taskExecutor, perLeafResults, leafReaderContexts);
if (topK.scoreDocs.length > 0
&& perLeafResults.size() > 1
+ // only re-enter if we used the optimistic collection
+ && knnCollectorManagerInner instanceof OptimisticKnnCollectorManager
// don't re-enter the search if we early terminated
&& topK.totalHits.relation() == TotalHits.Relation.EQUAL_TO) {
float minTopKScore = topK.scoreDocs[topK.scoreDocs.length - 1].score;
- TimeLimitingKnnCollectorManager knnCollectorManagerInner =
+ TimeLimitingKnnCollectorManager knnCollectorManagerPhase2 =
new TimeLimitingKnnCollectorManager(
new ReentrantKnnCollectorManager(
new TopKnnCollectorManager(k, indexSearcher), perLeafResults),
@@ -134,7 +136,7 @@ abstract class AbstractKnnVectorQuery extends Query {
if (perLeaf.scoreDocs.length > 0
&& perLeaf.scoreDocs[perLeaf.scoreDocs.length - 1].score >= minTopKScore) {
// All this leaf's hits are at or above the global topK min score; explore it further
- tasks.add(() -> searchLeaf(ctx, filterWeight, knnCollectorManagerInner));
+ tasks.add(() -> searchLeaf(ctx, filterWeight, knnCollectorManagerPhase2));
} else {
// This leaf is tapped out; discard the context from the active list so we maintain
// correspondence between tasks and leaves
@@ -256,17 +258,27 @@ abstract class AbstractKnnVectorQuery extends Query {
}
protected KnnCollectorManager getKnnCollectorManager(int k, IndexSearcher searcher) {
- KnnCollectorManager manager =
- (visitedLimit, strategy, context) -> {
- @SuppressWarnings("resource")
- float leafProportion =
- context.reader().maxDoc() / (float) context.parent.reader().maxDoc();
- int perLeafTopK = perLeafTopKCalculation(k, leafProportion);
- // if we divided by zero above, leafProportion can be NaN and then this would be 0
- assert perLeafTopK > 0;
- return new TopKnnCollector(perLeafTopK, visitedLimit, strategy);
- };
- return manager;
+ return new OptimisticKnnCollectorManager(k);
+ }
+
+ static class OptimisticKnnCollectorManager implements KnnCollectorManager {
+ private final int k;
+
+ OptimisticKnnCollectorManager(int k) {
+ this.k = k;
+ }
+
+ @Override
+ public KnnCollector newCollector(
+ int visitedLimit, KnnSearchStrategy searchStrategy, LeafReaderContext context)
+ throws IOException {
+ @SuppressWarnings("resource")
+ float leafProportion = context.reader().maxDoc() / (float) context.parent.reader().maxDoc();
+ int perLeafTopK = perLeafTopKCalculation(k, leafProportion);
+ // if we divided by zero above, leafProportion can be NaN and then this would be 0
+ assert perLeafTopK > 0;
+ return new TopKnnCollector(perLeafTopK, visitedLimit, searchStrategy);
+ }
}
/*
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 bf059b7b369..37c42f92803 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
@@ -22,8 +22,10 @@ import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
+import java.util.HashSet;
import java.util.List;
import java.util.Map;
+import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import org.apache.lucene.document.Document;
@@ -403,6 +405,49 @@ abstract class ParentBlockJoinKnnVectorQueryTestCase extends LuceneTestCase {
}
}
+ public void testTwoSegments() throws IOException {
+ // see https://github.com/apache/lucene/issues/15005
+ int dim = random().nextInt(1, 10);
+ try (Directory d = newDirectory()) {
+ RandomIndexWriter writer = new RandomIndexWriter(random(), d, newIndexWriterConfig());
+ writer.addDocuments(createFamily("a", 2, dim));
+ writer.addDocuments(createFamily("b", 3, dim));
+ writer.commit();
+ writer.addDocuments(createFamily("c", 1, dim));
+ writer.close();
+ try (IndexReader reader = DirectoryReader.open(d)) {
+ BitSetProducer parentFilter =
+ new QueryBitSetProducer(new TermQuery(new Term("docType", "_parent")));
+ IndexSearcher searcher = newSearcher(reader);
+ Query query = getParentJoinKnnQuery("field", randomVector(dim), null, 3, parentFilter);
+ TopDocs results = searcher.search(query, 3);
+ assertEquals(3, results.scoreDocs.length);
+ assertTrue(results.totalHits.value() >= results.scoreDocs.length);
+ Set<String> resultParentIds = new HashSet<>();
+ for (ScoreDoc scoreDoc : results.scoreDocs) {
+ String parentId = reader.storedFields().document(scoreDoc.doc).get("parentId");
+ assertFalse(resultParentIds.contains(parentId));
+ resultParentIds.add(parentId);
+ }
+ assertEquals(Set.of("a", "b", "c"), resultParentIds);
+ }
+ }
+ }
+
+ private List<Document> createFamily(String parentId, int size, int dim) {
+ List<Document> family = new ArrayList<>();
+ for (int i = 0; i < size; i++) {
+ Document doc = new Document();
+ doc.add(getKnnVectorField("field", randomVector(dim)));
+ doc.add(new StoredField("parentId", parentId));
+ family.add(doc);
+ }
+ Document doc = new Document();
+ doc.add(new StringField("docType", "_parent", Field.Store.NO));
+ family.add(doc);
+ return family;
+ }
+
/** Tests with random vectors, number of documents, etc. Uses RandomIndexWriter. */
public void testRandom() throws IOException {
int numDocs = atLeast(100);
diff --git a/lucene/join/src/test/org/apache/lucene/search/join/TestParentBlockJoinByteKnnVectorQuery.java b/lucene/join/src/test/org/apache/lucene/search/join/TestParentBlockJoinByteKnnVectorQuery.java
index 2f5a50b0d7e..4a15d1ec913 100644
--- a/lucene/join/src/test/org/apache/lucene/search/join/TestParentBlockJoinByteKnnVectorQuery.java
+++ b/lucene/join/src/test/org/apache/lucene/search/join/TestParentBlockJoinByteKnnVectorQuery.java
@@ -110,7 +110,7 @@ public class TestParentBlockJoinByteKnnVectorQuery extends ParentBlockJoinKnnVec
@Override
float[] randomVector(int dim) {
- BytesRef v = TestUtil.randomBinaryTerm(random(), TestUtil.nextInt(random(), 1, 100));
+ BytesRef v = TestUtil.randomBinaryTerm(random(), dim);
// clip at -127 to avoid overflow
for (int i = v.offset; i < v.offset + v.length; i++) {
if (v.bytes[i] == -128) {