You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@lucene.apache.org by to...@apache.org on 2025/04/02 14:20:18 UTC
(lucene) branch main updated: Add a HNSW collector that exits early when nearest neighbor queue saturates (#14094)
This is an automated email from the ASF dual-hosted git repository.
tommaso 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 525bf34bfdf Add a HNSW collector that exits early when nearest neighbor queue saturates (#14094)
525bf34bfdf is described below
commit 525bf34bfdfc16cc220d326d2cf30541f1afef29
Author: Tommaso Teofili <to...@elastic.co>
AuthorDate: Wed Apr 2 16:20:13 2025 +0200
Add a HNSW collector that exits early when nearest neighbor queue saturates (#14094)
* Add a HNSW early termination based on nearest neighbor queue saturation
Co-authored-by: Benjamin Trent <be...@gmail.com>
---
lucene/CHANGES.txt | 2 +
.../search/HnswQueueSaturationCollector.java | 101 +++++++++
.../lucene/search/PatienceKnnVectorQuery.java | 237 +++++++++++++++++++++
.../lucene/search/knn/KnnSearchStrategy.java | 40 ++++
.../util/hnsw/FilteredHnswGraphSearcher.java | 3 +
.../apache/lucene/util/hnsw/HnswGraphBuilder.java | 2 +-
.../apache/lucene/util/hnsw/HnswGraphSearcher.java | 3 +
.../search/TestHnswQueueSaturationCollector.java | 108 ++++++++++
.../lucene/search/TestPatienceByteVectorQuery.java | 114 ++++++++++
.../search/TestPatienceFloatVectorQuery.java | 105 +++++++++
10 files changed, 714 insertions(+), 1 deletion(-)
diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index fd3f1596b06..f80e14e5e7b 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -143,6 +143,8 @@ New Features
* GITHUB#14412: Allow skip cache factor to be updated dynamically. (Sagar Upadhyaya)
+* GITHUB#14094: New KNN query that early terminates when HNSW nearest neighbor queue saturates. (Tommaso Teofili)
+
Improvements
---------------------
diff --git a/lucene/core/src/java/org/apache/lucene/search/HnswQueueSaturationCollector.java b/lucene/core/src/java/org/apache/lucene/search/HnswQueueSaturationCollector.java
new file mode 100644
index 00000000000..6acc85fcee6
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/search/HnswQueueSaturationCollector.java
@@ -0,0 +1,101 @@
+/*
+ * 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.search;
+
+import org.apache.lucene.search.knn.KnnSearchStrategy;
+
+/**
+ * A {@link KnnCollector.Decorator} that early exits when nearest neighbor queue keeps saturating
+ * beyond a 'patience' parameter. This records the rate of collection of new nearest neighbors in
+ * the {@code delegate} KnnCollector queue, at each HNSW node candidate visit. Once it saturates for
+ * a number of consecutive node visits (e.g., the patience parameter), this early terminates.
+ *
+ * @lucene.experimental
+ */
+public class HnswQueueSaturationCollector extends KnnCollector.Decorator {
+
+ private final KnnCollector delegate;
+ private final double saturationThreshold;
+ private final int patience;
+ private boolean patienceFinished;
+ private int countSaturated;
+ private int previousQueueSize;
+ private int currentQueueSize;
+
+ HnswQueueSaturationCollector(KnnCollector delegate, double saturationThreshold, int patience) {
+ super(delegate);
+ this.delegate = delegate;
+ this.previousQueueSize = 0;
+ this.currentQueueSize = 0;
+ this.countSaturated = 0;
+ this.patienceFinished = false;
+ this.saturationThreshold = saturationThreshold;
+ this.patience = patience;
+ }
+
+ @Override
+ public boolean earlyTerminated() {
+ return delegate.earlyTerminated() || patienceFinished;
+ }
+
+ @Override
+ public boolean collect(int docId, float similarity) {
+ boolean collect = delegate.collect(docId, similarity);
+ if (collect) {
+ currentQueueSize++;
+ }
+ return collect;
+ }
+
+ @Override
+ public TopDocs topDocs() {
+ TopDocs topDocs;
+ if (patienceFinished && delegate.earlyTerminated() == false) {
+ // this avoids re-running exact search in the filtered scenario when patience is exhausted
+ TopDocs delegateDocs = delegate.topDocs();
+ TotalHits totalHits =
+ new TotalHits(delegateDocs.totalHits.value(), TotalHits.Relation.EQUAL_TO);
+ topDocs = new TopDocs(totalHits, delegateDocs.scoreDocs);
+ } else {
+ topDocs = delegate.topDocs();
+ }
+ return topDocs;
+ }
+
+ public void nextCandidate() {
+ double queueSaturation =
+ (double) Math.min(currentQueueSize, previousQueueSize) / currentQueueSize;
+ previousQueueSize = currentQueueSize;
+ if (queueSaturation >= saturationThreshold) {
+ countSaturated++;
+ } else {
+ countSaturated = 0;
+ }
+ if (countSaturated > patience) {
+ patienceFinished = true;
+ }
+ }
+
+ @Override
+ public KnnSearchStrategy getSearchStrategy() {
+ KnnSearchStrategy delegateStrategy = delegate.getSearchStrategy();
+ assert delegateStrategy instanceof KnnSearchStrategy.Hnsw;
+ return new KnnSearchStrategy.Patience(
+ this, ((KnnSearchStrategy.Hnsw) delegateStrategy).filteredSearchThreshold());
+ }
+}
diff --git a/lucene/core/src/java/org/apache/lucene/search/PatienceKnnVectorQuery.java b/lucene/core/src/java/org/apache/lucene/search/PatienceKnnVectorQuery.java
new file mode 100644
index 00000000000..ac2d401cbbe
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/search/PatienceKnnVectorQuery.java
@@ -0,0 +1,237 @@
+/*
+ * 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.search;
+
+import java.io.IOException;
+import java.util.Objects;
+import org.apache.lucene.index.FieldInfo;
+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.Bits;
+
+/**
+ * This is a version of knn vector query that exits early when HNSW queue saturates over a {@code
+ * #saturationThreshold} for more than {@code #patience} times.
+ *
+ * <p>See <a
+ * href="https://cs.uwaterloo.ca/~jimmylin/publications/Teofili_Lin_ECIR2025.pdf">"Patience in
+ * Proximity: A Simple Early Termination Strategy for HNSW Graph Traversal in Approximate k-Nearest
+ * Neighbor Search"</a> (Teofili and Lin). In ECIR '25: Proceedings of the 47th European Conference
+ * on Information Retrieval.
+ *
+ * @lucene.experimental
+ */
+public class PatienceKnnVectorQuery extends AbstractKnnVectorQuery {
+
+ private static final double DEFAULT_SATURATION_THRESHOLD = 0.995d;
+
+ private final int patience;
+ private final double saturationThreshold;
+
+ final AbstractKnnVectorQuery delegate;
+
+ /**
+ * Construct a new PatienceKnnVectorQuery instance for a float vector field
+ *
+ * @param knnQuery the knn query to be seeded
+ * @param saturationThreshold the early exit saturation threshold
+ * @param patience the patience parameter
+ * @return a new PatienceKnnVectorQuery instance
+ * @lucene.experimental
+ */
+ public static PatienceKnnVectorQuery fromFloatQuery(
+ KnnFloatVectorQuery knnQuery, double saturationThreshold, int patience) {
+ return new PatienceKnnVectorQuery(knnQuery, saturationThreshold, patience);
+ }
+
+ /**
+ * Construct a new PatienceKnnVectorQuery instance for a float vector field
+ *
+ * @param knnQuery the knn query to be seeded
+ * @return a new PatienceKnnVectorQuery instance
+ * @lucene.experimental
+ */
+ public static PatienceKnnVectorQuery fromFloatQuery(KnnFloatVectorQuery knnQuery) {
+ return new PatienceKnnVectorQuery(
+ knnQuery, DEFAULT_SATURATION_THRESHOLD, defaultPatience(knnQuery));
+ }
+
+ /**
+ * Construct a new PatienceKnnVectorQuery instance for a byte vector field
+ *
+ * @param knnQuery the knn query to be seeded
+ * @param saturationThreshold the early exit saturation threshold
+ * @param patience the patience parameter
+ * @return a new PatienceKnnVectorQuery instance
+ * @lucene.experimental
+ */
+ public static PatienceKnnVectorQuery fromByteQuery(
+ KnnByteVectorQuery knnQuery, double saturationThreshold, int patience) {
+ return new PatienceKnnVectorQuery(knnQuery, saturationThreshold, patience);
+ }
+
+ /**
+ * Construct a new PatienceKnnVectorQuery instance for a byte vector field
+ *
+ * @param knnQuery the knn query to be seeded
+ * @return a new PatienceKnnVectorQuery instance
+ * @lucene.experimental
+ */
+ public static PatienceKnnVectorQuery fromByteQuery(KnnByteVectorQuery knnQuery) {
+ return new PatienceKnnVectorQuery(
+ knnQuery, DEFAULT_SATURATION_THRESHOLD, defaultPatience(knnQuery));
+ }
+
+ /**
+ * Construct a new PatienceKnnVectorQuery instance for seeded vector field
+ *
+ * @param knnQuery the knn query to be seeded
+ * @param saturationThreshold the early exit saturation threshold
+ * @param patience the patience parameter
+ * @return a new PatienceKnnVectorQuery instance
+ * @lucene.experimental
+ */
+ public static PatienceKnnVectorQuery fromSeededQuery(
+ SeededKnnVectorQuery knnQuery, double saturationThreshold, int patience) {
+ return new PatienceKnnVectorQuery(knnQuery, saturationThreshold, patience);
+ }
+
+ /**
+ * Construct a new PatienceKnnVectorQuery instance for seeded vector field
+ *
+ * @param knnQuery the knn query to be seeded
+ * @return a new PatienceKnnVectorQuery instance
+ * @lucene.experimental
+ */
+ public static PatienceKnnVectorQuery fromSeededQuery(SeededKnnVectorQuery knnQuery) {
+ return new PatienceKnnVectorQuery(
+ knnQuery, DEFAULT_SATURATION_THRESHOLD, defaultPatience(knnQuery));
+ }
+
+ PatienceKnnVectorQuery(
+ AbstractKnnVectorQuery knnQuery, double saturationThreshold, int patience) {
+ super(knnQuery.field, knnQuery.k, knnQuery.filter, knnQuery.searchStrategy);
+ this.delegate = knnQuery;
+ this.saturationThreshold = saturationThreshold;
+ this.patience = patience;
+ }
+
+ private static int defaultPatience(AbstractKnnVectorQuery delegate) {
+ return Math.max(7, (int) (delegate.k * 0.3));
+ }
+
+ @Override
+ public String toString(String field) {
+ return "PatienceKnnVectorQuery{"
+ + "saturationThreshold="
+ + saturationThreshold
+ + ", patience="
+ + patience
+ + ", delegate="
+ + delegate
+ + '}';
+ }
+
+ @Override
+ protected KnnCollectorManager getKnnCollectorManager(int k, IndexSearcher searcher) {
+ return delegate.getKnnCollectorManager(k, searcher);
+ }
+
+ @Override
+ protected TopDocs approximateSearch(
+ LeafReaderContext context,
+ Bits acceptDocs,
+ int visitedLimit,
+ KnnCollectorManager knnCollectorManager)
+ throws IOException {
+ return delegate.approximateSearch(
+ context, acceptDocs, visitedLimit, new PatienceCollectorManager(knnCollectorManager));
+ }
+
+ @Override
+ protected TopDocs exactSearch(
+ LeafReaderContext context, DocIdSetIterator acceptIterator, QueryTimeout queryTimeout)
+ throws IOException {
+ return delegate.exactSearch(context, acceptIterator, queryTimeout);
+ }
+
+ @Override
+ protected TopDocs mergeLeafResults(TopDocs[] perLeafResults) {
+ return delegate.mergeLeafResults(perLeafResults);
+ }
+
+ @Override
+ public void visit(QueryVisitor visitor) {
+ delegate.visit(visitor);
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ if (!super.equals(o)) return false;
+ PatienceKnnVectorQuery that = (PatienceKnnVectorQuery) o;
+ return saturationThreshold == that.saturationThreshold
+ && patience == that.patience
+ && Objects.equals(delegate, that.delegate);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(super.hashCode(), saturationThreshold, patience, delegate);
+ }
+
+ @Override
+ public String getField() {
+ return delegate.getField();
+ }
+
+ @Override
+ public int getK() {
+ return delegate.getK();
+ }
+
+ @Override
+ public Query getFilter() {
+ return delegate.getFilter();
+ }
+
+ @Override
+ VectorScorer createVectorScorer(LeafReaderContext context, FieldInfo fi) throws IOException {
+ return delegate.createVectorScorer(context, fi);
+ }
+
+ class PatienceCollectorManager implements KnnCollectorManager {
+ final KnnCollectorManager knnCollectorManager;
+
+ PatienceCollectorManager(KnnCollectorManager knnCollectorManager) {
+ this.knnCollectorManager = knnCollectorManager;
+ }
+
+ @Override
+ public KnnCollector newCollector(
+ int visitLimit, KnnSearchStrategy searchStrategy, LeafReaderContext ctx)
+ throws IOException {
+ return new HnswQueueSaturationCollector(
+ knnCollectorManager.newCollector(visitLimit, searchStrategy, ctx),
+ saturationThreshold,
+ patience);
+ }
+ }
+}
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 bcb938d5c65..50b88591c53 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
@@ -18,6 +18,7 @@ package org.apache.lucene.search.knn;
import java.util.Objects;
import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.search.HnswQueueSaturationCollector;
/**
* KnnSearchStrategy is a strategy for kNN search, providing additional search strategy
@@ -40,6 +41,9 @@ public abstract class KnnSearchStrategy {
@Override
public abstract int hashCode();
+ /** Signal processing of the next block of vectors. */
+ public abstract void nextVectorsBlock();
+
/**
* A strategy for kNN search that uses HNSW
*
@@ -90,6 +94,9 @@ public abstract class KnnSearchStrategy {
public int hashCode() {
return Objects.hashCode(filteredSearchThreshold);
}
+
+ @Override
+ public void nextVectorsBlock() {}
}
/**
@@ -156,5 +163,38 @@ public abstract class KnnSearchStrategy {
public int hashCode() {
return Objects.hash(entryPoints, numberOfEntryPoints, originalStrategy);
}
+
+ @Override
+ public void nextVectorsBlock() {}
+ }
+
+ /**
+ * A strategy for kNN search on HNSW that early exits when nearest neighbor collection rate
+ * saturates.
+ *
+ * @lucene.experimental
+ */
+ public static class Patience extends Hnsw {
+ private final HnswQueueSaturationCollector collector;
+
+ public Patience(HnswQueueSaturationCollector collector, int filteredSearchThreshold) {
+ super(filteredSearchThreshold);
+ this.collector = collector;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ return super.equals(obj) && Objects.equals(collector, ((Patience) obj).collector);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(super.filteredSearchThreshold, collector);
+ }
+
+ @Override
+ public void nextVectorsBlock() {
+ collector.nextCandidate();
+ }
}
}
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
index ae39c612afa..10e5725d81e 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/FilteredHnswGraphSearcher.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/FilteredHnswGraphSearcher.java
@@ -201,6 +201,9 @@ public class FilteredHnswGraphSearcher extends HnswGraphSearcher {
}
}
}
+ if (results.getSearchStrategy() != null) {
+ results.getSearchStrategy().nextVectorsBlock();
+ }
}
}
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 5106c981a9b..56f7bc62e44 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
@@ -634,7 +634,7 @@ public class HnswGraphBuilder implements HnswBuilder {
@Override
public KnnSearchStrategy getSearchStrategy() {
- throw new IllegalArgumentException("Should not use unique strategy during graph building");
+ return null;
}
}
}
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 cee873446dc..6f16ae70b53 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
@@ -312,6 +312,9 @@ public class HnswGraphSearcher extends AbstractHnswGraphSearcher {
}
}
}
+ if (results.getSearchStrategy() != null) {
+ results.getSearchStrategy().nextVectorsBlock();
+ }
}
}
diff --git a/lucene/core/src/test/org/apache/lucene/search/TestHnswQueueSaturationCollector.java b/lucene/core/src/test/org/apache/lucene/search/TestHnswQueueSaturationCollector.java
new file mode 100644
index 00000000000..9f6a98f3812
--- /dev/null
+++ b/lucene/core/src/test/org/apache/lucene/search/TestHnswQueueSaturationCollector.java
@@ -0,0 +1,108 @@
+/*
+ * 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.search;
+
+import java.util.Random;
+import org.apache.lucene.tests.util.LuceneTestCase;
+import org.junit.Test;
+
+/** Tests for {@link HnswQueueSaturationCollector} */
+public class TestHnswQueueSaturationCollector extends LuceneTestCase {
+
+ @Test
+ public void testDelegate() {
+ Random random = random();
+ int numDocs = 100;
+ int k = random.nextInt(1, 10);
+ KnnCollector delegate = new TopKnnCollector(k, numDocs);
+ HnswQueueSaturationCollector queueSaturationCollector =
+ new HnswQueueSaturationCollector(delegate, 1, Integer.MAX_VALUE);
+ for (int i = 0; i < random.nextInt(numDocs); i++) {
+ queueSaturationCollector.collect(random.nextInt(numDocs), random.nextFloat(1.0f));
+ }
+ assertEquals(delegate.k(), queueSaturationCollector.k());
+ assertEquals(delegate.visitedCount(), queueSaturationCollector.visitedCount());
+ assertEquals(delegate.visitLimit(), queueSaturationCollector.visitLimit());
+ assertEquals(
+ delegate.minCompetitiveSimilarity(),
+ queueSaturationCollector.minCompetitiveSimilarity(),
+ 1e-3);
+ }
+
+ @Test
+ public void testEarlyExpectedExit() {
+ int numDocs = 1000;
+ int k = 10;
+ KnnCollector delegate = new TopKnnCollector(k, numDocs);
+ HnswQueueSaturationCollector queueSaturationCollector =
+ new HnswQueueSaturationCollector(delegate, 0.9, 10);
+ for (int i = 0; i < numDocs; i++) {
+ queueSaturationCollector.collect(i, 1.0f - i * 1e-3f);
+ if (i % 10 == 0) {
+ queueSaturationCollector.nextCandidate();
+ }
+ if (queueSaturationCollector.earlyTerminated()) {
+ assertEquals(120, i);
+ break;
+ }
+ }
+ }
+
+ @Test
+ public void testDelegateVsSaturateEarlyExit() {
+ Random random = random();
+ int numDocs = 10000;
+ int k = random.nextInt(1, 100);
+ KnnCollector delegate = new TopKnnCollector(k, numDocs);
+ HnswQueueSaturationCollector queueSaturationCollector =
+ new HnswQueueSaturationCollector(delegate, 0.5, 1);
+ for (int i = 0; i < random.nextInt(numDocs); i++) {
+ queueSaturationCollector.collect(random.nextInt(numDocs), random.nextFloat(1.0f));
+ if (i % 10 == 0) {
+ queueSaturationCollector.nextCandidate();
+ }
+ boolean earlyTerminatedSaturation = queueSaturationCollector.earlyTerminated();
+ boolean earlyTerminatedDelegate = delegate.earlyTerminated();
+ assertTrue(earlyTerminatedSaturation || !earlyTerminatedDelegate);
+ }
+ }
+
+ @Test
+ public void testEarlyExitRelation() {
+ Random random = random();
+ int numDocs = 10000;
+ int k = random.nextInt(100);
+ KnnCollector delegate = new TopKnnCollector(k, random.nextInt(numDocs));
+ HnswQueueSaturationCollector queueSaturationCollector =
+ new HnswQueueSaturationCollector(delegate, 0.5, 1);
+ for (int i = 0; i < random.nextInt(numDocs); i++) {
+ queueSaturationCollector.collect(random.nextInt(numDocs), random.nextFloat(1.0f));
+ if (i % 10 == 0) {
+ queueSaturationCollector.nextCandidate();
+ }
+ if (delegate.earlyTerminated()) {
+ TopDocs topDocs = queueSaturationCollector.topDocs();
+ assertEquals(TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO, topDocs.totalHits.relation());
+ }
+ if (queueSaturationCollector.earlyTerminated()) {
+ TopDocs topDocs = queueSaturationCollector.topDocs();
+ assertEquals(TotalHits.Relation.EQUAL_TO, topDocs.totalHits.relation());
+ break;
+ }
+ }
+ }
+}
diff --git a/lucene/core/src/test/org/apache/lucene/search/TestPatienceByteVectorQuery.java b/lucene/core/src/test/org/apache/lucene/search/TestPatienceByteVectorQuery.java
new file mode 100644
index 00000000000..60979c27c02
--- /dev/null
+++ b/lucene/core/src/test/org/apache/lucene/search/TestPatienceByteVectorQuery.java
@@ -0,0 +1,114 @@
+/*
+ * 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.search;
+
+import static org.apache.lucene.search.TestKnnByteVectorQuery.floatToBytes;
+
+import java.io.IOException;
+import org.apache.lucene.document.Field;
+import org.apache.lucene.document.KnnByteVectorField;
+import org.apache.lucene.index.DirectoryReader;
+import org.apache.lucene.index.IndexReader;
+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;
+
+public class TestPatienceByteVectorQuery extends BaseKnnVectorQueryTestCase {
+
+ @Override
+ PatienceKnnVectorQuery getKnnVectorQuery(String field, float[] query, int k, Query queryFilter) {
+ return PatienceKnnVectorQuery.fromByteQuery(
+ new KnnByteVectorQuery(field, floatToBytes(query), k, queryFilter));
+ }
+
+ @Override
+ AbstractKnnVectorQuery getThrowingKnnVectorQuery(String field, float[] vec, int k, Query query) {
+ return PatienceKnnVectorQuery.fromByteQuery(
+ new ThrowingKnnVectorQuery(field, floatToBytes(vec), k, query));
+ }
+
+ @Override
+ AbstractKnnVectorQuery getCappedResultsThrowingKnnVectorQuery(
+ String field, float[] vec, int k, Query query, int maxResults) {
+ return new TestKnnByteVectorQuery.CappedResultsThrowingKnnVectorQuery(
+ field, floatToBytes(vec), k, query, maxResults);
+ }
+
+ @Override
+ float[] randomVector(int dim) {
+ byte[] b = TestVectorUtil.randomVectorBytes(dim);
+ float[] v = new float[b.length];
+ int vi = 0;
+ for (int i = 0; i < v.length; i++) {
+ v[vi++] = b[i];
+ }
+ return v;
+ }
+
+ @Override
+ Field getKnnVectorField(
+ String name, float[] vector, VectorSimilarityFunction similarityFunction) {
+ return new KnnByteVectorField(name, floatToBytes(vector), similarityFunction);
+ }
+
+ @Override
+ Field getKnnVectorField(String name, float[] vector) {
+ return new KnnByteVectorField(name, floatToBytes(vector), VectorSimilarityFunction.EUCLIDEAN);
+ }
+
+ public void testToString() throws IOException {
+ try (Directory indexStore =
+ getIndexStore("field", new float[] {0, 1}, new float[] {1, 2}, new float[] {0, 0});
+ IndexReader reader = DirectoryReader.open(indexStore)) {
+ AbstractKnnVectorQuery query = getKnnVectorQuery("field", new float[] {0.0f, 1.0f}, 10);
+ assertEquals(
+ "PatienceKnnVectorQuery{saturationThreshold=0.995, patience=7, delegate=KnnByteVectorQuery:field[0,...][10]}",
+ query.toString("ignored"));
+
+ assertDocScoreQueryToString(query.rewrite(newSearcher(reader)));
+
+ // test with filter
+ Query filter = new TermQuery(new Term("id", "text"));
+ query = getKnnVectorQuery("field", new float[] {0.0f, 1.0f}, 10, filter);
+ assertEquals(
+ "PatienceKnnVectorQuery{saturationThreshold=0.995, patience=7, delegate=KnnByteVectorQuery:field[0,...][10][id:text]}",
+ query.toString("ignored"));
+ }
+ }
+
+ static class ThrowingKnnVectorQuery extends KnnByteVectorQuery {
+
+ public ThrowingKnnVectorQuery(String field, byte[] target, int k, Query filter) {
+ super(field, target, k, filter, new KnnSearchStrategy.Hnsw(0));
+ }
+
+ @Override
+ protected TopDocs exactSearch(
+ LeafReaderContext context, DocIdSetIterator acceptIterator, QueryTimeout queryTimeout) {
+ throw new UnsupportedOperationException("exact search is not supported");
+ }
+
+ @Override
+ public String toString(String field) {
+ return "ThrowingKnnVectorQuery{" + field + "}";
+ }
+ }
+}
diff --git a/lucene/core/src/test/org/apache/lucene/search/TestPatienceFloatVectorQuery.java b/lucene/core/src/test/org/apache/lucene/search/TestPatienceFloatVectorQuery.java
new file mode 100644
index 00000000000..21be86b30ca
--- /dev/null
+++ b/lucene/core/src/test/org/apache/lucene/search/TestPatienceFloatVectorQuery.java
@@ -0,0 +1,105 @@
+/*
+ * 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.search;
+
+import java.io.IOException;
+import org.apache.lucene.document.Field;
+import org.apache.lucene.document.KnnFloatVectorField;
+import org.apache.lucene.index.DirectoryReader;
+import org.apache.lucene.index.IndexReader;
+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;
+
+public class TestPatienceFloatVectorQuery extends BaseKnnVectorQueryTestCase {
+
+ @Override
+ PatienceKnnVectorQuery getKnnVectorQuery(String field, float[] query, int k, Query queryFilter) {
+ return PatienceKnnVectorQuery.fromFloatQuery(
+ new KnnFloatVectorQuery(field, query, k, queryFilter));
+ }
+
+ @Override
+ AbstractKnnVectorQuery getThrowingKnnVectorQuery(String field, float[] vec, int k, Query query) {
+ return PatienceKnnVectorQuery.fromFloatQuery(new ThrowingKnnVectorQuery(field, vec, k, query));
+ }
+
+ @Override
+ AbstractKnnVectorQuery getCappedResultsThrowingKnnVectorQuery(
+ String field, float[] vec, int k, Query query, int maxResults) {
+ return new TestKnnFloatVectorQuery.CappedResultsThrowingKnnVectorQuery(
+ field, vec, k, query, maxResults);
+ }
+
+ @Override
+ float[] randomVector(int dim) {
+ return TestVectorUtil.randomVector(dim);
+ }
+
+ @Override
+ Field getKnnVectorField(
+ String name, float[] vector, VectorSimilarityFunction similarityFunction) {
+ return new KnnFloatVectorField(name, vector, similarityFunction);
+ }
+
+ @Override
+ Field getKnnVectorField(String name, float[] vector) {
+ return new KnnFloatVectorField(name, vector);
+ }
+
+ public void testToString() throws IOException {
+ try (Directory indexStore =
+ getIndexStore("field", new float[] {0, 1}, new float[] {1, 2}, new float[] {0, 0});
+ IndexReader reader = DirectoryReader.open(indexStore)) {
+ AbstractKnnVectorQuery query = getKnnVectorQuery("field", new float[] {0.0f, 1.0f}, 10);
+ assertEquals(
+ "PatienceKnnVectorQuery{saturationThreshold=0.995, patience=7, delegate=KnnFloatVectorQuery:field[0.0,...][10]}",
+ query.toString("ignored"));
+
+ assertDocScoreQueryToString(query.rewrite(newSearcher(reader)));
+
+ // test with filter
+ Query filter = new TermQuery(new Term("id", "text"));
+ query = getKnnVectorQuery("field", new float[] {0.0f, 1.0f}, 10, filter);
+ assertEquals(
+ "PatienceKnnVectorQuery{saturationThreshold=0.995, patience=7, delegate=KnnFloatVectorQuery:field[0.0,...][10][id:text]}",
+ query.toString("ignored"));
+ }
+ }
+
+ static class ThrowingKnnVectorQuery extends KnnFloatVectorQuery {
+
+ public ThrowingKnnVectorQuery(String field, float[] target, int k, Query filter) {
+ super(field, target, k, filter, new KnnSearchStrategy.Hnsw(0));
+ }
+
+ @Override
+ protected TopDocs exactSearch(
+ LeafReaderContext context, DocIdSetIterator acceptIterator, QueryTimeout queryTimeout) {
+ throw new UnsupportedOperationException("exact search is not supported");
+ }
+
+ @Override
+ public String toString(String field) {
+ return "ThrowingKnnVectorQuery{" + field + "}";
+ }
+ }
+}