You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@lucene.apache.org by be...@apache.org on 2025/08/25 16:54:28 UTC
(lucene) branch branch_10x updated: Add optimistic collection to DiversifyingNearestChildren* vector queries (#15063)
This is an automated email from the ASF dual-hosted git repository.
benwtrent 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 586dc103361 Add optimistic collection to DiversifyingNearestChildren* vector queries (#15063)
586dc103361 is described below
commit 586dc103361c99fefa2f67946855691a073e805c
Author: Benjamin Trent <be...@gmail.com>
AuthorDate: Mon Aug 25 12:44:08 2025 -0400
Add optimistic collection to DiversifyingNearestChildren* vector queries (#15063)
Optimistic knn search addresses a major issue where we return inconsistent results due to race conditions in the shared queue previously used over multi-segment search.
It not only returns consistent results, but is generally better overall when it comes to recall & latency.
This PR refactors the optimistic querying in AbstractKnnVectorQuery so that more sub-queries can take advantage of this logic.
Additionally, this refactors PatienceKnnVectorQuery to better fit this API and utilize the underlying multi-segment search logic.
Here are the results for nested vs. baseline. Note how baseline recall fluctuates, where it is consistent in candidate.
closes: https://github.com/apache/lucene/issues/15059
---
lucene/CHANGES.txt | 5 ++-
.../lucene/search/AbstractKnnVectorQuery.java | 37 ++++++++++++++--------
.../search/HnswQueueSaturationCollector.java | 8 ++++-
.../lucene/search/PatienceKnnVectorQuery.java | 24 ++++++++++++--
.../lucene/search/knn/KnnCollectorManager.java | 21 ++++++++++++
.../lucene/search/knn/KnnSearchStrategy.java | 4 ++-
.../lucene/search/knn/TopKnnCollectorManager.java | 29 ++++++++---------
...rsifyingNearestChildrenKnnCollectorManager.java | 32 +++++++++++--------
8 files changed, 110 insertions(+), 50 deletions(-)
diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index b879803abed..da522f379b7 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -63,12 +63,15 @@ Improvements
* GITHUB#14737: Clean up redundant code and comments in query node classes. (Stefan Vodita)
-* GITHUB#14816: Expose search strategy in KNN query
+* GITHUB#14816: Expose search strategy in KNN query (Tommaso Teofili)
* GITHUB#14226: Use optimistic-with-checking KNN Query execution strategy in place of cross-thread
global queue min-score checking. Improves performance and consistency. (Mike Sokolov, Dzung Bui,
Ben Trent)
+* GITHUB#15063: Use optimistic-with-checking KNN Query strategy for `DiversifyingChildren*` queries.
+ (Ben Trent)
+
Optimizations
---------------------
* GITHUB#14932: Switched to GroupVarInt Encoding for HNSW Graph edges, added backwards compatibility (Akira Lonske)
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 cb6529a4baa..b3b3853c5cf 100644
--- a/lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java
+++ b/lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java
@@ -103,28 +103,30 @@ abstract class AbstractKnnVectorQuery extends Query {
filterWeight = null;
}
- KnnCollectorManager knnCollectorManagerInner = getKnnCollectorManager(k, indexSearcher);
- TimeLimitingKnnCollectorManager knnCollectorManager =
- new TimeLimitingKnnCollectorManager(knnCollectorManagerInner, indexSearcher.getTimeout());
+ KnnCollectorManager knnCollectorManager = getKnnCollectorManager(k, indexSearcher);
+ OptimisticKnnCollectorManager optimisticCollectorManager =
+ new OptimisticKnnCollectorManager(k, knnCollectorManager);
+ TimeLimitingKnnCollectorManager timeLimitingKnnCollectorManager =
+ new TimeLimitingKnnCollectorManager(optimisticCollectorManager, indexSearcher.getTimeout());
TaskExecutor taskExecutor = indexSearcher.getTaskExecutor();
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));
+ tasks.add(() -> searchLeaf(context, filterWeight, timeLimitingKnnCollectorManager));
}
Map<Integer, TopDocs> perLeafResults = new HashMap<>();
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
+ && knnCollectorManager.isOptimistic()
// 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 knnCollectorManagerPhase2 =
new TimeLimitingKnnCollectorManager(
new ReentrantKnnCollectorManager(
- new TopKnnCollectorManager(k, indexSearcher), perLeafResults),
+ getKnnCollectorManager(k, indexSearcher), perLeafResults),
indexSearcher.getTimeout());
Iterator<LeafReaderContext> ctxIter = leafReaderContexts.iterator();
while (ctxIter.hasNext()) {
@@ -234,26 +236,33 @@ abstract class AbstractKnnVectorQuery extends Query {
}
protected KnnCollectorManager getKnnCollectorManager(int k, IndexSearcher searcher) {
- return new OptimisticKnnCollectorManager(k);
+ return new TopKnnCollectorManager(k, searcher);
}
static class OptimisticKnnCollectorManager implements KnnCollectorManager {
private final int k;
+ private final KnnCollectorManager delegate;
- OptimisticKnnCollectorManager(int k) {
+ OptimisticKnnCollectorManager(int k, KnnCollectorManager delegate) {
this.k = k;
+ this.delegate = delegate;
}
@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);
+ // The delegate supports optimistic collection
+ if (delegate.isOptimistic()) {
+ @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 delegate.newOptimisticCollector(visitedLimit, searchStrategy, context, perLeafTopK);
+ }
+ // We don't support optimistic collection, so just do regular execution path
+ return delegate.newCollector(visitedLimit, searchStrategy, context);
}
}
diff --git a/lucene/core/src/java/org/apache/lucene/search/HnswQueueSaturationCollector.java b/lucene/core/src/java/org/apache/lucene/search/HnswQueueSaturationCollector.java
index a67b28689f8..9148316651f 100644
--- a/lucene/core/src/java/org/apache/lucene/search/HnswQueueSaturationCollector.java
+++ b/lucene/core/src/java/org/apache/lucene/search/HnswQueueSaturationCollector.java
@@ -99,7 +99,13 @@ public class HnswQueueSaturationCollector extends KnnCollector.Decorator {
return new KnnSearchStrategy.Patience(this, hnswStrategy.filteredSearchThreshold());
} else if (delegateStrategy instanceof KnnSearchStrategy.Seeded seededStrategy) {
if (seededStrategy.originalStrategy() instanceof KnnSearchStrategy.Hnsw hnswStrategy) {
- return new KnnSearchStrategy.Patience(this, hnswStrategy.filteredSearchThreshold());
+ // rewrap the underlying HNSW strategy with patience
+ // this way we still use the seeded entry points, filter threshold,
+ // and can utilize patience thresholds
+ KnnSearchStrategy.Patience patienceStrategy =
+ new KnnSearchStrategy.Patience(this, hnswStrategy.filteredSearchThreshold());
+ return new KnnSearchStrategy.Seeded(
+ seededStrategy.entryPoints(), seededStrategy.numberOfEntryPoints(), patienceStrategy);
}
}
return delegateStrategy;
diff --git a/lucene/core/src/java/org/apache/lucene/search/PatienceKnnVectorQuery.java b/lucene/core/src/java/org/apache/lucene/search/PatienceKnnVectorQuery.java
index f911161da0f..a1d95d2cc87 100644
--- a/lucene/core/src/java/org/apache/lucene/search/PatienceKnnVectorQuery.java
+++ b/lucene/core/src/java/org/apache/lucene/search/PatienceKnnVectorQuery.java
@@ -190,7 +190,7 @@ public class PatienceKnnVectorQuery extends AbstractKnnVectorQuery {
@Override
protected KnnCollectorManager getKnnCollectorManager(int k, IndexSearcher searcher) {
- return delegate.getKnnCollectorManager(k, searcher);
+ return new PatienceCollectorManager(delegate.getKnnCollectorManager(k, searcher));
}
@Override
@@ -200,8 +200,7 @@ public class PatienceKnnVectorQuery extends AbstractKnnVectorQuery {
int visitedLimit,
KnnCollectorManager knnCollectorManager)
throws IOException {
- return delegate.approximateSearch(
- context, acceptDocs, visitedLimit, new PatienceCollectorManager(knnCollectorManager));
+ return delegate.approximateSearch(context, acceptDocs, visitedLimit, knnCollectorManager);
}
@Override
@@ -273,6 +272,25 @@ public class PatienceKnnVectorQuery extends AbstractKnnVectorQuery {
saturationThreshold,
patience);
}
+
+ @Override
+ public KnnCollector newOptimisticCollector(
+ int visitLimit, KnnSearchStrategy searchStrategy, LeafReaderContext ctx, int k)
+ throws IOException {
+ if (knnCollectorManager.isOptimistic()) {
+ return new HnswQueueSaturationCollector(
+ knnCollectorManager.newOptimisticCollector(visitLimit, searchStrategy, ctx, k),
+ saturationThreshold,
+ patience);
+ } else {
+ return null;
+ }
+ }
+
+ @Override
+ public boolean isOptimistic() {
+ return knnCollectorManager.isOptimistic();
+ }
}
@Override
diff --git a/lucene/core/src/java/org/apache/lucene/search/knn/KnnCollectorManager.java b/lucene/core/src/java/org/apache/lucene/search/knn/KnnCollectorManager.java
index 6e3219f3b66..6c3f5b50c04 100644
--- a/lucene/core/src/java/org/apache/lucene/search/knn/KnnCollectorManager.java
+++ b/lucene/core/src/java/org/apache/lucene/search/knn/KnnCollectorManager.java
@@ -38,4 +38,25 @@ public interface KnnCollectorManager {
KnnCollector newCollector(
int visitedLimit, KnnSearchStrategy searchStrategy, LeafReaderContext context)
throws IOException;
+
+ /**
+ * Return a new {@link KnnCollector} instance, generally with a specific k value, scaled per leaf
+ * statistics
+ *
+ * @param visitedLimit the maximum number of nodes that the search is allowed to visit
+ * @param searchStrategy the optional search strategy configuration
+ * @param context the leaf reader context
+ * @param k the number of neighbors to collect, this is the expected number of results
+ * @return a new KnnCollector instance
+ * @throws IOException if there is an error creating the collector
+ */
+ default KnnCollector newOptimisticCollector(
+ int visitedLimit, KnnSearchStrategy searchStrategy, LeafReaderContext context, int k)
+ throws IOException {
+ return null;
+ }
+
+ default boolean isOptimistic() {
+ return false;
+ }
}
diff --git a/lucene/core/src/java/org/apache/lucene/search/knn/KnnSearchStrategy.java b/lucene/core/src/java/org/apache/lucene/search/knn/KnnSearchStrategy.java
index 0e46a25f231..792f69c37c2 100644
--- a/lucene/core/src/java/org/apache/lucene/search/knn/KnnSearchStrategy.java
+++ b/lucene/core/src/java/org/apache/lucene/search/knn/KnnSearchStrategy.java
@@ -165,7 +165,9 @@ public abstract class KnnSearchStrategy {
}
@Override
- public void nextVectorsBlock() {}
+ public void nextVectorsBlock() {
+ originalStrategy.nextVectorsBlock();
+ }
}
/**
diff --git a/lucene/core/src/java/org/apache/lucene/search/knn/TopKnnCollectorManager.java b/lucene/core/src/java/org/apache/lucene/search/knn/TopKnnCollectorManager.java
index cd3d1cd8e64..8394bacc5c2 100644
--- a/lucene/core/src/java/org/apache/lucene/search/knn/TopKnnCollectorManager.java
+++ b/lucene/core/src/java/org/apache/lucene/search/knn/TopKnnCollectorManager.java
@@ -22,24 +22,15 @@ import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.KnnCollector;
import org.apache.lucene.search.TopKnnCollector;
-import org.apache.lucene.util.hnsw.BlockingFloatHeap;
-/**
- * TopKnnCollectorManager responsible for creating {@link TopKnnCollector} instances. When
- * concurrency is supported, the {@link BlockingFloatHeap} is used to track the global top scores
- * collected across all leaves.
- */
+/** TopKnnCollectorManager responsible for creating {@link TopKnnCollector} instances. */
public class TopKnnCollectorManager implements KnnCollectorManager {
// the number of docs to collect
private final int k;
- // the global score queue used to track the top scores collected across all leaves
- private final BlockingFloatHeap globalScoreQueue;
public TopKnnCollectorManager(int k, IndexSearcher indexSearcher) {
- boolean isMultiSegments = indexSearcher.getIndexReader().leaves().size() > 1;
this.k = k;
- this.globalScoreQueue = isMultiSegments ? new BlockingFloatHeap(k) : null;
}
/**
@@ -52,11 +43,17 @@ public class TopKnnCollectorManager implements KnnCollectorManager {
public KnnCollector newCollector(
int visitedLimit, KnnSearchStrategy searchStrategy, LeafReaderContext context)
throws IOException {
- if (globalScoreQueue == null) {
- return new TopKnnCollector(k, visitedLimit, searchStrategy);
- } else {
- return new MultiLeafKnnCollector(
- k, globalScoreQueue, new TopKnnCollector(k, visitedLimit, searchStrategy));
- }
+ return new TopKnnCollector(k, visitedLimit, searchStrategy);
+ }
+
+ @Override
+ public KnnCollector newOptimisticCollector(
+ int visitedLimit, KnnSearchStrategy searchStrategy, LeafReaderContext context, int k) {
+ return new TopKnnCollector(k, visitedLimit, searchStrategy);
+ }
+
+ @Override
+ public boolean isOptimistic() {
+ return true;
}
}
diff --git a/lucene/join/src/java/org/apache/lucene/search/join/DiversifyingNearestChildrenKnnCollectorManager.java b/lucene/join/src/java/org/apache/lucene/search/join/DiversifyingNearestChildrenKnnCollectorManager.java
index db0e38eafe8..7ebacdd8194 100644
--- a/lucene/join/src/java/org/apache/lucene/search/join/DiversifyingNearestChildrenKnnCollectorManager.java
+++ b/lucene/join/src/java/org/apache/lucene/search/join/DiversifyingNearestChildrenKnnCollectorManager.java
@@ -23,9 +23,7 @@ import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.KnnCollector;
import org.apache.lucene.search.knn.KnnCollectorManager;
import org.apache.lucene.search.knn.KnnSearchStrategy;
-import org.apache.lucene.search.knn.MultiLeafKnnCollector;
import org.apache.lucene.util.BitSet;
-import org.apache.lucene.util.hnsw.BlockingFloatHeap;
/**
* DiversifyingNearestChildrenKnnCollectorManager responsible for creating {@link
@@ -37,7 +35,6 @@ public class DiversifyingNearestChildrenKnnCollectorManager implements KnnCollec
private final int k;
// filter identifying the parent documents.
private final BitSetProducer parentsFilter;
- private final BlockingFloatHeap globalScoreQueue;
/**
* Constructor
@@ -49,8 +46,6 @@ public class DiversifyingNearestChildrenKnnCollectorManager implements KnnCollec
int k, BitSetProducer parentsFilter, IndexSearcher indexSearcher) {
this.k = k;
this.parentsFilter = parentsFilter;
- this.globalScoreQueue =
- indexSearcher.getIndexReader().leaves().size() > 1 ? new BlockingFloatHeap(k) : null;
}
/**
@@ -67,15 +62,24 @@ public class DiversifyingNearestChildrenKnnCollectorManager implements KnnCollec
if (parentBitSet == null) {
return null;
}
- if (globalScoreQueue == null) {
- return new DiversifyingNearestChildrenKnnCollector(
- k, visitedLimit, searchStrategy, parentBitSet);
- } else {
- return new MultiLeafKnnCollector(
- k,
- globalScoreQueue,
- new DiversifyingNearestChildrenKnnCollector(
- k, visitedLimit, searchStrategy, parentBitSet));
+ return new DiversifyingNearestChildrenKnnCollector(
+ k, visitedLimit, searchStrategy, parentBitSet);
+ }
+
+ @Override
+ public KnnCollector newOptimisticCollector(
+ int visitedLimit, KnnSearchStrategy searchStrategy, LeafReaderContext context, int k)
+ throws IOException {
+ BitSet parentBitSet = parentsFilter.getBitSet(context);
+ if (parentBitSet == null) {
+ return null;
}
+ return new DiversifyingNearestChildrenKnnCollector(
+ k, visitedLimit, searchStrategy, parentBitSet);
+ }
+
+ @Override
+ public boolean isOptimistic() {
+ return true;
}
}