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/07/17 19:09:55 UTC
(lucene) 01/04: OptimisticKnnVectorQuery (#14226)
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
commit 61424c489e65b545f62dd138d45efce1bfafedfb
Author: Michael Sokolov <so...@falutin.net>
AuthorDate: Mon May 5 13:59:30 2025 -0400
OptimisticKnnVectorQuery (#14226)
* OptimisticKnnVectorQuery that re-enters search for segments with with strong results
* fix JDK23 lint
* simplify OptimisticKnnVectorQuery to 2 rounds only
* move OptimisticKnnVectorQuery into AbstractKnnVectorQuery
* tidy
* Use TopKnnCollector for second pass so we correctly collect total top K
---------
Co-authored-by: Michael Sokolov <so...@amazon.com>
---
.../lucene/search/AbstractKnnVectorQuery.java | 128 +++++++++++++++++++--
.../org/apache/lucene/search/DocAndScoreQuery.java | 18 ++-
.../apache/lucene/search/SeededKnnVectorQuery.java | 8 +-
.../lucene/search/BaseKnnVectorQueryTestCase.java | 2 +-
.../lucene/search/TestKnnFloatVectorQuery.java | 4 +-
.../search/TestSeededKnnByteVectorQuery.java | 4 +-
6 files changed, 147 insertions(+), 17 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 57e2859e5da..9501e8ca0b2 100644
--- a/lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java
+++ b/lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java
@@ -18,12 +18,17 @@ package org.apache.lucene.search;
import java.io.IOException;
import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Iterator;
import java.util.List;
+import java.util.Map;
import java.util.Objects;
import java.util.concurrent.Callable;
import org.apache.lucene.codecs.KnnVectorsReader;
+import org.apache.lucene.codecs.lucene90.IndexedDISI;
import org.apache.lucene.index.FieldInfo;
import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.KnnVectorValues;
import org.apache.lucene.index.LeafReader;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.QueryTimeout;
@@ -50,6 +55,9 @@ import org.apache.lucene.util.FixedBitSet;
abstract class AbstractKnnVectorQuery extends Query {
private static final TopDocs NO_RESULTS = TopDocsCollector.EMPTY_TOPDOCS;
+ // Constant controlling the degree of additional result exploration done during
+ // pro-rata search of segments.
+ private static final int LAMBDA = 16;
protected final String field;
protected final int k;
@@ -102,22 +110,60 @@ abstract class AbstractKnnVectorQuery extends Query {
new TimeLimitingKnnCollectorManager(
getKnnCollectorManager(k, indexSearcher), indexSearcher.getTimeout());
TaskExecutor taskExecutor = indexSearcher.getTaskExecutor();
- List<LeafReaderContext> leafReaderContexts = reader.leaves();
+ List<LeafReaderContext> leafReaderContexts = new ArrayList<>(reader.leaves());
List<Callable<TopDocs>> tasks = new ArrayList<>(leafReaderContexts.size());
for (LeafReaderContext context : leafReaderContexts) {
tasks.add(() -> searchLeaf(context, filterWeight, knnCollectorManager));
}
- TopDocs[] perLeafResults = taskExecutor.invokeAll(tasks).toArray(TopDocs[]::new);
-
- // Merge sort the results
- TopDocs topK = mergeLeafResults(perLeafResults);
+ Map<Integer, TopDocs> perLeafResults = new HashMap<>();
+ TopDocs topK = runSearchTasks(tasks, taskExecutor, perLeafResults, leafReaderContexts);
+ if (topK.scoreDocs.length > 0 && perLeafResults.size() > 1) {
+ float minTopKScore = topK.scoreDocs[topK.scoreDocs.length - 1].score;
+ TimeLimitingKnnCollectorManager knnCollectorManagerInner =
+ new TimeLimitingKnnCollectorManager(
+ new ReentrantKnnCollectorManager(
+ new TopKnnCollectorManager(k, indexSearcher), perLeafResults),
+ indexSearcher.getTimeout());
+ Iterator<LeafReaderContext> ctxIter = leafReaderContexts.iterator();
+ while (ctxIter.hasNext()) {
+ LeafReaderContext ctx = ctxIter.next();
+ TopDocs perLeaf = perLeafResults.get(ctx.ord);
+ 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));
+ } else {
+ // This leaf is tapped out; discard the context from the active list so we maintain
+ // correspondence between tasks and leaves
+ ctxIter.remove();
+ }
+ }
+ assert leafReaderContexts.size() == tasks.size();
+ assert perLeafResults.size() == reader.leaves().size();
+ topK = runSearchTasks(tasks, taskExecutor, perLeafResults, leafReaderContexts);
+ }
if (topK.scoreDocs.length == 0) {
return new MatchNoDocsQuery();
}
return DocAndScoreQuery.createDocAndScoreQuery(reader, topK);
}
- private TopDocs searchLeaf(
+ private TopDocs runSearchTasks(
+ List<Callable<TopDocs>> tasks,
+ TaskExecutor taskExecutor,
+ Map<Integer, TopDocs> perLeafResults,
+ List<LeafReaderContext> leafReaderContexts)
+ throws IOException {
+ List<TopDocs> taskResults = taskExecutor.invokeAll(tasks);
+ for (int i = 0; i < taskResults.size(); i++) {
+ perLeafResults.put(leafReaderContexts.get(i).ord, taskResults.get(i));
+ }
+ tasks.clear();
+ // Merge sort the results
+ return mergeLeafResults(perLeafResults.values().toArray(TopDocs[]::new));
+ }
+
+ protected TopDocs searchLeaf(
LeafReaderContext ctx,
Weight filterWeight,
TimeLimitingKnnCollectorManager timeLimitingKnnCollectorManager)
@@ -203,7 +249,29 @@ abstract class AbstractKnnVectorQuery extends Query {
}
protected KnnCollectorManager getKnnCollectorManager(int k, IndexSearcher searcher) {
- return new TopKnnCollectorManager(k, 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;
+ }
+
+ /*
+ * Returns perLeafTopK, the expected number (K * leafProportion) of hits in a leaf with the given
+ * proportion of the entire index, plus three standard deviations of a binomial distribution. Math
+ * says there is a 95% probability that this segment's contribution to the global top K hits are
+ * <= perLeafTopK.
+ */
+ private static int perLeafTopKCalculation(int k, float leafProportion) {
+ return (int)
+ Math.max(
+ 1, k * leafProportion + LAMBDA * Math.sqrt(k * leafProportion * (1 - leafProportion)));
}
protected abstract TopDocs approximateSearch(
@@ -283,6 +351,52 @@ abstract class AbstractKnnVectorQuery extends Query {
return TopDocs.merge(k, perLeafResults);
}
+ // forked from SeededKnnVectorQuery.SeededCollectorManager
+ private class ReentrantKnnCollectorManager implements KnnCollectorManager {
+ final KnnCollectorManager knnCollectorManager;
+ final Map<Integer, TopDocs> perLeafResults;
+
+ ReentrantKnnCollectorManager(
+ KnnCollectorManager knnCollectorManager, Map<Integer, TopDocs> perLeafResults) {
+ this.knnCollectorManager = knnCollectorManager;
+ this.perLeafResults = perLeafResults;
+ }
+
+ @Override
+ public KnnCollector newCollector(
+ int visitLimit, KnnSearchStrategy searchStrategy, LeafReaderContext ctx)
+ throws IOException {
+ KnnCollector delegateCollector =
+ knnCollectorManager.newCollector(visitLimit, searchStrategy, ctx);
+ TopDocs seedTopDocs = perLeafResults.get(ctx.ord);
+ VectorScorer scorer = createVectorScorer(ctx, ctx.reader().getFieldInfos().fieldInfo(field));
+ if (seedTopDocs.totalHits.value() == 0 || scorer == null) {
+ // shouldn't happen - we only come here when there are results
+ assert false;
+ // on the other hand, it should be safe to return no results?
+ return delegateCollector;
+ }
+ DocIdSetIterator vectorIterator = scorer.iterator();
+ // Handle sparse
+ if (vectorIterator instanceof IndexedDISI indexedDISI) {
+ vectorIterator = IndexedDISI.asDocIndexIterator(indexedDISI);
+ }
+ // Most underlying iterators are indexed, so we can map the seed docs to the vector docs
+ if (vectorIterator instanceof KnnVectorValues.DocIndexIterator indexIterator) {
+ DocIdSetIterator seedDocs =
+ new SeededKnnVectorQuery.MappedDISI(
+ indexIterator, new SeededKnnVectorQuery.TopDocsDISI(seedTopDocs, ctx));
+ return knnCollectorManager.newCollector(
+ visitLimit,
+ new KnnSearchStrategy.Seeded(seedDocs, seedTopDocs.scoreDocs.length, searchStrategy),
+ ctx);
+ }
+ // could lead to an infinite loop if this ever happens
+ assert false;
+ return delegateCollector;
+ }
+ }
+
@Override
public void visit(QueryVisitor visitor) {
if (visitor.acceptField(field)) {
diff --git a/lucene/core/src/java/org/apache/lucene/search/DocAndScoreQuery.java b/lucene/core/src/java/org/apache/lucene/search/DocAndScoreQuery.java
index c51e66502e6..d86249c778b 100644
--- a/lucene/core/src/java/org/apache/lucene/search/DocAndScoreQuery.java
+++ b/lucene/core/src/java/org/apache/lucene/search/DocAndScoreQuery.java
@@ -33,6 +33,7 @@ class DocAndScoreQuery extends Query {
private final float[] scores;
private final float maxScore;
private final int[] segmentStarts;
+ private final long visited;
private final Object contextIdentity;
/**
@@ -44,15 +45,23 @@ class DocAndScoreQuery extends Query {
* document in each segment. If a segment has no matching documents, it should be assigned the
* index of the next segment that does. There should be a final entry that is always
* docs.length-1.
+ * @param visited the number of graph nodes that were visited, and for which vector distance
+ * scores were evaluated.
* @param contextIdentity an object identifying the reader context that was used to build this
* query
*/
DocAndScoreQuery(
- int[] docs, float[] scores, float maxScore, int[] segmentStarts, Object contextIdentity) {
+ int[] docs,
+ float[] scores,
+ float maxScore,
+ int[] segmentStarts,
+ long visited,
+ Object contextIdentity) {
this.docs = docs;
this.scores = scores;
this.maxScore = maxScore;
this.segmentStarts = segmentStarts;
+ this.visited = visited;
this.contextIdentity = contextIdentity;
}
@@ -159,6 +168,10 @@ class DocAndScoreQuery extends Query {
};
}
+ public long visited() {
+ return visited;
+ }
+
@Override
public String toString(String field) {
return "DocAndScoreQuery[" + docs[0] + ",...][" + scores[0] + ",...]," + maxScore;
@@ -199,7 +212,8 @@ class DocAndScoreQuery extends Query {
scores[i] = topK.scoreDocs[i].score;
}
int[] segmentStarts = findSegmentStarts(reader.leaves(), docs);
- return new DocAndScoreQuery(docs, scores, maxScore, segmentStarts, reader.getContext().id());
+ return new DocAndScoreQuery(
+ docs, scores, maxScore, segmentStarts, topK.totalHits.value(), reader.getContext().id());
}
static int[] findSegmentStarts(List<LeafReaderContext> leaves, int[] docs) {
diff --git a/lucene/core/src/java/org/apache/lucene/search/SeededKnnVectorQuery.java b/lucene/core/src/java/org/apache/lucene/search/SeededKnnVectorQuery.java
index 5b8b37b1c53..2b43b6c0b63 100644
--- a/lucene/core/src/java/org/apache/lucene/search/SeededKnnVectorQuery.java
+++ b/lucene/core/src/java/org/apache/lucene/search/SeededKnnVectorQuery.java
@@ -215,11 +215,11 @@ public class SeededKnnVectorQuery extends AbstractKnnVectorQuery {
return delegate.createVectorScorer(context, fi);
}
- private static class MappedDISI extends DocIdSetIterator {
+ static class MappedDISI extends DocIdSetIterator {
KnnVectorValues.DocIndexIterator indexedDISI;
DocIdSetIterator sourceDISI;
- private MappedDISI(KnnVectorValues.DocIndexIterator indexedDISI, DocIdSetIterator sourceDISI) {
+ MappedDISI(KnnVectorValues.DocIndexIterator indexedDISI, DocIdSetIterator sourceDISI) {
this.indexedDISI = indexedDISI;
this.sourceDISI = sourceDISI;
}
@@ -261,11 +261,11 @@ public class SeededKnnVectorQuery extends AbstractKnnVectorQuery {
}
}
- private static class TopDocsDISI extends DocIdSetIterator {
+ static class TopDocsDISI extends DocIdSetIterator {
private final int[] sortedDocIds;
private int idx = -1;
- private TopDocsDISI(TopDocs topDocs, LeafReaderContext ctx) {
+ TopDocsDISI(TopDocs topDocs, LeafReaderContext ctx) {
sortedDocIds = new int[topDocs.scoreDocs.length];
for (int i = 0; i < topDocs.scoreDocs.length; i++) {
// Remove the doc base as added by the collector
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 d5151069bef..ffcb6ce7a3c 100644
--- a/lucene/core/src/test/org/apache/lucene/search/BaseKnnVectorQueryTestCase.java
+++ b/lucene/core/src/test/org/apache/lucene/search/BaseKnnVectorQueryTestCase.java
@@ -513,7 +513,7 @@ abstract class BaseKnnVectorQueryTestCase extends LuceneTestCase {
assertRandomConsistency(false);
}
- @AwaitsFix(bugUrl = "https://github.com/apache/lucene/issues/14180")
+ // @AwaitsFix(bugUrl = "https://github.com/apache/lucene/issues/14180")
public void testRandomConsistencyMultiThreaded() throws IOException {
assertRandomConsistency(true);
}
diff --git a/lucene/core/src/test/org/apache/lucene/search/TestKnnFloatVectorQuery.java b/lucene/core/src/test/org/apache/lucene/search/TestKnnFloatVectorQuery.java
index 2262513b1f3..4914c13f44e 100644
--- a/lucene/core/src/test/org/apache/lucene/search/TestKnnFloatVectorQuery.java
+++ b/lucene/core/src/test/org/apache/lucene/search/TestKnnFloatVectorQuery.java
@@ -234,11 +234,13 @@ public class TestKnnFloatVectorQuery extends BaseKnnVectorQueryTestCase {
int[] segments = DocAndScoreQuery.findSegmentStarts(indexReader.leaves(), docs);
DocAndScoreQuery query =
- new DocAndScoreQuery(docs, scores, maxScore, segments, indexReader.getContext().id());
+ new DocAndScoreQuery(
+ docs, scores, maxScore, segments, scoreDocs.length, indexReader.getContext().id());
final Weight w = query.createWeight(searcher, ScoreMode.TOP_SCORES, 1.0f);
TopDocs topDocs = searcher.search(query, 100);
assertEquals(scoreDocs.length, topDocs.totalHits.value());
+ assertEquals(query.visited(), topDocs.totalHits.value());
assertEquals(TotalHits.Relation.EQUAL_TO, topDocs.totalHits.relation());
Arrays.sort(topDocs.scoreDocs, Comparator.comparingInt(scoreDoc -> scoreDoc.doc));
assertEquals(scoreDocs.length, topDocs.scoreDocs.length);
diff --git a/lucene/core/src/test/org/apache/lucene/search/TestSeededKnnByteVectorQuery.java b/lucene/core/src/test/org/apache/lucene/search/TestSeededKnnByteVectorQuery.java
index 025910fc037..169f09db2f6 100644
--- a/lucene/core/src/test/org/apache/lucene/search/TestSeededKnnByteVectorQuery.java
+++ b/lucene/core/src/test/org/apache/lucene/search/TestSeededKnnByteVectorQuery.java
@@ -275,10 +275,10 @@ public class TestSeededKnnByteVectorQuery extends BaseKnnVectorQueryTestCase {
@Override
public KnnCollector newCollector(
- int numVisisted, KnnSearchStrategy searchStrategy, LeafReaderContext context)
+ int numVisited, KnnSearchStrategy searchStrategy, LeafReaderContext context)
throws IOException {
KnnCollector knnCollector =
- knnCollectorManager.newCollector(numVisisted, searchStrategy, context);
+ knnCollectorManager.newCollector(numVisited, searchStrategy, context);
if (knnCollector.getSearchStrategy() instanceof KnnSearchStrategy.Seeded seeded) {
if (seedCalls == null && seeded.numberOfEntryPoints() > 0) {
fail("Expected non-seeded collector but received: " + knnCollector);