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/02/13 16:46:16 UTC
(lucene) branch main updated: Add new Acorn-esque filtered HNSW search heuristic (#14160)
This is an automated email from the ASF dual-hosted git repository.
benwtrent 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 0614fe9ebd7 Add new Acorn-esque filtered HNSW search heuristic (#14160)
0614fe9ebd7 is described below
commit 0614fe9ebd7b46fede610f0bb59b2817de72e7c3
Author: Benjamin Trent <be...@gmail.com>
AuthorDate: Thu Feb 13 11:46:11 2025 -0500
Add new Acorn-esque filtered HNSW search heuristic (#14160)
This is a continuation and completion of the work started by @benchaplin in https://github.com/apache/lucene/pull/14085
The algorithm is fairly simple:
- Only score and then explore vectors that actually match the filtering criteria
- Since this will make the graph even sparser, the search spread is increased to also include the candidate's neighbor neighbors (e.g. generally maxConn * maxConn exploration)
- Additionally, even more scored candidates for a given NSW are considered to combat the increased sparsity
Some of the changes to the baseline Acorn algorithm are:
- There is some general threshold of filtering that bypasses this algorithm altogether. Default suggestion is 60%.
- The number of additional neighbors explored is predicated on the percentage of the immediate neighborhood that is filtered out
- More extended neighborhoods will be explored if the filter is exceptionally restrictive. This attempts to find valid vectors to score and explore for every candidate.
- Only look at the extended neighbors if less than 90% of the current neighborhood matches the filter.
Since this changes the behavior significantly, and there are still some weird edge cases, I am exposing as a parameter within a new idea called `KnnSearchStrategy` that the collectors can provide. This strategy object can be provided to the queries.
closes: https://github.com/apache/lucene/issues/13940
---
lucene/CHANGES.txt | 7 +
.../lucene90/Lucene90HnswVectorsReader.java | 15 +-
.../lucene90/Lucene90OnHeapHnswGraph.java | 15 +-
.../lucene91/Lucene91HnswVectorsReader.java | 17 +-
.../lucene91/Lucene91OnHeapHnswGraph.java | 15 +-
.../lucene92/Lucene92HnswVectorsReader.java | 17 +-
.../lucene94/Lucene94HnswVectorsReader.java | 17 +-
.../lucene95/Lucene95HnswVectorsReader.java | 9 +-
.../lucene94/Lucene94HnswVectorsWriter.java | 5 +
.../lucene95/Lucene95HnswVectorsWriter.java | 7 +-
.../codecs/lucene99/Lucene99HnswVectorsReader.java | 30 ++-
.../codecs/lucene99/Lucene99HnswVectorsWriter.java | 5 +
.../lucene/search/AbstractKnnVectorQuery.java | 16 +-
.../search/AbstractVectorSimilarityQuery.java | 3 +
.../apache/lucene/search/KnnByteVectorQuery.java | 26 ++-
.../org/apache/lucene/search/KnnCollector.java | 4 +-
.../apache/lucene/search/KnnFloatVectorQuery.java | 26 ++-
.../apache/lucene/search/SeededKnnVectorQuery.java | 6 +-
.../lucene/search/VectorSimilarityCollector.java | 3 +-
.../lucene/search/knn/KnnSearchStrategy.java | 105 ++++++++-
.../util/hnsw/FilteredHnswGraphSearcher.java | 252 +++++++++++++++++++++
.../org/apache/lucene/util/hnsw/HnswGraph.java | 14 +-
.../apache/lucene/util/hnsw/HnswGraphBuilder.java | 6 +
.../apache/lucene/util/hnsw/HnswGraphSearcher.java | 69 +++++-
.../apache/lucene/util/hnsw/OnHeapHnswGraph.java | 5 +
.../lucene/search/BaseKnnVectorQueryTestCase.java | 16 +-
.../lucene/search/TestKnnByteVectorQuery.java | 3 +-
.../lucene/search/TestKnnFloatVectorQuery.java | 3 +-
.../search/TestSeededKnnByteVectorQuery.java | 4 +-
.../search/TestSeededKnnFloatVectorQuery.java | 2 +-
.../apache/lucene/util/hnsw/HnswGraphTestCase.java | 7 +-
.../org/apache/lucene/util/hnsw/TestHnswUtil.java | 15 +-
.../DiversifyingChildrenByteKnnVectorQuery.java | 31 ++-
.../DiversifyingChildrenFloatKnnVectorQuery.java | 33 ++-
34 files changed, 712 insertions(+), 96 deletions(-)
diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index 59c585e67e6..126b7d3fd22 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -11,6 +11,9 @@ API Changes
* GITHUB#14165: TieredMergePolicy's maxMergeAtOnce parameter was removed. (Adrien Grand)
+* GITHUB#14160: Adjusts the default HNSW search behavior so that filter optimized search is enabled
+ when 60% or less of vectors pass the pre-filter. (Ben Chaplin, Ben Trent)
+
New Features
---------------------
* GITHUB#14097: Binary partitioning merge policy over float-valued vector field. (Mike Sokolov)
@@ -100,6 +103,10 @@ Optimizations
for fewer objects to be created during indexing and simplifies internally used iterfaces.
(Ben Trent)
+* GITHUB#14160: Adds new filtered HNSW search heuristic based on ACORN-1 algorithm. This improves runtime of filtered HNSW
+ searches as much as 5x. By default this new heuristic isn't used, but is available via the new experimental
+ KnnSearchStrategy interface in knn queries (Ben Chaplin, Ben Trent)
+
* GITHUB#14193: Add Automata.makeCharSet() and makeCharClass() that return minimal DFA
for lists of characters and ranges. Use them in RegExp parser. (Robert Muir)
diff --git a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene90/Lucene90HnswVectorsReader.java b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene90/Lucene90HnswVectorsReader.java
index 698ccff49f7..c586f065dc0 100644
--- a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene90/Lucene90HnswVectorsReader.java
+++ b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene90/Lucene90HnswVectorsReader.java
@@ -463,11 +463,21 @@ public final class Lucene90HnswVectorsReader extends KnnVectorsReader {
arcUpTo = 0;
}
+ @Override
+ public int maxConn() {
+ return UNKNOWN_MAX_CONN;
+ }
+
@Override
public int size() {
return entry.size();
}
+ @Override
+ public int neighborCount() {
+ return arcCount;
+ }
+
@Override
public int nextNeighbor() throws IOException {
if (arcUpTo >= arcCount) {
@@ -488,11 +498,6 @@ public final class Lucene90HnswVectorsReader extends KnnVectorsReader {
throw new UnsupportedOperationException();
}
- @Override
- public int maxConn() {
- throw new UnsupportedOperationException();
- }
-
@Override
public NodesIterator getNodesOnLevel(int level) {
throw new UnsupportedOperationException();
diff --git a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene90/Lucene90OnHeapHnswGraph.java b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene90/Lucene90OnHeapHnswGraph.java
index 835293946dd..1749a77a524 100644
--- a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene90/Lucene90OnHeapHnswGraph.java
+++ b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene90/Lucene90OnHeapHnswGraph.java
@@ -189,18 +189,23 @@ public final class Lucene90OnHeapHnswGraph extends HnswGraph {
}
@Override
- public int numLevels() {
- throw new UnsupportedOperationException();
+ public int neighborCount() {
+ return cur.size();
}
@Override
- public int entryNode() {
+ public int maxConn() {
+ return maxConn;
+ }
+
+ @Override
+ public int numLevels() {
throw new UnsupportedOperationException();
}
@Override
- public int maxConn() {
- return maxConn;
+ public int entryNode() {
+ throw new UnsupportedOperationException();
}
@Override
diff --git a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene91/Lucene91HnswVectorsReader.java b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene91/Lucene91HnswVectorsReader.java
index 80a71196c44..18c105a3e01 100644
--- a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene91/Lucene91HnswVectorsReader.java
+++ b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene91/Lucene91HnswVectorsReader.java
@@ -493,6 +493,7 @@ public final class Lucene91HnswVectorsReader extends KnnVectorsReader {
final int entryNode;
final int size;
final long bytesForConns;
+ final int maxConn;
int arcCount;
int arcUpTo;
@@ -506,6 +507,17 @@ public final class Lucene91HnswVectorsReader extends KnnVectorsReader {
this.size = entry.size();
this.graphOffsetsByLevel = entry.graphOffsetsByLevel;
this.bytesForConns = Math.multiplyExact(Math.addExact(entry.maxConn, 1L), Integer.BYTES);
+ this.maxConn = entry.maxConn;
+ }
+
+ @Override
+ public int maxConn() {
+ return maxConn;
+ }
+
+ @Override
+ public int neighborCount() {
+ return arcCount;
}
@Override
@@ -548,11 +560,6 @@ public final class Lucene91HnswVectorsReader extends KnnVectorsReader {
return entryNode;
}
- @Override
- public int maxConn() {
- return (int) bytesForConns / Integer.BYTES - 1;
- }
-
@Override
public NodesIterator getNodesOnLevel(int level) {
if (level == 0) {
diff --git a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene91/Lucene91OnHeapHnswGraph.java b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene91/Lucene91OnHeapHnswGraph.java
index 3ebbeee8d8a..7effb2dd053 100644
--- a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene91/Lucene91OnHeapHnswGraph.java
+++ b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene91/Lucene91OnHeapHnswGraph.java
@@ -131,6 +131,16 @@ public final class Lucene91OnHeapHnswGraph extends HnswGraph {
upto = -1;
}
+ @Override
+ public int maxConn() {
+ return maxConn;
+ }
+
+ @Override
+ public int neighborCount() {
+ return cur.size();
+ }
+
@Override
public int nextNeighbor() {
if (++upto < cur.size()) {
@@ -160,11 +170,6 @@ public final class Lucene91OnHeapHnswGraph extends HnswGraph {
return entryNode;
}
- @Override
- public int maxConn() {
- return maxConn;
- }
-
@Override
public NodesIterator getNodesOnLevel(int level) {
if (level == 0) {
diff --git a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene92/Lucene92HnswVectorsReader.java b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene92/Lucene92HnswVectorsReader.java
index 213b3fff4f2..1b9ff59dec5 100644
--- a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene92/Lucene92HnswVectorsReader.java
+++ b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene92/Lucene92HnswVectorsReader.java
@@ -401,6 +401,7 @@ public final class Lucene92HnswVectorsReader extends KnnVectorsReader {
final int size;
final long bytesForConns;
final long bytesForConns0;
+ private final int maxConn;
int arcCount;
int arcUpTo;
@@ -416,6 +417,7 @@ public final class Lucene92HnswVectorsReader extends KnnVectorsReader {
this.bytesForConns = Math.multiplyExact(Math.addExact(entry.M, 1L), Integer.BYTES);
this.bytesForConns0 =
Math.multiplyExact(Math.addExact(Math.multiplyExact(entry.M, 2L), 1), Integer.BYTES);
+ this.maxConn = entry.M;
}
@Override
@@ -449,6 +451,16 @@ public final class Lucene92HnswVectorsReader extends KnnVectorsReader {
return arc;
}
+ @Override
+ public int neighborCount() {
+ return arcCount;
+ }
+
+ @Override
+ public int maxConn() {
+ return maxConn;
+ }
+
@Override
public int numLevels() {
return numLevels;
@@ -459,11 +471,6 @@ public final class Lucene92HnswVectorsReader extends KnnVectorsReader {
return entryNode;
}
- @Override
- public int maxConn() {
- return (int) bytesForConns / Integer.BYTES - 1;
- }
-
@Override
public NodesIterator getNodesOnLevel(int level) {
if (level == 0) {
diff --git a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene94/Lucene94HnswVectorsReader.java b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene94/Lucene94HnswVectorsReader.java
index e386c2126be..61ad082da16 100644
--- a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene94/Lucene94HnswVectorsReader.java
+++ b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene94/Lucene94HnswVectorsReader.java
@@ -448,7 +448,7 @@ public final class Lucene94HnswVectorsReader extends KnnVectorsReader {
final int size;
final long bytesForConns;
final long bytesForConns0;
- final int maxConn;
+ private final int maxConn;
int arcCount;
int arcUpTo;
@@ -464,7 +464,7 @@ public final class Lucene94HnswVectorsReader extends KnnVectorsReader {
this.bytesForConns = Math.multiplyExact(Math.addExact(entry.M, 1L), Integer.BYTES);
this.bytesForConns0 =
Math.multiplyExact(Math.addExact(Math.multiplyExact(entry.M, 2L), 1), Integer.BYTES);
- maxConn = entry.M;
+ this.maxConn = entry.M;
}
@Override
@@ -504,13 +504,18 @@ public final class Lucene94HnswVectorsReader extends KnnVectorsReader {
}
@Override
- public int maxConn() {
- return maxConn;
+ public int entryNode() {
+ return entryNode;
}
@Override
- public int entryNode() {
- return entryNode;
+ public int neighborCount() {
+ return arcCount;
+ }
+
+ @Override
+ public int maxConn() {
+ return maxConn;
}
@Override
diff --git a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene95/Lucene95HnswVectorsReader.java b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene95/Lucene95HnswVectorsReader.java
index 99fc82bfc54..3c68a8a6353 100644
--- a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene95/Lucene95HnswVectorsReader.java
+++ b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene95/Lucene95HnswVectorsReader.java
@@ -465,6 +465,7 @@ public final class Lucene95HnswVectorsReader extends KnnVectorsReader implements
int arc;
private final DirectMonotonicReader graphLevelNodeOffsets;
private final long[] graphLevelNodeIndexOffsets;
+ private final int maxConn;
// Allocated to be M*2 to track the current neighbors being explored
private int[] currentNeighborsBuffer;
@@ -480,6 +481,7 @@ public final class Lucene95HnswVectorsReader extends KnnVectorsReader implements
this.graphLevelNodeOffsets =
DirectMonotonicReader.getInstance(entry.offsetsMeta, addressesData);
this.currentNeighborsBuffer = new int[entry.M * 2];
+ this.maxConn = entry.M;
graphLevelNodeIndexOffsets = new long[numLevels];
graphLevelNodeIndexOffsets[0] = 0;
for (int i = 1; i < numLevels; i++) {
@@ -537,9 +539,14 @@ public final class Lucene95HnswVectorsReader extends KnnVectorsReader implements
return entryNode;
}
+ @Override
+ public int neighborCount() {
+ return arcCount;
+ }
+
@Override
public int maxConn() {
- return currentNeighborsBuffer.length / 2;
+ return maxConn;
}
@Override
diff --git a/lucene/backward-codecs/src/test/org/apache/lucene/backward_codecs/lucene94/Lucene94HnswVectorsWriter.java b/lucene/backward-codecs/src/test/org/apache/lucene/backward_codecs/lucene94/Lucene94HnswVectorsWriter.java
index ec7eb048c8f..fdc5d5da146 100644
--- a/lucene/backward-codecs/src/test/org/apache/lucene/backward_codecs/lucene94/Lucene94HnswVectorsWriter.java
+++ b/lucene/backward-codecs/src/test/org/apache/lucene/backward_codecs/lucene94/Lucene94HnswVectorsWriter.java
@@ -344,6 +344,11 @@ public final class Lucene94HnswVectorsWriter extends KnnVectorsWriter {
throw new UnsupportedOperationException("Not supported on a mock graph");
}
+ @Override
+ public int neighborCount() {
+ throw new UnsupportedOperationException("Not supported on a mock graph");
+ }
+
@Override
public NodesIterator getNodesOnLevel(int level) {
if (level == 0) {
diff --git a/lucene/backward-codecs/src/test/org/apache/lucene/backward_codecs/lucene95/Lucene95HnswVectorsWriter.java b/lucene/backward-codecs/src/test/org/apache/lucene/backward_codecs/lucene95/Lucene95HnswVectorsWriter.java
index 55fb3e8ea3e..943e21c560d 100644
--- a/lucene/backward-codecs/src/test/org/apache/lucene/backward_codecs/lucene95/Lucene95HnswVectorsWriter.java
+++ b/lucene/backward-codecs/src/test/org/apache/lucene/backward_codecs/lucene95/Lucene95HnswVectorsWriter.java
@@ -367,10 +367,15 @@ public final class Lucene95HnswVectorsWriter extends KnnVectorsWriter {
}
@Override
- public int maxConn() {
+ public int neighborCount() {
throw new UnsupportedOperationException("Not supported on a mock graph");
}
+ @Override
+ public int maxConn() {
+ return M;
+ }
+
@Override
public NodesIterator getNodesOnLevel(int level) {
if (level == 0) {
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsReader.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsReader.java
index b29f9da5b41..101ebd1eef4 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsReader.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsReader.java
@@ -43,6 +43,7 @@ import org.apache.lucene.store.IOContext;
import org.apache.lucene.store.IndexInput;
import org.apache.lucene.store.RandomAccessInput;
import org.apache.lucene.store.ReadAdvice;
+import org.apache.lucene.util.BitSet;
import org.apache.lucene.util.Bits;
import org.apache.lucene.util.IOSupplier;
import org.apache.lucene.util.IOUtils;
@@ -309,7 +310,6 @@ public final class Lucene99HnswVectorsReader extends KnnVectorsReader
Bits acceptDocs,
IOSupplier<RandomVectorScorer> scorerSupplier)
throws IOException {
-
if (fieldEntry.size() == 0 || knnCollector.k() == 0) {
return;
}
@@ -317,8 +317,23 @@ public final class Lucene99HnswVectorsReader extends KnnVectorsReader
final KnnCollector collector =
new OrdinalTranslatedKnnCollector(knnCollector, scorer::ordToDoc);
final Bits acceptedOrds = scorer.getAcceptOrds(acceptDocs);
- if (knnCollector.k() < scorer.maxOrd()) {
- HnswGraphSearcher.search(scorer, collector, getGraph(fieldEntry), acceptedOrds);
+ HnswGraph graph = getGraph(fieldEntry);
+ boolean doHnsw = knnCollector.k() < scorer.maxOrd();
+ // Take into account if quantized? E.g. some scorer cost?
+ int filteredDocCount = 0;
+ // The approximate number of vectors that would be visited if we did not filter
+ int unfilteredVisit = (int) (Math.log(graph.size()) * knnCollector.k());
+ if (acceptDocs instanceof BitSet bitSet) {
+ // Use approximate cardinality as this is good enough, but ensure we don't exceed the graph
+ // size as that is illogical
+ filteredDocCount = Math.min(bitSet.approximateCardinality(), graph.size());
+ if (unfilteredVisit >= filteredDocCount) {
+ doHnsw = false;
+ }
+ }
+ if (doHnsw) {
+ HnswGraphSearcher.search(
+ scorer, collector, getGraph(fieldEntry), acceptedOrds, filteredDocCount);
} else {
// if k is larger than the number of vectors, we can just iterate over all vectors
// and collect them
@@ -459,6 +474,7 @@ public final class Lucene99HnswVectorsReader extends KnnVectorsReader
int arcCount;
int arcUpTo;
int arc;
+ private final int maxConn;
private final DirectMonotonicReader graphLevelNodeOffsets;
private final long[] graphLevelNodeIndexOffsets;
// Allocated to be M*2 to track the current neighbors being explored
@@ -476,6 +492,7 @@ public final class Lucene99HnswVectorsReader extends KnnVectorsReader
this.graphLevelNodeOffsets =
DirectMonotonicReader.getInstance(entry.offsetsMeta, addressesData);
this.currentNeighborsBuffer = new int[entry.M * 2];
+ this.maxConn = entry.M;
graphLevelNodeIndexOffsets = new long[numLevels];
graphLevelNodeIndexOffsets[0] = 0;
for (int i = 1; i < numLevels; i++) {
@@ -522,6 +539,11 @@ public final class Lucene99HnswVectorsReader extends KnnVectorsReader
return arc;
}
+ @Override
+ public int neighborCount() {
+ return arcCount;
+ }
+
@Override
public int numLevels() throws IOException {
return numLevels;
@@ -529,7 +551,7 @@ public final class Lucene99HnswVectorsReader extends KnnVectorsReader
@Override
public int maxConn() {
- return currentNeighborsBuffer.length >> 1;
+ return maxConn;
}
@Override
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java
index 62c7f5c807b..5701c985c66 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java
@@ -304,6 +304,11 @@ public final class Lucene99HnswVectorsWriter extends KnnVectorsWriter {
throw new UnsupportedOperationException("Not supported on a mock graph");
}
+ @Override
+ public int neighborCount() {
+ throw new UnsupportedOperationException("Not supported on a mock graph");
+ }
+
@Override
public NodesIterator getNodesOnLevel(int level) {
if (level == 0) {
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 e9246a8b575..b53fddae178 100644
--- a/lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java
+++ b/lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java
@@ -32,6 +32,7 @@ import org.apache.lucene.index.LeafReader;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.QueryTimeout;
import org.apache.lucene.search.knn.KnnCollectorManager;
+import org.apache.lucene.search.knn.KnnSearchStrategy;
import org.apache.lucene.search.knn.TopKnnCollectorManager;
import org.apache.lucene.util.BitSet;
import org.apache.lucene.util.BitSetIterator;
@@ -56,14 +57,16 @@ abstract class AbstractKnnVectorQuery extends Query {
protected final String field;
protected final int k;
protected final Query filter;
+ protected final KnnSearchStrategy searchStrategy;
- public AbstractKnnVectorQuery(String field, int k, Query filter) {
+ AbstractKnnVectorQuery(String field, int k, Query filter, KnnSearchStrategy searchStrategy) {
this.field = Objects.requireNonNull(field, "field");
this.k = k;
if (k < 1) {
throw new IllegalArgumentException("k must be at least 1, got: " + k);
}
this.filter = filter;
+ this.searchStrategy = searchStrategy;
}
@Override
@@ -146,7 +149,11 @@ abstract class AbstractKnnVectorQuery extends Query {
// Perform the approximate kNN search
// We pass cost + 1 here to account for the edge case when we explore exactly cost vectors
TopDocs results = approximateSearch(ctx, acceptDocs, cost + 1, timeLimitingKnnCollectorManager);
- if (results.totalHits.relation() == TotalHits.Relation.EQUAL_TO
+ if ((results.totalHits.relation() == TotalHits.Relation.EQUAL_TO
+ // We know that there are more than `k` available docs, if we didn't even get `k`
+ // something weird
+ // happened, and we need to drop to exact search
+ && results.scoreDocs.length >= k)
// Return partial results only when timeout is met
|| (queryTimeout != null && queryTimeout.shouldExit())) {
return results;
@@ -302,7 +309,10 @@ abstract class AbstractKnnVectorQuery extends Query {
if (this == o) return true;
if (o == null || getClass() != o.getClass()) return false;
AbstractKnnVectorQuery that = (AbstractKnnVectorQuery) o;
- return k == that.k && Objects.equals(field, that.field) && Objects.equals(filter, that.filter);
+ return k == that.k
+ && Objects.equals(field, that.field)
+ && Objects.equals(filter, that.filter)
+ && Objects.equals(searchStrategy, that.searchStrategy);
}
@Override
diff --git a/lucene/core/src/java/org/apache/lucene/search/AbstractVectorSimilarityQuery.java b/lucene/core/src/java/org/apache/lucene/search/AbstractVectorSimilarityQuery.java
index 39c81651cf2..df11ae04232 100644
--- a/lucene/core/src/java/org/apache/lucene/search/AbstractVectorSimilarityQuery.java
+++ b/lucene/core/src/java/org/apache/lucene/search/AbstractVectorSimilarityQuery.java
@@ -25,6 +25,7 @@ import org.apache.lucene.index.LeafReader;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.QueryTimeout;
import org.apache.lucene.search.knn.KnnCollectorManager;
+import org.apache.lucene.search.knn.KnnSearchStrategy;
import org.apache.lucene.util.BitSet;
import org.apache.lucene.util.BitSetIterator;
import org.apache.lucene.util.Bits;
@@ -35,6 +36,8 @@ import org.apache.lucene.util.Bits;
* @lucene.experimental
*/
abstract class AbstractVectorSimilarityQuery extends Query {
+ // TODO, switch to optionally use the new strategy
+ static final KnnSearchStrategy.Hnsw DEFAULT_STRATEGY = new KnnSearchStrategy.Hnsw(0);
protected final String field;
protected final float traversalSimilarity, resultSimilarity;
protected final Query filter;
diff --git a/lucene/core/src/java/org/apache/lucene/search/KnnByteVectorQuery.java b/lucene/core/src/java/org/apache/lucene/search/KnnByteVectorQuery.java
index 51d74e0257d..231bcaed191 100644
--- a/lucene/core/src/java/org/apache/lucene/search/KnnByteVectorQuery.java
+++ b/lucene/core/src/java/org/apache/lucene/search/KnnByteVectorQuery.java
@@ -16,6 +16,8 @@
*/
package org.apache.lucene.search;
+import static org.apache.lucene.search.knn.KnnSearchStrategy.Hnsw.DEFAULT;
+
import java.io.IOException;
import java.util.Arrays;
import java.util.Objects;
@@ -26,6 +28,7 @@ import org.apache.lucene.index.FieldInfo;
import org.apache.lucene.index.LeafReader;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.knn.KnnCollectorManager;
+import org.apache.lucene.search.knn.KnnSearchStrategy;
import org.apache.lucene.util.ArrayUtil;
import org.apache.lucene.util.Bits;
@@ -72,7 +75,25 @@ public class KnnByteVectorQuery extends AbstractKnnVectorQuery {
* @throws IllegalArgumentException if <code>k</code> is less than 1
*/
public KnnByteVectorQuery(String field, byte[] target, int k, Query filter) {
- super(field, k, filter);
+ this(field, target, k, filter, DEFAULT);
+ }
+
+ /**
+ * Find the <code>k</code> nearest documents to the target vector according to the vectors in the
+ * given field. <code>target</code> vector.
+ *
+ * @param field a field that has been indexed as a {@link KnnByteVectorField}.
+ * @param target the target of the search
+ * @param k the number of documents to find
+ * @param filter a filter applied before the vector search
+ * @param searchStrategy the search strategy to use. If null, the default strategy will be used.
+ * The underlying format may not support all strategies and is free to ignore the requested
+ * strategy.
+ * @lucene.experimental
+ */
+ public KnnByteVectorQuery(
+ String field, byte[] target, int k, Query filter, KnnSearchStrategy searchStrategy) {
+ super(field, k, filter, searchStrategy);
this.target = Objects.requireNonNull(target, "target");
}
@@ -83,7 +104,8 @@ public class KnnByteVectorQuery extends AbstractKnnVectorQuery {
int visitedLimit,
KnnCollectorManager knnCollectorManager)
throws IOException {
- KnnCollector knnCollector = knnCollectorManager.newCollector(visitedLimit, null, context);
+ KnnCollector knnCollector =
+ knnCollectorManager.newCollector(visitedLimit, searchStrategy, context);
LeafReader reader = context.reader();
ByteVectorValues byteVectorValues = reader.getByteVectorValues(field);
if (byteVectorValues == null) {
diff --git a/lucene/core/src/java/org/apache/lucene/search/KnnCollector.java b/lucene/core/src/java/org/apache/lucene/search/KnnCollector.java
index 43086c4ad57..424711dd54f 100644
--- a/lucene/core/src/java/org/apache/lucene/search/KnnCollector.java
+++ b/lucene/core/src/java/org/apache/lucene/search/KnnCollector.java
@@ -91,9 +91,7 @@ public interface KnnCollector {
/**
* @return the search strategy used by this collector, can be null
*/
- default KnnSearchStrategy getSearchStrategy() {
- return null;
- }
+ KnnSearchStrategy getSearchStrategy();
/**
* KnnCollector.Decorator is the base class for decorators of KnnCollector objects, which extend
diff --git a/lucene/core/src/java/org/apache/lucene/search/KnnFloatVectorQuery.java b/lucene/core/src/java/org/apache/lucene/search/KnnFloatVectorQuery.java
index 8f8a1a72177..5156ff2816f 100644
--- a/lucene/core/src/java/org/apache/lucene/search/KnnFloatVectorQuery.java
+++ b/lucene/core/src/java/org/apache/lucene/search/KnnFloatVectorQuery.java
@@ -16,6 +16,8 @@
*/
package org.apache.lucene.search;
+import static org.apache.lucene.search.knn.KnnSearchStrategy.Hnsw.DEFAULT;
+
import java.io.IOException;
import java.util.Arrays;
import java.util.Objects;
@@ -26,6 +28,7 @@ import org.apache.lucene.index.FloatVectorValues;
import org.apache.lucene.index.LeafReader;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.knn.KnnCollectorManager;
+import org.apache.lucene.search.knn.KnnSearchStrategy;
import org.apache.lucene.util.ArrayUtil;
import org.apache.lucene.util.Bits;
import org.apache.lucene.util.VectorUtil;
@@ -73,7 +76,25 @@ public class KnnFloatVectorQuery extends AbstractKnnVectorQuery {
* @throws IllegalArgumentException if <code>k</code> is less than 1
*/
public KnnFloatVectorQuery(String field, float[] target, int k, Query filter) {
- super(field, k, filter);
+ this(field, target, k, filter, DEFAULT);
+ }
+
+ /**
+ * Find the <code>k</code> nearest documents to the target vector according to the vectors in the
+ * given field. <code>target</code> vector.
+ *
+ * @param field a field that has been indexed as a {@link KnnFloatVectorField}.
+ * @param target the target of the search
+ * @param k the number of documents to find
+ * @param filter a filter applied before the vector search
+ * @param searchStrategy the search strategy to use. If null, the default strategy will be used.
+ * The underlying format may not support all strategies and is free to ignore the requested
+ * strategy.
+ * @lucene.experimental
+ */
+ public KnnFloatVectorQuery(
+ String field, float[] target, int k, Query filter, KnnSearchStrategy searchStrategy) {
+ super(field, k, filter, searchStrategy);
this.target = VectorUtil.checkFinite(Objects.requireNonNull(target, "target"));
}
@@ -84,7 +105,8 @@ public class KnnFloatVectorQuery extends AbstractKnnVectorQuery {
int visitedLimit,
KnnCollectorManager knnCollectorManager)
throws IOException {
- KnnCollector knnCollector = knnCollectorManager.newCollector(visitedLimit, null, context);
+ KnnCollector knnCollector =
+ knnCollectorManager.newCollector(visitedLimit, searchStrategy, context);
LeafReader reader = context.reader();
FloatVectorValues floatVectorValues = reader.getFloatVectorValues(field);
if (floatVectorValues == null) {
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 bcdd4c6bdcf..98d8599f116 100644
--- a/lucene/core/src/java/org/apache/lucene/search/SeededKnnVectorQuery.java
+++ b/lucene/core/src/java/org/apache/lucene/search/SeededKnnVectorQuery.java
@@ -70,7 +70,7 @@ public class SeededKnnVectorQuery extends AbstractKnnVectorQuery {
}
SeededKnnVectorQuery(AbstractKnnVectorQuery knnQuery, Query seed, Weight seedWeight) {
- super(knnQuery.field, knnQuery.k, knnQuery.filter);
+ super(knnQuery.field, knnQuery.k, knnQuery.filter, knnQuery.searchStrategy);
this.delegate = knnQuery;
this.seed = Objects.requireNonNull(seed);
this.seedWeight = seedWeight;
@@ -315,7 +315,9 @@ public class SeededKnnVectorQuery extends AbstractKnnVectorQuery {
DocIdSetIterator seedDocs =
new MappedDISI(indexIterator, new TopDocsDISI(seedTopDocs, ctx));
return knnCollectorManager.newCollector(
- visitLimit, new KnnSearchStrategy.Seeded(seedDocs, seedTopDocs.scoreDocs.length), ctx);
+ visitLimit,
+ new KnnSearchStrategy.Seeded(seedDocs, seedTopDocs.scoreDocs.length, searchStrategy),
+ ctx);
}
return delegateCollector;
}
diff --git a/lucene/core/src/java/org/apache/lucene/search/VectorSimilarityCollector.java b/lucene/core/src/java/org/apache/lucene/search/VectorSimilarityCollector.java
index e7e10c67365..0c0036c20e4 100644
--- a/lucene/core/src/java/org/apache/lucene/search/VectorSimilarityCollector.java
+++ b/lucene/core/src/java/org/apache/lucene/search/VectorSimilarityCollector.java
@@ -40,7 +40,8 @@ class VectorSimilarityCollector extends AbstractKnnCollector {
*/
public VectorSimilarityCollector(
float traversalSimilarity, float resultSimilarity, long visitLimit) {
- super(1, visitLimit);
+ // TODO: add search strategy support
+ super(1, visitLimit, AbstractVectorSimilarityQuery.DEFAULT_STRATEGY);
if (traversalSimilarity > resultSimilarity) {
throw new IllegalArgumentException("traversalSimilarity should be <= resultSimilarity");
}
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 a1d5a7b1b57..bcb938d5c65 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
@@ -16,6 +16,7 @@
*/
package org.apache.lucene.search.knn;
+import java.util.Objects;
import org.apache.lucene.search.DocIdSetIterator;
/**
@@ -24,14 +25,85 @@ import org.apache.lucene.search.DocIdSetIterator;
*
* @lucene.experimental
*/
-public interface KnnSearchStrategy {
+public abstract class KnnSearchStrategy {
+ public static final int DEFAULT_FILTERED_SEARCH_THRESHOLD = 60;
- /** A strategy for kNN search that uses a set of entry points to start the search */
- class Seeded implements KnnSearchStrategy {
+ /** Override and implement search strategy instance equivalence properly in a subclass. */
+ @Override
+ public abstract boolean equals(Object obj);
+
+ /**
+ * Override and implement search strategy hash code properly in a subclass.
+ *
+ * @see #equals(Object)
+ */
+ @Override
+ public abstract int hashCode();
+
+ /**
+ * A strategy for kNN search that uses HNSW
+ *
+ * @lucene.experimental
+ */
+ public static class Hnsw extends KnnSearchStrategy {
+ public static final Hnsw DEFAULT = new Hnsw(DEFAULT_FILTERED_SEARCH_THRESHOLD);
+
+ private final int filteredSearchThreshold;
+
+ /**
+ * Create a new Hnsw strategy
+ *
+ * @param filteredSearchThreshold threshold for filtered search, a percentage value from 0 to
+ * 100 where 0 means never use filtered search and 100 means always use filtered search.
+ */
+ public Hnsw(int filteredSearchThreshold) {
+ if (filteredSearchThreshold < 0 || filteredSearchThreshold > 100) {
+ throw new IllegalArgumentException("filteredSearchThreshold must be >= 0 and <= 100");
+ }
+ this.filteredSearchThreshold = filteredSearchThreshold;
+ }
+
+ public int filteredSearchThreshold() {
+ return filteredSearchThreshold;
+ }
+
+ /**
+ * Whether to use filtered search based on the ratio of vectors that pass the filter
+ *
+ * @param ratioPassingFilter ratio of vectors that pass the filter
+ * @return true if filtered search should be used
+ */
+ public final boolean useFilteredSearch(float ratioPassingFilter) {
+ assert ratioPassingFilter >= 0 && ratioPassingFilter <= 1;
+ return ratioPassingFilter * 100 < filteredSearchThreshold;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ Hnsw hnsw = (Hnsw) o;
+ return filteredSearchThreshold == hnsw.filteredSearchThreshold;
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hashCode(filteredSearchThreshold);
+ }
+ }
+
+ /**
+ * A strategy for kNN search that uses a set of entry points to start the search
+ *
+ * @lucene.experimental
+ */
+ public static class Seeded extends KnnSearchStrategy {
private final DocIdSetIterator entryPoints;
private final int numberOfEntryPoints;
+ private final KnnSearchStrategy originalStrategy;
- public Seeded(DocIdSetIterator entryPoints, int numberOfEntryPoints) {
+ public Seeded(
+ DocIdSetIterator entryPoints, int numberOfEntryPoints, KnnSearchStrategy originalStrategy) {
if (numberOfEntryPoints < 0) {
throw new IllegalArgumentException("numberOfEntryPoints must be >= 0");
}
@@ -40,6 +112,7 @@ public interface KnnSearchStrategy {
throw new IllegalArgumentException("entryPoints must not be null");
}
this.entryPoints = entryPoints == null ? DocIdSetIterator.empty() : entryPoints;
+ this.originalStrategy = originalStrategy;
}
/**
@@ -59,5 +132,29 @@ public interface KnnSearchStrategy {
public int numberOfEntryPoints() {
return numberOfEntryPoints;
}
+
+ /**
+ * Original strategy to use after seeding
+ *
+ * @return original strategy
+ */
+ public KnnSearchStrategy originalStrategy() {
+ return originalStrategy;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ Seeded seeded = (Seeded) o;
+ return numberOfEntryPoints == seeded.numberOfEntryPoints
+ && Objects.equals(entryPoints, seeded.entryPoints)
+ && Objects.equals(originalStrategy, seeded.originalStrategy);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(entryPoints, numberOfEntryPoints, originalStrategy);
+ }
}
}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/FilteredHnswGraphSearcher.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/FilteredHnswGraphSearcher.java
new file mode 100644
index 00000000000..ae39c612afa
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/FilteredHnswGraphSearcher.java
@@ -0,0 +1,252 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.lucene.util.hnsw;
+
+import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
+
+import java.io.IOException;
+import org.apache.lucene.search.KnnCollector;
+import org.apache.lucene.util.BitSet;
+import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.FixedBitSet;
+import org.apache.lucene.util.SparseFixedBitSet;
+
+/**
+ * Searches an HNSW graph to find nearest neighbors to a query vector. This particular
+ * implementation is optimized for a filtered search, inspired by the ACORN-1 algorithm.
+ * https://arxiv.org/abs/2403.04871 However, this implementation is augmented in some ways, mainly:
+ *
+ * <ul>
+ * <li>It dynamically determines when the optimized filter step should occur based on some
+ * filtered lambda. This is done per small world
+ * <li>The graph searcher doesn't always explore all the extended neighborhood and the number of
+ * additional candidates is predicated on the original candidate's filtered percentage.
+ * </ul>
+ */
+public class FilteredHnswGraphSearcher extends HnswGraphSearcher {
+ // How many filtered candidates must be found to consider N-hop neighbors
+ private static final float EXPANDED_EXPLORATION_LAMBDA = 0.10f;
+
+ // How many extra neighbors to explore, used as a multiple to the candidates neighbor count
+ private final int maxExplorationMultiplier;
+ private final int minToScore;
+
+ /** Creates a new graph searcher. */
+ private FilteredHnswGraphSearcher(
+ NeighborQueue candidates, BitSet visited, int filterSize, HnswGraph graph) {
+ super(candidates, visited);
+ assert graph.maxConn() > 0 : "graph must have known max connections";
+ float filterRatio = filterSize / (float) graph.size();
+ this.maxExplorationMultiplier =
+ (int) Math.round(Math.min(1 / filterRatio, graph.maxConn() / 2.0));
+ // As the filter gets exceptionally restrictive, we must spread out the exploration
+ this.minToScore =
+ (int)
+ Math.round(
+ Math.min(
+ Math.max(0, 1.0 / filterRatio - (2.0 * graph.maxConn())), graph.maxConn()));
+ }
+
+ /**
+ * Creates a new filtered graph searcher.
+ *
+ * @param k the number of nearest neighbors to find
+ * @param graph the graph to search
+ * @param filterSize the number of vectors that pass the accepted ords filter
+ * @param acceptOrds the accepted ords filter
+ * @return a new graph searcher optimized for filtered search
+ */
+ public static FilteredHnswGraphSearcher create(
+ int k, HnswGraph graph, int filterSize, Bits acceptOrds) {
+ if (acceptOrds == null) {
+ throw new IllegalArgumentException("acceptOrds must not be null to used filtered search");
+ }
+ if (filterSize <= 0 || filterSize >= getGraphSize(graph)) {
+ throw new IllegalArgumentException("filterSize must be > 0 and < graph size");
+ }
+ return new FilteredHnswGraphSearcher(
+ new NeighborQueue(k, true), bitSet(filterSize, getGraphSize(graph), k), filterSize, graph);
+ }
+
+ private static BitSet bitSet(long filterSize, int graphSize, int topk) {
+ float percentFiltered = (float) filterSize / graphSize;
+ assert percentFiltered > 0.0f && percentFiltered < 1.0f;
+ double totalOps = Math.log(graphSize) * topk;
+ int approximateVisitation = (int) (totalOps / percentFiltered);
+ return bitSet(approximateVisitation, graphSize);
+ }
+
+ private static BitSet bitSet(int expectedBits, int totalBits) {
+ if (expectedBits < (totalBits >>> 7)) {
+ return new SparseFixedBitSet(totalBits);
+ } else {
+ return new FixedBitSet(totalBits);
+ }
+ }
+
+ @Override
+ void searchLevel(
+ KnnCollector results,
+ RandomVectorScorer scorer,
+ int level,
+ final int[] eps,
+ HnswGraph graph,
+ Bits acceptOrds)
+ throws IOException {
+ assert level == 0 : "Filtered search only works on the base level";
+
+ int size = getGraphSize(graph);
+
+ prepareScratchState();
+
+ for (int ep : eps) {
+ if (visited.getAndSet(ep) == false) {
+ if (results.earlyTerminated()) {
+ return;
+ }
+ float score = scorer.score(ep);
+ results.incVisitedCount(1);
+ candidates.add(ep, score);
+ if (acceptOrds.get(ep)) {
+ results.collect(ep, score);
+ }
+ }
+ }
+ // Collect the vectors to score and potentially add as candidates
+ IntArrayQueue toScore = new IntArrayQueue(graph.maxConn() * 2 * maxExplorationMultiplier);
+ IntArrayQueue toExplore = new IntArrayQueue(graph.maxConn() * 2 * maxExplorationMultiplier);
+ // A bound that holds the minimum similarity to the query vector that a candidate vector must
+ // have to be considered.
+ float minAcceptedSimilarity = Math.nextUp(results.minCompetitiveSimilarity());
+ while (candidates.size() > 0 && results.earlyTerminated() == false) {
+ // get the best candidate (closest or best scoring)
+ float topCandidateSimilarity = candidates.topScore();
+ if (minAcceptedSimilarity > topCandidateSimilarity) {
+ break;
+ }
+ int topCandidateNode = candidates.pop();
+ graph.seek(level, topCandidateNode);
+ int neighborCount = graph.neighborCount();
+ toScore.clear();
+ toExplore.clear();
+ int friendOrd;
+ while ((friendOrd = graph.nextNeighbor()) != NO_MORE_DOCS && toScore.isFull() == false) {
+ assert friendOrd < size : "friendOrd=" + friendOrd + "; size=" + size;
+ if (visited.getAndSet(friendOrd)) {
+ continue;
+ }
+ if (acceptOrds.get(friendOrd)) {
+ toScore.add(friendOrd);
+ } else {
+ toExplore.add(friendOrd);
+ }
+ }
+ // adjust to locally number of filtered candidates to explore
+ float filteredAmount = toExplore.count() / (float) neighborCount;
+ int maxToScoreCount =
+ (int) (neighborCount * Math.min(maxExplorationMultiplier, 1f / (1f - filteredAmount)));
+ int maxAdditionalToExploreCount = toExplore.capacity() - 1;
+ // There is enough filtered, or we don't have enough candidates to score and explore
+ int totalExplored = toScore.count() + toExplore.count();
+ if (toScore.count() < maxToScoreCount && filteredAmount > EXPANDED_EXPLORATION_LAMBDA) {
+ // Now we need to explore the neighbors of the neighbors
+ int exploreFriend;
+ while ((exploreFriend = toExplore.poll()) != NO_MORE_DOCS
+ // only explore initial additional neighborhood
+ && totalExplored < maxAdditionalToExploreCount
+ && toScore.count() < maxToScoreCount) {
+ graphSeek(graph, level, exploreFriend);
+ int friendOfAFriendOrd;
+ while ((friendOfAFriendOrd = graph.nextNeighbor()) != NO_MORE_DOCS
+ && toScore.count() < maxToScoreCount) {
+ if (visited.getAndSet(friendOfAFriendOrd)) {
+ continue;
+ }
+ totalExplored++;
+ if (acceptOrds.get(friendOfAFriendOrd)) {
+ toScore.add(friendOfAFriendOrd);
+ // If we have YET to find a minimum of number candidates, we will continue to explore
+ // until our max
+ } else if (totalExplored < maxAdditionalToExploreCount
+ && toScore.count() < minToScore) {
+ toExplore.add(friendOfAFriendOrd);
+ }
+ }
+ }
+ }
+ // Score the vectors and add them to the candidate list
+ int toScoreOrd;
+ while ((toScoreOrd = toScore.poll()) != NO_MORE_DOCS) {
+ float friendSimilarity = scorer.score(toScoreOrd);
+ results.incVisitedCount(1);
+ if (friendSimilarity > minAcceptedSimilarity) {
+ candidates.add(toScoreOrd, friendSimilarity);
+ if (results.collect(toScoreOrd, friendSimilarity)) {
+ minAcceptedSimilarity = Math.nextUp(results.minCompetitiveSimilarity());
+ }
+ }
+ }
+ }
+ }
+
+ private void prepareScratchState() {
+ candidates.clear();
+ visited.clear();
+ }
+
+ private static class IntArrayQueue {
+ private int[] nodes;
+ private int upto;
+ private int size;
+
+ IntArrayQueue(int capacity) {
+ nodes = new int[capacity];
+ }
+
+ int capacity() {
+ return nodes.length;
+ }
+
+ int count() {
+ return size - upto;
+ }
+
+ void add(int node) {
+ if (isFull()) {
+ throw new UnsupportedOperationException("Initial capacity should remain unchanged");
+ }
+ nodes[size++] = node;
+ }
+
+ boolean isFull() {
+ return size == nodes.length;
+ }
+
+ int poll() {
+ if (upto == size) {
+ return NO_MORE_DOCS;
+ }
+ return nodes[upto++];
+ }
+
+ void clear() {
+ upto = 0;
+ size = 0;
+ }
+ }
+}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraph.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraph.java
index 834b040bf87..2c945382ed5 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraph.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraph.java
@@ -51,6 +51,7 @@ import org.apache.lucene.internal.hppc.IntCursor;
* exclude deleted documents.
*/
public abstract class HnswGraph {
+ public static final int UNKNOWN_MAX_CONN = -1;
/** Sole constructor */
protected HnswGraph() {}
@@ -85,7 +86,7 @@ public abstract class HnswGraph {
public abstract int numLevels() throws IOException;
/** returns M, the maximum number of connections for a node. */
- public abstract int maxConn() throws IOException;
+ public abstract int maxConn();
/** Returns graph's entry point on the top level * */
public abstract int entryNode() throws IOException;
@@ -99,6 +100,8 @@ public abstract class HnswGraph {
*/
public abstract NodesIterator getNodesOnLevel(int level) throws IOException;
+ public abstract int neighborCount();
+
/** Empty graph value */
public static final HnswGraph EMPTY =
new HnswGraph() {
@@ -122,15 +125,20 @@ public abstract class HnswGraph {
}
@Override
- public int maxConn() {
+ public int entryNode() {
return 0;
}
@Override
- public int entryNode() {
+ public int neighborCount() {
return 0;
}
+ @Override
+ public int maxConn() {
+ return UNKNOWN_MAX_CONN;
+ }
+
@Override
public NodesIterator getNodesOnLevel(int level) {
return ArrayNodesIterator.EMPTY;
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphBuilder.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphBuilder.java
index 465d8c8c6eb..882c4edaec5 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphBuilder.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphBuilder.java
@@ -30,6 +30,7 @@ import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.Lock;
import org.apache.lucene.search.KnnCollector;
import org.apache.lucene.search.TopDocs;
+import org.apache.lucene.search.knn.KnnSearchStrategy;
import org.apache.lucene.util.FixedBitSet;
import org.apache.lucene.util.InfoStream;
import org.apache.lucene.util.hnsw.HnswUtil.Component;
@@ -611,5 +612,10 @@ public class HnswGraphBuilder implements HnswBuilder {
public TopDocs topDocs() {
throw new IllegalArgumentException();
}
+
+ @Override
+ public KnnSearchStrategy getSearchStrategy() {
+ throw new IllegalArgumentException("Should not use unique strategy during graph building");
+ }
}
}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphSearcher.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphSearcher.java
index 9278c252541..e0a56fe6dd7 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphSearcher.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphSearcher.java
@@ -37,9 +37,9 @@ public class HnswGraphSearcher extends AbstractHnswGraphSearcher {
* Scratch data structures that are used in each {@link #searchLevel} call. These can be expensive
* to allocate, so they're cleared and reused across calls.
*/
- private final NeighborQueue candidates;
+ protected final NeighborQueue candidates;
- private BitSet visited;
+ protected BitSet visited;
/**
* Creates a new graph searcher.
@@ -52,6 +52,28 @@ public class HnswGraphSearcher extends AbstractHnswGraphSearcher {
this.visited = visited;
}
+ /**
+ * See {@link HnswGraphSearcher#search(RandomVectorScorer, KnnCollector, HnswGraph, Bits, int)}
+ *
+ * @param scorer the scorer to compare the query with the nodes
+ * @param knnCollector a hnsw knn collector of top knn results to be returned
+ * @param graph the graph values. May represent the entire graph, or a level in a hierarchical
+ * graph.
+ * @param acceptOrds {@link Bits} that represents the allowed document ordinals to match, or
+ * {@code null} if they are all allowed to match.
+ */
+ public static void search(
+ RandomVectorScorer scorer, KnnCollector knnCollector, HnswGraph graph, Bits acceptOrds)
+ throws IOException {
+ int filteredDocCount = 0;
+ if (acceptOrds instanceof BitSet bitSet) {
+ // Use approximate cardinality as this is good enough, but ensure we don't exceed the graph
+ // size as that is illogical
+ filteredDocCount = Math.min(bitSet.approximateCardinality(), graph.size());
+ }
+ search(scorer, knnCollector, graph, acceptOrds, filteredDocCount);
+ }
+
/**
* Searches the HNSW graph for the nearest neighbors of a query vector. If entry points are
* directly provided via the knnCollector, then the search will be initialized at those points.
@@ -63,18 +85,49 @@ public class HnswGraphSearcher extends AbstractHnswGraphSearcher {
* graph.
* @param acceptOrds {@link Bits} that represents the allowed document ordinals to match, or
* {@code null} if they are all allowed to match.
+ * @param filteredDocCount the number of docs that pass the filter
*/
public static void search(
- RandomVectorScorer scorer, KnnCollector knnCollector, HnswGraph graph, Bits acceptOrds)
+ RandomVectorScorer scorer,
+ KnnCollector knnCollector,
+ HnswGraph graph,
+ Bits acceptOrds,
+ int filteredDocCount)
throws IOException {
- AbstractHnswGraphSearcher graphSearcher =
- new HnswGraphSearcher(
- new NeighborQueue(knnCollector.k(), true), new SparseFixedBitSet(getGraphSize(graph)));
+ assert filteredDocCount >= 0 && filteredDocCount <= graph.size();
+ KnnSearchStrategy.Hnsw hnswStrategy;
+ if (knnCollector.getSearchStrategy() instanceof KnnSearchStrategy.Hnsw hnsw) {
+ hnswStrategy = hnsw;
+ } else if (knnCollector.getSearchStrategy() instanceof KnnSearchStrategy.Seeded seeded
+ && seeded.originalStrategy() instanceof KnnSearchStrategy.Hnsw hnsw) {
+ hnswStrategy = hnsw;
+ } else {
+ hnswStrategy = KnnSearchStrategy.Hnsw.DEFAULT;
+ }
+ final AbstractHnswGraphSearcher innerSearcher;
+ // First, check if we should use a filtered searcher
+ if (acceptOrds != null
+ // We can only use filtered search if we know the maxConn
+ && graph.maxConn() != HnswGraph.UNKNOWN_MAX_CONN
+ && filteredDocCount > 0
+ && hnswStrategy.useFilteredSearch((float) filteredDocCount / graph.size())) {
+ innerSearcher =
+ FilteredHnswGraphSearcher.create(knnCollector.k(), graph, filteredDocCount, acceptOrds);
+ } else {
+ innerSearcher =
+ new HnswGraphSearcher(
+ new NeighborQueue(knnCollector.k(), true),
+ new SparseFixedBitSet(getGraphSize(graph)));
+ }
+ // Then, check if we the search strategy is seeded
+ final AbstractHnswGraphSearcher graphSearcher;
if (knnCollector.getSearchStrategy() instanceof KnnSearchStrategy.Seeded seeded
&& seeded.numberOfEntryPoints() > 0) {
graphSearcher =
SeededHnswGraphSearcher.fromEntryPoints(
- graphSearcher, seeded.numberOfEntryPoints(), seeded.entryPoints(), graph.size());
+ innerSearcher, seeded.numberOfEntryPoints(), seeded.entryPoints(), graph.size());
+ } else {
+ graphSearcher = innerSearcher;
}
graphSearcher.search(knnCollector, scorer, graph, acceptOrds);
}
@@ -276,7 +329,7 @@ public class HnswGraphSearcher extends AbstractHnswGraphSearcher {
return graph.nextNeighbor();
}
- private static int getGraphSize(HnswGraph graph) {
+ static int getGraphSize(HnswGraph graph) {
return graph.maxNodeId() + 1;
}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/OnHeapHnswGraph.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/OnHeapHnswGraph.java
index 61f30152787..9b55369d8e9 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/OnHeapHnswGraph.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/OnHeapHnswGraph.java
@@ -170,6 +170,11 @@ public final class OnHeapHnswGraph extends HnswGraph implements Accountable {
upto = -1;
}
+ @Override
+ public int neighborCount() {
+ return cur.size();
+ }
+
@Override
public int nextNeighbor() {
if (++upto < cur.size()) {
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 842426a1611..e2690d8d3bf 100644
--- a/lucene/core/src/test/org/apache/lucene/search/BaseKnnVectorQueryTestCase.java
+++ b/lucene/core/src/test/org/apache/lucene/search/BaseKnnVectorQueryTestCase.java
@@ -620,20 +620,6 @@ abstract class BaseKnnVectorQueryTestCase extends LuceneTestCase {
getThrowingKnnVectorQuery("field", randomVector(dimension), 10, filter1),
numDocs));
- // Test a restrictive filter and check we use exact search
- Query filter2 = IntPoint.newRangeQuery("tag", lower, lower + 6);
- results =
- searcher.search(
- getKnnVectorQuery("field", randomVector(dimension), 5, filter2), numDocs);
- assertEquals(5, results.totalHits.value());
- assertEquals(results.totalHits.value(), results.scoreDocs.length);
- expectThrows(
- UnsupportedOperationException.class,
- () ->
- searcher.search(
- getThrowingKnnVectorQuery("field", randomVector(dimension), 5, filter2),
- numDocs));
-
// Test an unrestrictive filter and check we use approximate search
Query filter3 = IntPoint.newRangeQuery("tag", lower, numDocs);
results =
@@ -653,7 +639,7 @@ abstract class BaseKnnVectorQueryTestCase extends LuceneTestCase {
}
// Test a filter that exhausts visitedLimit in upper levels, and switches to exact search
- Query filter4 = IntPoint.newRangeQuery("tag", lower, lower + 2);
+ Query filter4 = IntPoint.newRangeQuery("tag", lower, lower + 6);
expectThrows(
UnsupportedOperationException.class,
() ->
diff --git a/lucene/core/src/test/org/apache/lucene/search/TestKnnByteVectorQuery.java b/lucene/core/src/test/org/apache/lucene/search/TestKnnByteVectorQuery.java
index 21219e0e1d9..dd24b14e685 100644
--- a/lucene/core/src/test/org/apache/lucene/search/TestKnnByteVectorQuery.java
+++ b/lucene/core/src/test/org/apache/lucene/search/TestKnnByteVectorQuery.java
@@ -25,6 +25,7 @@ import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.QueryTimeout;
import org.apache.lucene.index.Term;
import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.search.knn.KnnSearchStrategy;
import org.apache.lucene.store.Directory;
import org.apache.lucene.util.TestVectorUtil;
@@ -112,7 +113,7 @@ public class TestKnnByteVectorQuery extends BaseKnnVectorQueryTestCase {
static class ThrowingKnnVectorQuery extends KnnByteVectorQuery {
public ThrowingKnnVectorQuery(String field, byte[] target, int k, Query filter) {
- super(field, target, k, filter);
+ super(field, target, k, filter, new KnnSearchStrategy.Hnsw(0));
}
@Override
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 ece2b385654..750f64d7c2b 100644
--- a/lucene/core/src/test/org/apache/lucene/search/TestKnnFloatVectorQuery.java
+++ b/lucene/core/src/test/org/apache/lucene/search/TestKnnFloatVectorQuery.java
@@ -37,6 +37,7 @@ import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.QueryTimeout;
import org.apache.lucene.index.Term;
import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.search.knn.KnnSearchStrategy;
import org.apache.lucene.store.Directory;
import org.apache.lucene.tests.index.RandomIndexWriter;
import org.apache.lucene.tests.util.LuceneTestCase;
@@ -262,7 +263,7 @@ public class TestKnnFloatVectorQuery extends BaseKnnVectorQueryTestCase {
static class ThrowingKnnVectorQuery extends KnnFloatVectorQuery {
public ThrowingKnnVectorQuery(String field, float[] target, int k, Query filter) {
- super(field, target, k, filter);
+ super(field, target, k, filter, new KnnSearchStrategy.Hnsw(0));
}
@Override
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 49ed29fcc08..3d1b1e5b37d 100644
--- a/lucene/core/src/test/org/apache/lucene/search/TestSeededKnnByteVectorQuery.java
+++ b/lucene/core/src/test/org/apache/lucene/search/TestSeededKnnByteVectorQuery.java
@@ -154,7 +154,7 @@ public class TestSeededKnnByteVectorQuery extends BaseKnnVectorQueryTestCase {
} else {
searcher.setTimeout(null);
}
- int k = random().nextInt(80) + 1;
+ int k = random().nextInt(10) + 1;
int n = random().nextInt(100) + 1;
// we may get fewer results than requested if there are deletions, but this test doesn't
// check that
@@ -301,7 +301,7 @@ public class TestSeededKnnByteVectorQuery extends BaseKnnVectorQueryTestCase {
private final KnnSearchStrategy.Seeded seeded;
public AssertingSeededStrategy(KnnSearchStrategy.Seeded seeded) {
- super(seeded.entryPoints(), seeded.numberOfEntryPoints());
+ super(seeded.entryPoints(), seeded.numberOfEntryPoints(), seeded.originalStrategy());
this.seeded = seeded;
}
diff --git a/lucene/core/src/test/org/apache/lucene/search/TestSeededKnnFloatVectorQuery.java b/lucene/core/src/test/org/apache/lucene/search/TestSeededKnnFloatVectorQuery.java
index c15b720b3d1..5280f9489e3 100644
--- a/lucene/core/src/test/org/apache/lucene/search/TestSeededKnnFloatVectorQuery.java
+++ b/lucene/core/src/test/org/apache/lucene/search/TestSeededKnnFloatVectorQuery.java
@@ -140,7 +140,7 @@ public class TestSeededKnnFloatVectorQuery extends BaseKnnVectorQueryTestCase {
} else {
searcher.setTimeout(null);
}
- int k = random().nextInt(80) + 1;
+ int k = random().nextInt(10) + 1;
int n = random().nextInt(100) + 1;
// we may get fewer results than requested if there are deletions, but this test doesn't
// check that
diff --git a/lucene/core/src/test/org/apache/lucene/util/hnsw/HnswGraphTestCase.java b/lucene/core/src/test/org/apache/lucene/util/hnsw/HnswGraphTestCase.java
index c6da6d30faf..91fb4779551 100644
--- a/lucene/core/src/test/org/apache/lucene/util/hnsw/HnswGraphTestCase.java
+++ b/lucene/core/src/test/org/apache/lucene/util/hnsw/HnswGraphTestCase.java
@@ -1336,7 +1336,12 @@ abstract class HnswGraphTestCase<T> extends LuceneTestCase {
}
@Override
- public int maxConn() throws IOException {
+ public int neighborCount() {
+ return delegate.neighborCount();
+ }
+
+ @Override
+ public int maxConn() {
return delegate.maxConn();
}
diff --git a/lucene/core/src/test/org/apache/lucene/util/hnsw/TestHnswUtil.java b/lucene/core/src/test/org/apache/lucene/util/hnsw/TestHnswUtil.java
index 7ea05150e14..84637d86ca7 100644
--- a/lucene/core/src/test/org/apache/lucene/util/hnsw/TestHnswUtil.java
+++ b/lucene/core/src/test/org/apache/lucene/util/hnsw/TestHnswUtil.java
@@ -250,11 +250,6 @@ public class TestHnswUtil extends LuceneTestCase {
return 0;
}
- @Override
- public int maxConn() {
- return 0;
- }
-
@Override
public String toString() {
StringBuilder buf = new StringBuilder();
@@ -273,6 +268,16 @@ public class TestHnswUtil extends LuceneTestCase {
return buf.toString();
}
+ @Override
+ public int neighborCount() {
+ return nodes[currentLevel][currentNode].length;
+ }
+
+ @Override
+ public int maxConn() {
+ return UNKNOWN_MAX_CONN;
+ }
+
@Override
public NodesIterator getNodesOnLevel(int level) {
diff --git a/lucene/join/src/java/org/apache/lucene/search/join/DiversifyingChildrenByteKnnVectorQuery.java b/lucene/join/src/java/org/apache/lucene/search/join/DiversifyingChildrenByteKnnVectorQuery.java
index a0b0f5d32ac..d76b2bed75a 100644
--- a/lucene/join/src/java/org/apache/lucene/search/join/DiversifyingChildrenByteKnnVectorQuery.java
+++ b/lucene/join/src/java/org/apache/lucene/search/join/DiversifyingChildrenByteKnnVectorQuery.java
@@ -16,6 +16,8 @@
*/
package org.apache.lucene.search.join;
+import static org.apache.lucene.search.knn.KnnSearchStrategy.Hnsw.DEFAULT;
+
import java.io.IOException;
import java.util.Arrays;
import java.util.Objects;
@@ -34,6 +36,7 @@ import org.apache.lucene.search.TopDocsCollector;
import org.apache.lucene.search.TotalHits;
import org.apache.lucene.search.VectorScorer;
import org.apache.lucene.search.knn.KnnCollectorManager;
+import org.apache.lucene.search.knn.KnnSearchStrategy;
import org.apache.lucene.util.BitSet;
import org.apache.lucene.util.Bits;
@@ -69,7 +72,30 @@ public class DiversifyingChildrenByteKnnVectorQuery extends KnnByteVectorQuery {
*/
public DiversifyingChildrenByteKnnVectorQuery(
String field, byte[] query, Query childFilter, int k, BitSetProducer parentsFilter) {
- super(field, query, k, childFilter);
+ this(field, query, childFilter, k, parentsFilter, DEFAULT);
+ }
+
+ /**
+ * Create a DiversifyingChildrenByteKnnVectorQuery.
+ *
+ * @param field the query field
+ * @param query the vector query
+ * @param childFilter the child filter
+ * @param k how many parent documents to return given the matching children
+ * @param parentsFilter Filter identifying the parent documents.
+ * @param searchStrategy the search strategy to use. If null, the default strategy will be used.
+ * The underlying format may not support all strategies and is free to ignore the requested
+ * strategy.
+ * @lucene.experimental
+ */
+ public DiversifyingChildrenByteKnnVectorQuery(
+ String field,
+ byte[] query,
+ Query childFilter,
+ int k,
+ BitSetProducer parentsFilter,
+ KnnSearchStrategy searchStrategy) {
+ super(field, query, k, childFilter, searchStrategy);
this.childFilter = childFilter;
this.parentsFilter = parentsFilter;
this.k = k;
@@ -144,7 +170,8 @@ public class DiversifyingChildrenByteKnnVectorQuery extends KnnByteVectorQuery {
KnnCollectorManager knnCollectorManager)
throws IOException {
ByteVectorValues.checkField(context.reader(), field);
- KnnCollector collector = knnCollectorManager.newCollector(visitedLimit, null, context);
+ KnnCollector collector =
+ knnCollectorManager.newCollector(visitedLimit, searchStrategy, context);
if (collector == null) {
return NO_RESULTS;
}
diff --git a/lucene/join/src/java/org/apache/lucene/search/join/DiversifyingChildrenFloatKnnVectorQuery.java b/lucene/join/src/java/org/apache/lucene/search/join/DiversifyingChildrenFloatKnnVectorQuery.java
index 3b05b392654..3cffa44772f 100644
--- a/lucene/join/src/java/org/apache/lucene/search/join/DiversifyingChildrenFloatKnnVectorQuery.java
+++ b/lucene/join/src/java/org/apache/lucene/search/join/DiversifyingChildrenFloatKnnVectorQuery.java
@@ -16,6 +16,8 @@
*/
package org.apache.lucene.search.join;
+import static org.apache.lucene.search.knn.KnnSearchStrategy.Hnsw.DEFAULT;
+
import java.io.IOException;
import java.util.Arrays;
import java.util.Objects;
@@ -34,6 +36,7 @@ import org.apache.lucene.search.TopDocsCollector;
import org.apache.lucene.search.TotalHits;
import org.apache.lucene.search.VectorScorer;
import org.apache.lucene.search.knn.KnnCollectorManager;
+import org.apache.lucene.search.knn.KnnSearchStrategy;
import org.apache.lucene.util.BitSet;
import org.apache.lucene.util.Bits;
@@ -59,7 +62,7 @@ public class DiversifyingChildrenFloatKnnVectorQuery extends KnnFloatVectorQuery
private final float[] query;
/**
- * Create a ToParentBlockJoinFloatVectorQuery.
+ * Create a DiversifyingChildrenFloatKnnVectorQuery.
*
* @param field the query field
* @param query the vector query
@@ -69,7 +72,30 @@ public class DiversifyingChildrenFloatKnnVectorQuery extends KnnFloatVectorQuery
*/
public DiversifyingChildrenFloatKnnVectorQuery(
String field, float[] query, Query childFilter, int k, BitSetProducer parentsFilter) {
- super(field, query, k, childFilter);
+ this(field, query, childFilter, k, parentsFilter, DEFAULT);
+ }
+
+ /**
+ * Create a DiversifyingChildrenFloatKnnVectorQuery.
+ *
+ * @param field the query field
+ * @param query the vector query
+ * @param childFilter the child filter
+ * @param k how many parent documents to return given the matching children
+ * @param parentsFilter Filter identifying the parent documents.
+ * @param searchStrategy the search strategy to use. If null, the default strategy will be used.
+ * The underlying format may not support all strategies and is free to ignore the requested
+ * strategy.
+ * @lucene.experimental
+ */
+ public DiversifyingChildrenFloatKnnVectorQuery(
+ String field,
+ float[] query,
+ Query childFilter,
+ int k,
+ BitSetProducer parentsFilter,
+ KnnSearchStrategy searchStrategy) {
+ super(field, query, k, childFilter, searchStrategy);
this.childFilter = childFilter;
this.parentsFilter = parentsFilter;
this.k = k;
@@ -143,7 +169,8 @@ public class DiversifyingChildrenFloatKnnVectorQuery extends KnnFloatVectorQuery
KnnCollectorManager knnCollectorManager)
throws IOException {
FloatVectorValues.checkField(context.reader(), field);
- KnnCollector collector = knnCollectorManager.newCollector(visitedLimit, null, context);
+ KnnCollector collector =
+ knnCollectorManager.newCollector(visitedLimit, searchStrategy, context);
if (collector == null) {
return NO_RESULTS;
}