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/03 16:13:45 UTC
(lucene) branch branch_10x updated (9965ec46814 -> 6f5451ca1aa)
This is an automated email from the ASF dual-hosted git repository.
benwtrent pushed a change to branch branch_10x
in repository https://gitbox.apache.org/repos/asf/lucene.git
from 9965ec46814 Filtered disjunctions may miss some top hits. (#14186)
new 1baccbd77b9 Adjust knn merge stability testing (#14172)
new 823f2c865cf Merge remote-tracking branch 'refs/remotes/origin/branch_10x' into branch_10x
new 6f5451ca1aa Adjusts Seeded knn searches to clean up user and internal interfaces (#14170)
The 3 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
lucene/CHANGES.txt | 4 +-
.../apache/lucene/search/AbstractKnnCollector.java | 17 +-
.../search/AbstractVectorSimilarityQuery.java | 4 +-
.../lucene/search/ByteVectorSimilarityQuery.java | 2 +-
.../lucene/search/FloatVectorSimilarityQuery.java | 2 +-
.../apache/lucene/search/KnnByteVectorQuery.java | 2 +-
.../org/apache/lucene/search/KnnCollector.java | 16 +-
.../apache/lucene/search/KnnFloatVectorQuery.java | 2 +-
.../lucene/search/SeededKnnByteVectorQuery.java | 97 -------
.../lucene/search/SeededKnnFloatVectorQuery.java | 97 -------
.../apache/lucene/search/SeededKnnVectorQuery.java | 321 +++++++++++++++++++++
.../search/TimeLimitingKnnCollectorManager.java | 7 +-
.../org/apache/lucene/search/TopKnnCollector.java | 11 +-
.../lucene/search/knn/EntryPointProvider.java | 28 --
.../lucene/search/knn/KnnCollectorManager.java | 5 +-
.../lucene/search/knn/KnnSearchStrategy.java | 63 ++++
.../lucene/search/knn/SeededKnnCollector.java | 48 ---
.../search/knn/SeededKnnCollectorManager.java | 177 ------------
.../lucene/search/knn/TopKnnCollectorManager.java | 9 +-
.../util/hnsw/AbstractHnswGraphSearcher.java | 84 ++++++
.../apache/lucene/util/hnsw/HnswGraphSearcher.java | 65 ++---
.../lucene/util/hnsw/SeededHnswGraphSearcher.java | 78 +++++
.../lucene/search/BaseKnnVectorQueryTestCase.java | 4 +-
.../search/TestSeededKnnByteVectorQuery.java | 192 +++++++++---
.../search/TestSeededKnnFloatVectorQuery.java | 124 ++++----
.../DiversifyingChildrenByteKnnVectorQuery.java | 2 +-
.../DiversifyingChildrenFloatKnnVectorQuery.java | 2 +-
.../DiversifyingNearestChildrenKnnCollector.java | 16 +-
...rsifyingNearestChildrenKnnCollectorManager.java | 11 +-
.../tests/index/BaseIndexFileFormatTestCase.java | 3 +
.../tests/index/BaseKnnVectorsFormatTestCase.java | 8 +
31 files changed, 892 insertions(+), 609 deletions(-)
delete mode 100644 lucene/core/src/java/org/apache/lucene/search/SeededKnnByteVectorQuery.java
delete mode 100644 lucene/core/src/java/org/apache/lucene/search/SeededKnnFloatVectorQuery.java
create mode 100644 lucene/core/src/java/org/apache/lucene/search/SeededKnnVectorQuery.java
delete mode 100644 lucene/core/src/java/org/apache/lucene/search/knn/EntryPointProvider.java
create mode 100644 lucene/core/src/java/org/apache/lucene/search/knn/KnnSearchStrategy.java
delete mode 100644 lucene/core/src/java/org/apache/lucene/search/knn/SeededKnnCollector.java
delete mode 100644 lucene/core/src/java/org/apache/lucene/search/knn/SeededKnnCollectorManager.java
create mode 100644 lucene/core/src/java/org/apache/lucene/util/hnsw/AbstractHnswGraphSearcher.java
create mode 100644 lucene/core/src/java/org/apache/lucene/util/hnsw/SeededHnswGraphSearcher.java
(lucene) 03/03: Adjusts Seeded knn searches to clean up user and internal interfaces (#14170)
Posted by be...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
benwtrent pushed a commit to branch branch_10x
in repository https://gitbox.apache.org/repos/asf/lucene.git
commit 6f5451ca1aa1817a3b7a4d24576a23b71e9a1f95
Author: Benjamin Trent <be...@gmail.com>
AuthorDate: Mon Feb 3 11:06:18 2025 -0500
Adjusts Seeded knn searches to clean up user and internal interfaces (#14170)
This is a bugfix and refactor for seeded knn searches.
First, Since we are using collectors, we don't actually need unique queries for every input type. Consequently, I have collapsed the two individual seeded queries into a single query that delegates to a provided kNN query. Then the collector manager is simply wrapped, so that the entry points can be provided.
Second, the interactions in the hnsw graph were not clear. Consequently, I did a minor refactor of HNSW searcher to have a "SeededSearcher", where instead of searching the graph for the entry points, it provides them directly.
Third, instead of continually overloading collectors, I opted to add a new "searchstrategy" value to KnnCollector. This way various strategies can be executed with different options. I think Seeded could eventually be replaced with something.
---
lucene/CHANGES.txt | 4 +-
.../apache/lucene/search/AbstractKnnCollector.java | 17 +-
.../search/AbstractVectorSimilarityQuery.java | 4 +-
.../lucene/search/ByteVectorSimilarityQuery.java | 2 +-
.../lucene/search/FloatVectorSimilarityQuery.java | 2 +-
.../apache/lucene/search/KnnByteVectorQuery.java | 2 +-
.../org/apache/lucene/search/KnnCollector.java | 16 +-
.../apache/lucene/search/KnnFloatVectorQuery.java | 2 +-
.../lucene/search/SeededKnnByteVectorQuery.java | 97 -------
.../lucene/search/SeededKnnFloatVectorQuery.java | 97 -------
.../apache/lucene/search/SeededKnnVectorQuery.java | 321 +++++++++++++++++++++
.../search/TimeLimitingKnnCollectorManager.java | 7 +-
.../org/apache/lucene/search/TopKnnCollector.java | 11 +-
.../lucene/search/knn/EntryPointProvider.java | 28 --
.../lucene/search/knn/KnnCollectorManager.java | 5 +-
.../lucene/search/knn/KnnSearchStrategy.java | 63 ++++
.../lucene/search/knn/SeededKnnCollector.java | 48 ---
.../search/knn/SeededKnnCollectorManager.java | 177 ------------
.../lucene/search/knn/TopKnnCollectorManager.java | 9 +-
.../util/hnsw/AbstractHnswGraphSearcher.java | 84 ++++++
.../apache/lucene/util/hnsw/HnswGraphSearcher.java | 65 ++---
.../lucene/util/hnsw/SeededHnswGraphSearcher.java | 78 +++++
.../lucene/search/BaseKnnVectorQueryTestCase.java | 4 +-
.../search/TestSeededKnnByteVectorQuery.java | 192 +++++++++---
.../search/TestSeededKnnFloatVectorQuery.java | 124 ++++----
.../DiversifyingChildrenByteKnnVectorQuery.java | 2 +-
.../DiversifyingChildrenFloatKnnVectorQuery.java | 2 +-
.../DiversifyingNearestChildrenKnnCollector.java | 16 +-
...rsifyingNearestChildrenKnnCollectorManager.java | 11 +-
29 files changed, 881 insertions(+), 609 deletions(-)
diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index ef8d89b31b8..8864e949ada 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -16,8 +16,8 @@ API Changes
New Features
---------------------
-* GITHUB#14084, GITHUB#13635, GITHUB#13634: Adds new `SeededKnnByteVectorQuery` and `SeededKnnFloatVectorQuery`
- queries. These queries allow for the vector search entry points to be initialized via a `seed` query. This follows
+* GITHUB#14084, GITHUB#13635, GITHUB#13634, GITHUB#14170: Adds new `SeededKnnVectorQuery` query.
+ These queries allow for the vector search entry points to be initialized via a `seed` query. This follows
the research provided via https://arxiv.org/abs/2307.16779. (Sean MacAvaney, Ben Trent).
diff --git a/lucene/core/src/java/org/apache/lucene/search/AbstractKnnCollector.java b/lucene/core/src/java/org/apache/lucene/search/AbstractKnnCollector.java
index 0b1b18d1dbd..34484559d0b 100644
--- a/lucene/core/src/java/org/apache/lucene/search/AbstractKnnCollector.java
+++ b/lucene/core/src/java/org/apache/lucene/search/AbstractKnnCollector.java
@@ -17,6 +17,8 @@
package org.apache.lucene.search;
+import org.apache.lucene.search.knn.KnnSearchStrategy;
+
/**
* AbstractKnnCollector is the default implementation for a knn collector used for gathering kNN
* results and providing topDocs from the gathered neighbors
@@ -25,11 +27,17 @@ public abstract class AbstractKnnCollector implements KnnCollector {
protected long visitedCount;
private final long visitLimit;
+ private final KnnSearchStrategy searchStrategy;
private final int k;
- protected AbstractKnnCollector(int k, long visitLimit) {
- this.visitLimit = visitLimit;
+ protected AbstractKnnCollector(int k, long visitLimit, KnnSearchStrategy searchStrategy) {
this.k = k;
+ this.searchStrategy = searchStrategy;
+ this.visitLimit = visitLimit;
+ }
+
+ protected AbstractKnnCollector(int k, long visitLimit) {
+ this(k, visitLimit, null);
}
@Override
@@ -68,4 +76,9 @@ public abstract class AbstractKnnCollector implements KnnCollector {
@Override
public abstract TopDocs topDocs();
+
+ @Override
+ public KnnSearchStrategy getSearchStrategy() {
+ return searchStrategy;
+ }
}
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 dde9ce76ac3..39c81651cf2 100644
--- a/lucene/core/src/java/org/apache/lucene/search/AbstractVectorSimilarityQuery.java
+++ b/lucene/core/src/java/org/apache/lucene/search/AbstractVectorSimilarityQuery.java
@@ -61,8 +61,8 @@ abstract class AbstractVectorSimilarityQuery extends Query {
}
protected KnnCollectorManager getKnnCollectorManager() {
- return (visitedLimit, context) ->
- new VectorSimilarityCollector(traversalSimilarity, resultSimilarity, visitedLimit);
+ return (visitLimit, searchStrategy, context) ->
+ new VectorSimilarityCollector(traversalSimilarity, resultSimilarity, visitLimit);
}
abstract VectorScorer createVectorScorer(LeafReaderContext context) throws IOException;
diff --git a/lucene/core/src/java/org/apache/lucene/search/ByteVectorSimilarityQuery.java b/lucene/core/src/java/org/apache/lucene/search/ByteVectorSimilarityQuery.java
index c547f1face7..90d7b1ee16e 100644
--- a/lucene/core/src/java/org/apache/lucene/search/ByteVectorSimilarityQuery.java
+++ b/lucene/core/src/java/org/apache/lucene/search/ByteVectorSimilarityQuery.java
@@ -113,7 +113,7 @@ public class ByteVectorSimilarityQuery extends AbstractVectorSimilarityQuery {
int visitLimit,
KnnCollectorManager knnCollectorManager)
throws IOException {
- KnnCollector collector = knnCollectorManager.newCollector(visitLimit, context);
+ KnnCollector collector = knnCollectorManager.newCollector(visitLimit, null, context);
context.reader().searchNearestVectors(field, target, collector, acceptDocs);
return collector.topDocs();
}
diff --git a/lucene/core/src/java/org/apache/lucene/search/FloatVectorSimilarityQuery.java b/lucene/core/src/java/org/apache/lucene/search/FloatVectorSimilarityQuery.java
index 4c7078ac140..0a440faefce 100644
--- a/lucene/core/src/java/org/apache/lucene/search/FloatVectorSimilarityQuery.java
+++ b/lucene/core/src/java/org/apache/lucene/search/FloatVectorSimilarityQuery.java
@@ -115,7 +115,7 @@ public class FloatVectorSimilarityQuery extends AbstractVectorSimilarityQuery {
int visitLimit,
KnnCollectorManager knnCollectorManager)
throws IOException {
- KnnCollector collector = knnCollectorManager.newCollector(visitLimit, context);
+ KnnCollector collector = knnCollectorManager.newCollector(visitLimit, null, context);
context.reader().searchNearestVectors(field, target, collector, acceptDocs);
return collector.topDocs();
}
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 05157ab65cb..51d74e0257d 100644
--- a/lucene/core/src/java/org/apache/lucene/search/KnnByteVectorQuery.java
+++ b/lucene/core/src/java/org/apache/lucene/search/KnnByteVectorQuery.java
@@ -83,7 +83,7 @@ public class KnnByteVectorQuery extends AbstractKnnVectorQuery {
int visitedLimit,
KnnCollectorManager knnCollectorManager)
throws IOException {
- KnnCollector knnCollector = knnCollectorManager.newCollector(visitedLimit, context);
+ KnnCollector knnCollector = knnCollectorManager.newCollector(visitedLimit, null, 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 f694d8f7085..43086c4ad57 100644
--- a/lucene/core/src/java/org/apache/lucene/search/KnnCollector.java
+++ b/lucene/core/src/java/org/apache/lucene/search/KnnCollector.java
@@ -17,6 +17,8 @@
package org.apache.lucene.search;
+import org.apache.lucene.search.knn.KnnSearchStrategy;
+
/**
* KnnCollector is a knn collector used for gathering kNN results and providing topDocs from the
* gathered neighbors
@@ -86,6 +88,13 @@ public interface KnnCollector {
*/
TopDocs topDocs();
+ /**
+ * @return the search strategy used by this collector, can be null
+ */
+ default KnnSearchStrategy getSearchStrategy() {
+ return null;
+ }
+
/**
* KnnCollector.Decorator is the base class for decorators of KnnCollector objects, which extend
* the object with new behaviors.
@@ -93,7 +102,7 @@ public interface KnnCollector {
* @lucene.experimental
*/
abstract class Decorator implements KnnCollector {
- private final KnnCollector collector;
+ protected final KnnCollector collector;
public Decorator(KnnCollector collector) {
this.collector = collector;
@@ -138,5 +147,10 @@ public interface KnnCollector {
public TopDocs topDocs() {
return collector.topDocs();
}
+
+ @Override
+ public KnnSearchStrategy getSearchStrategy() {
+ return collector.getSearchStrategy();
+ }
}
}
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 c7d6fdb3608..8f8a1a72177 100644
--- a/lucene/core/src/java/org/apache/lucene/search/KnnFloatVectorQuery.java
+++ b/lucene/core/src/java/org/apache/lucene/search/KnnFloatVectorQuery.java
@@ -84,7 +84,7 @@ public class KnnFloatVectorQuery extends AbstractKnnVectorQuery {
int visitedLimit,
KnnCollectorManager knnCollectorManager)
throws IOException {
- KnnCollector knnCollector = knnCollectorManager.newCollector(visitedLimit, context);
+ KnnCollector knnCollector = knnCollectorManager.newCollector(visitedLimit, null, context);
LeafReader reader = context.reader();
FloatVectorValues floatVectorValues = reader.getFloatVectorValues(field);
if (floatVectorValues == null) {
diff --git a/lucene/core/src/java/org/apache/lucene/search/SeededKnnByteVectorQuery.java b/lucene/core/src/java/org/apache/lucene/search/SeededKnnByteVectorQuery.java
deleted file mode 100644
index 980b6869c34..00000000000
--- a/lucene/core/src/java/org/apache/lucene/search/SeededKnnByteVectorQuery.java
+++ /dev/null
@@ -1,97 +0,0 @@
-/*
- * 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.ByteVectorValues;
-import org.apache.lucene.search.knn.KnnCollectorManager;
-import org.apache.lucene.search.knn.SeededKnnCollectorManager;
-
-/**
- * This is a version of knn byte vector query that provides a query seed to initiate the vector
- * search. NOTE: The underlying format is free to ignore the provided seed
- *
- * <p>See <a href="https://dl.acm.org/doi/10.1145/3539618.3591715">"Lexically-Accelerated Dense
- * Retrieval"</a> (Kulkarni, Hrishikesh and MacAvaney, Sean and Goharian, Nazli and Frieder, Ophir).
- * In SIGIR '23: Proceedings of the 46th International ACM SIGIR Conference on Research and
- * Development in Information Retrieval Pages 152 - 162
- *
- * @lucene.experimental
- */
-public class SeededKnnByteVectorQuery extends KnnByteVectorQuery {
- final Query seed;
- final Weight seedWeight;
-
- /**
- * Construct a new SeededKnnByteVectorQuery instance
- *
- * @param field knn byte vector field to query
- * @param target the query vector
- * @param k number of neighbors to return
- * @param filter a filter on the neighbors to return
- * @param seed a query seed to initiate the vector format search
- */
- public SeededKnnByteVectorQuery(String field, byte[] target, int k, Query filter, Query seed) {
- super(field, target, k, filter);
- this.seed = Objects.requireNonNull(seed);
- this.seedWeight = null;
- }
-
- SeededKnnByteVectorQuery(String field, byte[] target, int k, Query filter, Weight seedWeight) {
- super(field, target, k, filter);
- this.seed = null;
- this.seedWeight = Objects.requireNonNull(seedWeight);
- }
-
- @Override
- public Query rewrite(IndexSearcher indexSearcher) throws IOException {
- if (seedWeight != null) {
- return super.rewrite(indexSearcher);
- }
- BooleanQuery.Builder booleanSeedQueryBuilder =
- new BooleanQuery.Builder()
- .add(seed, BooleanClause.Occur.MUST)
- .add(new FieldExistsQuery(field), BooleanClause.Occur.FILTER);
- if (filter != null) {
- booleanSeedQueryBuilder.add(filter, BooleanClause.Occur.FILTER);
- }
- Query seedRewritten = indexSearcher.rewrite(booleanSeedQueryBuilder.build());
- Weight seedWeight = indexSearcher.createWeight(seedRewritten, ScoreMode.TOP_SCORES, 1f);
- SeededKnnByteVectorQuery rewritten =
- new SeededKnnByteVectorQuery(field, target, k, filter, seedWeight);
- return rewritten.rewrite(indexSearcher);
- }
-
- @Override
- protected KnnCollectorManager getKnnCollectorManager(int k, IndexSearcher searcher) {
- if (seedWeight == null) {
- throw new UnsupportedOperationException("must be rewritten before constructing manager");
- }
- return new SeededKnnCollectorManager(
- super.getKnnCollectorManager(k, searcher),
- seedWeight,
- k,
- leaf -> {
- ByteVectorValues vv = leaf.getByteVectorValues(field);
- if (vv == null) {
- ByteVectorValues.checkField(leaf.getContext().reader(), field);
- }
- return vv;
- });
- }
-}
diff --git a/lucene/core/src/java/org/apache/lucene/search/SeededKnnFloatVectorQuery.java b/lucene/core/src/java/org/apache/lucene/search/SeededKnnFloatVectorQuery.java
deleted file mode 100644
index 02a33bdcdef..00000000000
--- a/lucene/core/src/java/org/apache/lucene/search/SeededKnnFloatVectorQuery.java
+++ /dev/null
@@ -1,97 +0,0 @@
-/*
- * 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.FloatVectorValues;
-import org.apache.lucene.search.knn.KnnCollectorManager;
-import org.apache.lucene.search.knn.SeededKnnCollectorManager;
-
-/**
- * This is a version of knn float vector query that provides a query seed to initiate the vector
- * search. NOTE: The underlying format is free to ignore the provided seed.
- *
- * <p>See <a href="https://dl.acm.org/doi/10.1145/3539618.3591715">"Lexically-Accelerated Dense
- * Retrieval"</a> (Kulkarni, Hrishikesh and MacAvaney, Sean and Goharian, Nazli and Frieder, Ophir).
- * In SIGIR '23: Proceedings of the 46th International ACM SIGIR Conference on Research and
- * Development in Information Retrieval Pages 152 - 162
- *
- * @lucene.experimental
- */
-public class SeededKnnFloatVectorQuery extends KnnFloatVectorQuery {
- final Query seed;
- final Weight seedWeight;
-
- /**
- * Construct a new SeededKnnFloatVectorQuery instance
- *
- * @param field knn float vector field to query
- * @param target the query vector
- * @param k number of neighbors to return
- * @param filter a filter on the neighbors to return
- * @param seed a query seed to initiate the vector format search
- */
- public SeededKnnFloatVectorQuery(String field, float[] target, int k, Query filter, Query seed) {
- super(field, target, k, filter);
- this.seed = Objects.requireNonNull(seed);
- this.seedWeight = null;
- }
-
- SeededKnnFloatVectorQuery(String field, float[] target, int k, Query filter, Weight seedWeight) {
- super(field, target, k, filter);
- this.seed = null;
- this.seedWeight = Objects.requireNonNull(seedWeight);
- }
-
- @Override
- public Query rewrite(IndexSearcher indexSearcher) throws IOException {
- if (seedWeight != null) {
- return super.rewrite(indexSearcher);
- }
- BooleanQuery.Builder booleanSeedQueryBuilder =
- new BooleanQuery.Builder()
- .add(seed, BooleanClause.Occur.MUST)
- .add(new FieldExistsQuery(field), BooleanClause.Occur.FILTER);
- if (filter != null) {
- booleanSeedQueryBuilder.add(filter, BooleanClause.Occur.FILTER);
- }
- Query seedRewritten = indexSearcher.rewrite(booleanSeedQueryBuilder.build());
- Weight seedWeight = indexSearcher.createWeight(seedRewritten, ScoreMode.TOP_SCORES, 1f);
- SeededKnnFloatVectorQuery rewritten =
- new SeededKnnFloatVectorQuery(field, target, k, filter, seedWeight);
- return rewritten.rewrite(indexSearcher);
- }
-
- @Override
- protected KnnCollectorManager getKnnCollectorManager(int k, IndexSearcher searcher) {
- if (seedWeight == null) {
- throw new UnsupportedOperationException("must be rewritten before constructing manager");
- }
- return new SeededKnnCollectorManager(
- super.getKnnCollectorManager(k, searcher),
- seedWeight,
- k,
- leaf -> {
- FloatVectorValues vv = leaf.getFloatVectorValues(field);
- if (vv == null) {
- FloatVectorValues.checkField(leaf.getContext().reader(), field);
- }
- return vv;
- });
- }
-}
diff --git a/lucene/core/src/java/org/apache/lucene/search/SeededKnnVectorQuery.java b/lucene/core/src/java/org/apache/lucene/search/SeededKnnVectorQuery.java
new file mode 100644
index 00000000000..946fad57732
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/search/SeededKnnVectorQuery.java
@@ -0,0 +1,321 @@
+/*
+ * 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.Arrays;
+import java.util.Objects;
+import org.apache.lucene.codecs.lucene90.IndexedDISI;
+import org.apache.lucene.index.FieldInfo;
+import org.apache.lucene.index.KnnVectorValues;
+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.Bits;
+
+/**
+ * This is a version of knn vector query that provides a query seed to initiate the vector search.
+ * NOTE: The underlying format is free to ignore the provided seed
+ *
+ * <p>See <a href="https://dl.acm.org/doi/10.1145/3539618.3591715">"Lexically-Accelerated Dense
+ * Retrieval"</a> (Kulkarni, Hrishikesh and MacAvaney, Sean and Goharian, Nazli and Frieder, Ophir).
+ * In SIGIR '23: Proceedings of the 46th International ACM SIGIR Conference on Research and
+ * Development in Information Retrieval Pages 152 - 162
+ *
+ * @lucene.experimental
+ */
+public class SeededKnnVectorQuery extends AbstractKnnVectorQuery {
+ final Query seed;
+ final Weight seedWeight;
+ final AbstractKnnVectorQuery delegate;
+
+ /**
+ * Construct a new SeededKnnVectorQuery instance for a float vector field
+ *
+ * @param knnQuery the knn query to be seeded
+ * @param seed a query seed to initiate the vector format search
+ * @return a new SeededKnnVectorQuery instance
+ * @lucene.experimental
+ */
+ public static SeededKnnVectorQuery fromFloatQuery(KnnFloatVectorQuery knnQuery, Query seed) {
+ return new SeededKnnVectorQuery(knnQuery, seed, null);
+ }
+
+ /**
+ * Construct a new SeededKnnVectorQuery instance for a byte vector field
+ *
+ * @param knnQuery the knn query to be seeded
+ * @param seed a query seed to initiate the vector format search
+ * @return a new SeededKnnVectorQuery instance
+ * @lucene.experimental
+ */
+ public static SeededKnnVectorQuery fromByteQuery(KnnByteVectorQuery knnQuery, Query seed) {
+ return new SeededKnnVectorQuery(knnQuery, seed, null);
+ }
+
+ SeededKnnVectorQuery(AbstractKnnVectorQuery knnQuery, Query seed, Weight seedWeight) {
+ super(knnQuery.field, knnQuery.k, knnQuery.filter);
+ this.delegate = knnQuery;
+ this.seed = Objects.requireNonNull(seed);
+ this.seedWeight = seedWeight;
+ }
+
+ @Override
+ public String toString(String field) {
+ return "SeededKnnVectorQuery{"
+ + "seed="
+ + seed
+ + ", seedWeight="
+ + seedWeight
+ + ", delegate="
+ + delegate
+ + '}';
+ }
+
+ @Override
+ public Query rewrite(IndexSearcher indexSearcher) throws IOException {
+ if (seedWeight != null) {
+ return super.rewrite(indexSearcher);
+ }
+ SeededKnnVectorQuery rewritten =
+ new SeededKnnVectorQuery(delegate, seed, createSeedWeight(indexSearcher));
+ return rewritten.rewrite(indexSearcher);
+ }
+
+ Weight createSeedWeight(IndexSearcher indexSearcher) throws IOException {
+ BooleanQuery.Builder booleanSeedQueryBuilder =
+ new BooleanQuery.Builder()
+ .add(seed, BooleanClause.Occur.MUST)
+ .add(new FieldExistsQuery(field), BooleanClause.Occur.FILTER);
+ if (filter != null) {
+ booleanSeedQueryBuilder.add(filter, BooleanClause.Occur.FILTER);
+ }
+ Query seedRewritten = indexSearcher.rewrite(booleanSeedQueryBuilder.build());
+ return indexSearcher.createWeight(seedRewritten, ScoreMode.TOP_SCORES, 1f);
+ }
+
+ @Override
+ protected TopDocs approximateSearch(
+ LeafReaderContext context,
+ Bits acceptDocs,
+ int visitedLimit,
+ KnnCollectorManager knnCollectorManager)
+ throws IOException {
+ return delegate.approximateSearch(
+ context, acceptDocs, visitedLimit, new SeededCollectorManager(knnCollectorManager));
+ }
+
+ @Override
+ protected KnnCollectorManager getKnnCollectorManager(int k, IndexSearcher searcher) {
+ return delegate.getKnnCollectorManager(k, searcher);
+ }
+
+ @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;
+ SeededKnnVectorQuery that = (SeededKnnVectorQuery) o;
+ return Objects.equals(seed, that.seed)
+ && Objects.equals(seedWeight, that.seedWeight)
+ && Objects.equals(delegate, that.delegate);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(super.hashCode(), seed, seedWeight, 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);
+ }
+
+ private static class MappedDISI extends DocIdSetIterator {
+ KnnVectorValues.DocIndexIterator indexedDISI;
+ DocIdSetIterator sourceDISI;
+
+ private MappedDISI(KnnVectorValues.DocIndexIterator indexedDISI, DocIdSetIterator sourceDISI) {
+ this.indexedDISI = indexedDISI;
+ this.sourceDISI = sourceDISI;
+ }
+
+ /**
+ * Advances the source iterator to the first document number that is greater than or equal to
+ * the provided target and returns the corresponding index.
+ */
+ @Override
+ public int advance(int target) throws IOException {
+ int newTarget = sourceDISI.advance(target);
+ if (newTarget != NO_MORE_DOCS) {
+ indexedDISI.advance(newTarget);
+ }
+ return docID();
+ }
+
+ @Override
+ public long cost() {
+ return sourceDISI.cost();
+ }
+
+ @Override
+ public int docID() {
+ if (indexedDISI.docID() == NO_MORE_DOCS || sourceDISI.docID() == NO_MORE_DOCS) {
+ return NO_MORE_DOCS;
+ }
+ return indexedDISI.index();
+ }
+
+ /** Advances to the next document in the source iterator and returns the corresponding index. */
+ @Override
+ public int nextDoc() throws IOException {
+ int newTarget = sourceDISI.nextDoc();
+ if (newTarget != NO_MORE_DOCS) {
+ indexedDISI.advance(newTarget);
+ }
+ return docID();
+ }
+ }
+
+ private static class TopDocsDISI extends DocIdSetIterator {
+ private final int[] sortedDocIds;
+ private int idx = -1;
+
+ private TopDocsDISI(TopDocs topDocs) {
+ sortedDocIds = new int[topDocs.scoreDocs.length];
+ for (int i = 0; i < topDocs.scoreDocs.length; i++) {
+ sortedDocIds[i] = topDocs.scoreDocs[i].doc;
+ }
+ Arrays.sort(sortedDocIds);
+ }
+
+ @Override
+ public int advance(int target) throws IOException {
+ return slowAdvance(target);
+ }
+
+ @Override
+ public long cost() {
+ return sortedDocIds.length;
+ }
+
+ @Override
+ public int docID() {
+ if (idx == -1) {
+ return -1;
+ } else if (idx >= sortedDocIds.length) {
+ return DocIdSetIterator.NO_MORE_DOCS;
+ } else {
+ return sortedDocIds[idx];
+ }
+ }
+
+ @Override
+ public int nextDoc() {
+ idx += 1;
+ return docID();
+ }
+ }
+
+ class SeededCollectorManager implements KnnCollectorManager {
+ final KnnCollectorManager knnCollectorManager;
+
+ SeededCollectorManager(KnnCollectorManager knnCollectorManager) {
+ this.knnCollectorManager = knnCollectorManager;
+ }
+
+ @Override
+ public KnnCollector newCollector(
+ int visitLimit, KnnSearchStrategy searchStrategy, LeafReaderContext ctx)
+ throws IOException {
+ TopScoreDocCollector seedCollector =
+ new TopScoreDocCollectorManager(k, null, Integer.MAX_VALUE).newCollector();
+ final LeafReader leafReader = ctx.reader();
+ final LeafCollector leafCollector = seedCollector.getLeafCollector(ctx);
+ if (leafCollector != null) {
+ try {
+ BulkScorer scorer = seedWeight.bulkScorer(ctx);
+ if (scorer != null) {
+ scorer.score(
+ leafCollector,
+ leafReader.getLiveDocs(),
+ 0 /* min */,
+ DocIdSetIterator.NO_MORE_DOCS /* max */);
+ }
+ } catch (
+ @SuppressWarnings("unused")
+ CollectionTerminatedException e) {
+ }
+ leafCollector.finish();
+ }
+ KnnCollector delegateCollector =
+ knnCollectorManager.newCollector(visitLimit, searchStrategy, ctx);
+ TopDocs seedTopDocs = seedCollector.topDocs();
+ VectorScorer scorer =
+ delegate.createVectorScorer(ctx, leafReader.getFieldInfos().fieldInfo(field));
+ if (seedTopDocs.totalHits.value() == 0 || scorer == null) {
+ return delegateCollector;
+ }
+ DocIdSetIterator vectorIterator = scorer.iterator();
+ // Handle sparse
+ if (vectorIterator instanceof IndexedDISI indexedDISI) {
+ vectorIterator = IndexedDISI.asDocIndexIterator(indexedDISI);
+ }
+ // Most underlying iterators are indexed, so we can map the seed docs to the vector docs
+ if (vectorIterator instanceof KnnVectorValues.DocIndexIterator indexIterator) {
+ DocIdSetIterator seedDocs = new MappedDISI(indexIterator, new TopDocsDISI(seedTopDocs));
+ return knnCollectorManager.newCollector(
+ visitLimit, new KnnSearchStrategy.Seeded(seedDocs, seedTopDocs.scoreDocs.length), ctx);
+ }
+ return delegateCollector;
+ }
+ }
+}
diff --git a/lucene/core/src/java/org/apache/lucene/search/TimeLimitingKnnCollectorManager.java b/lucene/core/src/java/org/apache/lucene/search/TimeLimitingKnnCollectorManager.java
index 2dc2f035b90..08cc317fae1 100644
--- a/lucene/core/src/java/org/apache/lucene/search/TimeLimitingKnnCollectorManager.java
+++ b/lucene/core/src/java/org/apache/lucene/search/TimeLimitingKnnCollectorManager.java
@@ -20,6 +20,7 @@ import java.io.IOException;
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;
/** A {@link KnnCollectorManager} that collects results with a timeout. */
public class TimeLimitingKnnCollectorManager implements KnnCollectorManager {
@@ -37,8 +38,10 @@ public class TimeLimitingKnnCollectorManager implements KnnCollectorManager {
}
@Override
- public KnnCollector newCollector(int visitedLimit, LeafReaderContext context) throws IOException {
- KnnCollector collector = delegate.newCollector(visitedLimit, context);
+ public KnnCollector newCollector(
+ int visitedLimit, KnnSearchStrategy searchStrategy, LeafReaderContext context)
+ throws IOException {
+ KnnCollector collector = delegate.newCollector(visitedLimit, searchStrategy, context);
if (queryTimeout == null) {
return collector;
}
diff --git a/lucene/core/src/java/org/apache/lucene/search/TopKnnCollector.java b/lucene/core/src/java/org/apache/lucene/search/TopKnnCollector.java
index cef8eb18024..488215b945d 100644
--- a/lucene/core/src/java/org/apache/lucene/search/TopKnnCollector.java
+++ b/lucene/core/src/java/org/apache/lucene/search/TopKnnCollector.java
@@ -17,6 +17,7 @@
package org.apache.lucene.search;
+import org.apache.lucene.search.knn.KnnSearchStrategy;
import org.apache.lucene.util.hnsw.NeighborQueue;
/**
@@ -29,12 +30,18 @@ public class TopKnnCollector extends AbstractKnnCollector {
protected final NeighborQueue queue;
+ /** Constructor, with default search strategy */
+ public TopKnnCollector(int k, int visitLimit) {
+ this(k, visitLimit, null);
+ }
+
/**
* @param k the number of neighbors to collect
* @param visitLimit how many vector nodes the results are allowed to visit
+ * @param searchStrategy the search strategy to use
*/
- public TopKnnCollector(int k, int visitLimit) {
- super(k, visitLimit);
+ public TopKnnCollector(int k, int visitLimit, KnnSearchStrategy searchStrategy) {
+ super(k, visitLimit, searchStrategy);
this.queue = new NeighborQueue(k, false);
}
diff --git a/lucene/core/src/java/org/apache/lucene/search/knn/EntryPointProvider.java b/lucene/core/src/java/org/apache/lucene/search/knn/EntryPointProvider.java
deleted file mode 100644
index 9e7b44b571d..00000000000
--- a/lucene/core/src/java/org/apache/lucene/search/knn/EntryPointProvider.java
+++ /dev/null
@@ -1,28 +0,0 @@
-/*
- * 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.knn;
-
-import org.apache.lucene.search.DocIdSetIterator;
-
-/** Provides entry points for the kNN search */
-public interface EntryPointProvider {
- /** Iterator of valid entry points for the kNN search */
- DocIdSetIterator entryPoints();
-
- /** Number of valid entry points for the kNN search */
- int numberOfEntryPoints();
-}
diff --git a/lucene/core/src/java/org/apache/lucene/search/knn/KnnCollectorManager.java b/lucene/core/src/java/org/apache/lucene/search/knn/KnnCollectorManager.java
index 5e1eeecc840..6e3219f3b66 100644
--- a/lucene/core/src/java/org/apache/lucene/search/knn/KnnCollectorManager.java
+++ b/lucene/core/src/java/org/apache/lucene/search/knn/KnnCollectorManager.java
@@ -32,7 +32,10 @@ public interface KnnCollectorManager {
* Return a new {@link KnnCollector} instance.
*
* @param visitedLimit the maximum number of nodes that the search is allowed to visit
+ * @param searchStrategy the optional search strategy configuration
* @param context the leaf reader context
*/
- KnnCollector newCollector(int visitedLimit, LeafReaderContext context) throws IOException;
+ KnnCollector newCollector(
+ int visitedLimit, KnnSearchStrategy searchStrategy, LeafReaderContext context)
+ throws IOException;
}
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
new file mode 100644
index 00000000000..a1d5a7b1b57
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/search/knn/KnnSearchStrategy.java
@@ -0,0 +1,63 @@
+/*
+ * 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.knn;
+
+import org.apache.lucene.search.DocIdSetIterator;
+
+/**
+ * KnnSearchStrategy is a strategy for kNN search, providing additional search strategy
+ * configuration
+ *
+ * @lucene.experimental
+ */
+public interface KnnSearchStrategy {
+
+ /** A strategy for kNN search that uses a set of entry points to start the search */
+ class Seeded implements KnnSearchStrategy {
+ private final DocIdSetIterator entryPoints;
+ private final int numberOfEntryPoints;
+
+ public Seeded(DocIdSetIterator entryPoints, int numberOfEntryPoints) {
+ if (numberOfEntryPoints < 0) {
+ throw new IllegalArgumentException("numberOfEntryPoints must be >= 0");
+ }
+ this.numberOfEntryPoints = numberOfEntryPoints;
+ if (numberOfEntryPoints > 0 && entryPoints == null) {
+ throw new IllegalArgumentException("entryPoints must not be null");
+ }
+ this.entryPoints = entryPoints == null ? DocIdSetIterator.empty() : entryPoints;
+ }
+
+ /**
+ * Iterator of valid entry points for the kNN search
+ *
+ * @return DocIdSetIterator of entry points
+ */
+ public DocIdSetIterator entryPoints() {
+ return entryPoints;
+ }
+
+ /**
+ * Number of valid entry points for the kNN search
+ *
+ * @return number of entry points
+ */
+ public int numberOfEntryPoints() {
+ return numberOfEntryPoints;
+ }
+ }
+}
diff --git a/lucene/core/src/java/org/apache/lucene/search/knn/SeededKnnCollector.java b/lucene/core/src/java/org/apache/lucene/search/knn/SeededKnnCollector.java
deleted file mode 100644
index c3c4f62901e..00000000000
--- a/lucene/core/src/java/org/apache/lucene/search/knn/SeededKnnCollector.java
+++ /dev/null
@@ -1,48 +0,0 @@
-/*
- * 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.knn;
-
-import org.apache.lucene.search.DocIdSetIterator;
-import org.apache.lucene.search.KnnCollector;
-
-/**
- * A {@link KnnCollector} that provides seeded knn collection. See usage in {@link
- * SeededKnnCollectorManager}.
- *
- * @lucene.experimental
- */
-class SeededKnnCollector extends KnnCollector.Decorator implements EntryPointProvider {
- private final DocIdSetIterator entryPoints;
- private final int numberOfEntryPoints;
-
- SeededKnnCollector(
- KnnCollector collector, DocIdSetIterator entryPoints, int numberOfEntryPoints) {
- super(collector);
- this.entryPoints = entryPoints;
- this.numberOfEntryPoints = numberOfEntryPoints;
- }
-
- @Override
- public DocIdSetIterator entryPoints() {
- return entryPoints;
- }
-
- @Override
- public int numberOfEntryPoints() {
- return numberOfEntryPoints;
- }
-}
diff --git a/lucene/core/src/java/org/apache/lucene/search/knn/SeededKnnCollectorManager.java b/lucene/core/src/java/org/apache/lucene/search/knn/SeededKnnCollectorManager.java
deleted file mode 100644
index 7631db6e302..00000000000
--- a/lucene/core/src/java/org/apache/lucene/search/knn/SeededKnnCollectorManager.java
+++ /dev/null
@@ -1,177 +0,0 @@
-/*
- * 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.knn;
-
-import java.io.IOException;
-import java.util.Arrays;
-import org.apache.lucene.index.KnnVectorValues;
-import org.apache.lucene.index.LeafReader;
-import org.apache.lucene.index.LeafReaderContext;
-import org.apache.lucene.search.BulkScorer;
-import org.apache.lucene.search.CollectionTerminatedException;
-import org.apache.lucene.search.DocIdSetIterator;
-import org.apache.lucene.search.KnnCollector;
-import org.apache.lucene.search.LeafCollector;
-import org.apache.lucene.search.TopDocs;
-import org.apache.lucene.search.TopScoreDocCollector;
-import org.apache.lucene.search.TopScoreDocCollectorManager;
-import org.apache.lucene.search.Weight;
-import org.apache.lucene.util.IOFunction;
-
-/**
- * A {@link KnnCollectorManager} that provides seeded knn collection. See usage in {@link
- * org.apache.lucene.search.SeededKnnFloatVectorQuery} and {@link
- * org.apache.lucene.search.SeededKnnByteVectorQuery}.
- */
-public class SeededKnnCollectorManager implements KnnCollectorManager {
- private final KnnCollectorManager delegate;
- private final Weight seedWeight;
- private final int k;
- private final IOFunction<LeafReader, KnnVectorValues> vectorValuesSupplier;
-
- public SeededKnnCollectorManager(
- KnnCollectorManager delegate,
- Weight seedWeight,
- int k,
- IOFunction<LeafReader, KnnVectorValues> vectorValuesSupplier) {
- this.delegate = delegate;
- this.seedWeight = seedWeight;
- this.k = k;
- this.vectorValuesSupplier = vectorValuesSupplier;
- }
-
- @Override
- public KnnCollector newCollector(int visitedLimit, LeafReaderContext ctx) throws IOException {
- // Execute the seed query
- TopScoreDocCollector seedCollector =
- new TopScoreDocCollectorManager(k, null, Integer.MAX_VALUE).newCollector();
- final LeafReader leafReader = ctx.reader();
- final LeafCollector leafCollector = seedCollector.getLeafCollector(ctx);
- if (leafCollector != null) {
- try {
- BulkScorer scorer = seedWeight.bulkScorer(ctx);
- if (scorer != null) {
- scorer.score(
- leafCollector,
- leafReader.getLiveDocs(),
- 0 /* min */,
- DocIdSetIterator.NO_MORE_DOCS /* max */);
- }
- } catch (
- @SuppressWarnings("unused")
- CollectionTerminatedException e) {
- }
- leafCollector.finish();
- }
-
- TopDocs seedTopDocs = seedCollector.topDocs();
- KnnVectorValues vectorValues = vectorValuesSupplier.apply(leafReader);
- final KnnCollector delegateCollector = delegate.newCollector(visitedLimit, ctx);
- if (seedTopDocs.totalHits.value() == 0 || vectorValues == null) {
- return delegateCollector;
- }
- KnnVectorValues.DocIndexIterator indexIterator = vectorValues.iterator();
- DocIdSetIterator seedDocs = new MappedDISI(indexIterator, new TopDocsDISI(seedTopDocs));
- return new SeededKnnCollector(delegateCollector, seedDocs, seedTopDocs.scoreDocs.length);
- }
-
- private static class MappedDISI extends DocIdSetIterator {
- KnnVectorValues.DocIndexIterator indexedDISI;
- DocIdSetIterator sourceDISI;
-
- private MappedDISI(KnnVectorValues.DocIndexIterator indexedDISI, DocIdSetIterator sourceDISI) {
- this.indexedDISI = indexedDISI;
- this.sourceDISI = sourceDISI;
- }
-
- /**
- * Advances the source iterator to the first document number that is greater than or equal to
- * the provided target and returns the corresponding index.
- */
- @Override
- public int advance(int target) throws IOException {
- int newTarget = sourceDISI.advance(target);
- if (newTarget != NO_MORE_DOCS) {
- indexedDISI.advance(newTarget);
- }
- return docID();
- }
-
- @Override
- public long cost() {
- return sourceDISI.cost();
- }
-
- @Override
- public int docID() {
- if (indexedDISI.docID() == NO_MORE_DOCS || sourceDISI.docID() == NO_MORE_DOCS) {
- return NO_MORE_DOCS;
- }
- return indexedDISI.index();
- }
-
- /** Advances to the next document in the source iterator and returns the corresponding index. */
- @Override
- public int nextDoc() throws IOException {
- int newTarget = sourceDISI.nextDoc();
- if (newTarget != NO_MORE_DOCS) {
- indexedDISI.advance(newTarget);
- }
- return docID();
- }
- }
-
- private static class TopDocsDISI extends DocIdSetIterator {
- private final int[] sortedDocIds;
- private int idx = -1;
-
- private TopDocsDISI(TopDocs topDocs) {
- sortedDocIds = new int[topDocs.scoreDocs.length];
- for (int i = 0; i < topDocs.scoreDocs.length; i++) {
- sortedDocIds[i] = topDocs.scoreDocs[i].doc;
- }
- Arrays.sort(sortedDocIds);
- }
-
- @Override
- public int advance(int target) throws IOException {
- return slowAdvance(target);
- }
-
- @Override
- public long cost() {
- return sortedDocIds.length;
- }
-
- @Override
- public int docID() {
- if (idx == -1) {
- return -1;
- } else if (idx >= sortedDocIds.length) {
- return DocIdSetIterator.NO_MORE_DOCS;
- } else {
- return sortedDocIds[idx];
- }
- }
-
- @Override
- public int nextDoc() {
- idx += 1;
- return docID();
- }
- }
-}
diff --git a/lucene/core/src/java/org/apache/lucene/search/knn/TopKnnCollectorManager.java b/lucene/core/src/java/org/apache/lucene/search/knn/TopKnnCollectorManager.java
index a522d568078..cd3d1cd8e64 100644
--- a/lucene/core/src/java/org/apache/lucene/search/knn/TopKnnCollectorManager.java
+++ b/lucene/core/src/java/org/apache/lucene/search/knn/TopKnnCollectorManager.java
@@ -49,11 +49,14 @@ public class TopKnnCollectorManager implements KnnCollectorManager {
* @param context the leaf reader context
*/
@Override
- public KnnCollector newCollector(int visitedLimit, LeafReaderContext context) throws IOException {
+ public KnnCollector newCollector(
+ int visitedLimit, KnnSearchStrategy searchStrategy, LeafReaderContext context)
+ throws IOException {
if (globalScoreQueue == null) {
- return new TopKnnCollector(k, visitedLimit);
+ return new TopKnnCollector(k, visitedLimit, searchStrategy);
} else {
- return new MultiLeafKnnCollector(k, globalScoreQueue, new TopKnnCollector(k, visitedLimit));
+ return new MultiLeafKnnCollector(
+ k, globalScoreQueue, new TopKnnCollector(k, visitedLimit, searchStrategy));
}
}
}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/AbstractHnswGraphSearcher.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/AbstractHnswGraphSearcher.java
new file mode 100644
index 00000000000..ab602e45d64
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/AbstractHnswGraphSearcher.java
@@ -0,0 +1,84 @@
+/*
+ * 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 java.io.IOException;
+import org.apache.lucene.search.KnnCollector;
+import org.apache.lucene.util.Bits;
+
+/**
+ * AbstractHnswGraphSearcher is the base class for HnswGraphSearcher implementations.
+ *
+ * @lucene.experimental
+ */
+abstract class AbstractHnswGraphSearcher {
+
+ static final int UNK_EP = -1;
+
+ /**
+ * Search a given level of the graph starting at the given entry points.
+ *
+ * @param results the collector to collect the results
+ * @param scorer the scorer to compare the query with the nodes
+ * @param level the level of the graph to search
+ * @param eps the entry points to start the search from
+ * @param graph the HNSWGraph
+ * @param acceptOrds the ordinals to accept for the results
+ */
+ abstract void searchLevel(
+ KnnCollector results,
+ RandomVectorScorer scorer,
+ int level,
+ final int[] eps,
+ HnswGraph graph,
+ Bits acceptOrds)
+ throws IOException;
+
+ /**
+ * Function to find the best entry point from which to search the zeroth graph layer.
+ *
+ * @param scorer the scorer to compare the query with the nodes
+ * @param graph the HNSWGraph
+ * @param collector the knn result collector
+ * @return the best entry point, `-1` indicates graph entry node not set, or visitation limit
+ * exceeded
+ * @throws IOException When accessing the vectors or graph fails
+ */
+ abstract int[] findBestEntryPoint(
+ RandomVectorScorer scorer, HnswGraph graph, KnnCollector collector) throws IOException;
+
+ /**
+ * Search the graph for the given scorer. Gathering results in the provided collector that pass
+ * the provided acceptOrds.
+ *
+ * @param results the collector to collect the results
+ * @param scorer the scorer to compare the query with the nodes
+ * @param graph the HNSWGraph
+ * @param acceptOrds the ordinals to accept for the results
+ * @throws IOException When accessing the vectors or graph fails
+ */
+ public void search(
+ KnnCollector results, RandomVectorScorer scorer, HnswGraph graph, Bits acceptOrds)
+ throws IOException {
+ int[] eps = findBestEntryPoint(scorer, graph, results);
+ assert eps != null && eps.length > 0;
+ if (eps[0] == UNK_EP) {
+ return;
+ }
+ searchLevel(results, scorer, 0, eps, graph, acceptOrds);
+ }
+}
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 e8f0d316fd8..a414bb9dfe0 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
@@ -20,10 +20,9 @@ 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.DocIdSetIterator;
import org.apache.lucene.search.KnnCollector;
import org.apache.lucene.search.TopKnnCollector;
-import org.apache.lucene.search.knn.EntryPointProvider;
+import org.apache.lucene.search.knn.KnnSearchStrategy;
import org.apache.lucene.util.BitSet;
import org.apache.lucene.util.Bits;
import org.apache.lucene.util.FixedBitSet;
@@ -33,7 +32,7 @@ import org.apache.lucene.util.SparseFixedBitSet;
* Searches an HNSW graph to find nearest neighbors to a query vector. For more background on the
* search algorithm, see {@link HnswGraph}.
*/
-public class HnswGraphSearcher {
+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.
@@ -59,7 +58,7 @@ public class HnswGraphSearcher {
* Otherwise, the search will discover the best entry point per the normal HNSW search algorithm.
*
* @param scorer the scorer to compare the query with the nodes
- * @param knnCollector a collector of top knn results to be returned
+ * @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
@@ -68,33 +67,16 @@ public class HnswGraphSearcher {
public static void search(
RandomVectorScorer scorer, KnnCollector knnCollector, HnswGraph graph, Bits acceptOrds)
throws IOException {
- HnswGraphSearcher graphSearcher =
+ AbstractHnswGraphSearcher graphSearcher =
new HnswGraphSearcher(
new NeighborQueue(knnCollector.k(), true), new SparseFixedBitSet(getGraphSize(graph)));
- final int[] entryPoints;
- if (knnCollector instanceof EntryPointProvider epp) {
- if (epp.numberOfEntryPoints() <= 0) {
- throw new IllegalArgumentException("The number of entry points must be > 0");
- }
- DocIdSetIterator eps = epp.entryPoints();
- entryPoints = new int[epp.numberOfEntryPoints()];
- int idx = 0;
- while (idx < entryPoints.length) {
- int entryPointOrdInt = eps.nextDoc();
- if (entryPointOrdInt == NO_MORE_DOCS) {
- throw new IllegalArgumentException(
- "The number of entry points provided is less than the number of entry points requested");
- }
- assert entryPointOrdInt < getGraphSize(graph);
- entryPoints[idx++] = entryPointOrdInt;
- }
- // This is an invalid case, but we should check it
- assert entryPoints.length > 0;
- // We use provided entry point ordinals to search the complete graph (level 0)
- graphSearcher.searchLevel(knnCollector, scorer, 0, entryPoints, graph, acceptOrds);
- } else {
- search(scorer, knnCollector, graph, graphSearcher, acceptOrds);
+ if (knnCollector.getSearchStrategy() instanceof KnnSearchStrategy.Seeded seeded
+ && seeded.numberOfEntryPoints() > 0) {
+ graphSearcher =
+ SeededHnswGraphSearcher.fromEntryPoints(
+ graphSearcher, seeded.numberOfEntryPoints(), seeded.entryPoints(), graph.size());
}
+ graphSearcher.search(knnCollector, scorer, graph, acceptOrds);
}
/**
@@ -112,27 +94,14 @@ public class HnswGraphSearcher {
public static KnnCollector search(
RandomVectorScorer scorer, int topK, OnHeapHnswGraph graph, Bits acceptOrds, int visitedLimit)
throws IOException {
- KnnCollector knnCollector = new TopKnnCollector(topK, visitedLimit);
+ KnnCollector knnCollector = new TopKnnCollector(topK, visitedLimit, null);
OnHeapHnswGraphSearcher graphSearcher =
new OnHeapHnswGraphSearcher(
new NeighborQueue(topK, true), new SparseFixedBitSet(getGraphSize(graph)));
- search(scorer, knnCollector, graph, graphSearcher, acceptOrds);
+ graphSearcher.search(knnCollector, scorer, graph, acceptOrds);
return knnCollector;
}
- private static void search(
- RandomVectorScorer scorer,
- KnnCollector knnCollector,
- HnswGraph graph,
- HnswGraphSearcher graphSearcher,
- Bits acceptOrds)
- throws IOException {
- int ep = graphSearcher.findBestEntryPoint(scorer, graph, knnCollector);
- if (ep != -1) {
- graphSearcher.searchLevel(knnCollector, scorer, 0, new int[] {ep}, graph, acceptOrds);
- }
- }
-
/**
* Searches for the nearest neighbors of a query vector in a given level.
*
@@ -166,11 +135,12 @@ public class HnswGraphSearcher {
* exceeded
* @throws IOException When accessing the vector fails
*/
- private int findBestEntryPoint(RandomVectorScorer scorer, HnswGraph graph, KnnCollector collector)
+ @Override
+ int[] findBestEntryPoint(RandomVectorScorer scorer, HnswGraph graph, KnnCollector collector)
throws IOException {
int currentEp = graph.entryNode();
if (currentEp == -1 || graph.numLevels() == 1) {
- return currentEp;
+ return new int[] {currentEp};
}
int size = getGraphSize(graph);
prepareScratchState(size);
@@ -191,7 +161,7 @@ public class HnswGraphSearcher {
continue;
}
if (collector.earlyTerminated()) {
- return -1;
+ return new int[] {UNK_EP};
}
float friendSimilarity = scorer.score(friendOrd);
collector.incVisitedCount(1);
@@ -203,7 +173,7 @@ public class HnswGraphSearcher {
}
}
}
- return collector.earlyTerminated() ? -1 : currentEp;
+ return collector.earlyTerminated() ? new int[] {UNK_EP} : new int[] {currentEp};
}
/**
@@ -212,6 +182,7 @@ public class HnswGraphSearcher {
* score/comparison value, will be at the top of the heap, while the closest neighbor will be the
* last to be popped.
*/
+ @Override
void searchLevel(
KnnCollector results,
RandomVectorScorer scorer,
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/SeededHnswGraphSearcher.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/SeededHnswGraphSearcher.java
new file mode 100644
index 00000000000..50c102602c2
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/SeededHnswGraphSearcher.java
@@ -0,0 +1,78 @@
+/*
+ * 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.DocIdSetIterator;
+import org.apache.lucene.search.KnnCollector;
+import org.apache.lucene.util.Bits;
+
+/**
+ * A {@link HnswGraphSearcher} that uses a set of seed ordinals to initiate the search.
+ *
+ * @lucene.experimental
+ */
+final class SeededHnswGraphSearcher extends AbstractHnswGraphSearcher {
+
+ private final AbstractHnswGraphSearcher delegate;
+ private final int[] seedOrds;
+
+ static SeededHnswGraphSearcher fromEntryPoints(
+ AbstractHnswGraphSearcher delegate, int numEps, DocIdSetIterator eps, int graphSize)
+ throws IOException {
+ if (numEps <= 0) {
+ throw new IllegalArgumentException("The number of entry points must be > 0");
+ }
+ int[] entryPoints = new int[numEps];
+ int idx = 0;
+ while (idx < entryPoints.length) {
+ int entryPointOrdInt = eps.nextDoc();
+ if (entryPointOrdInt == NO_MORE_DOCS) {
+ throw new IllegalArgumentException(
+ "The number of entry points provided is less than the number of entry points requested");
+ }
+ assert entryPointOrdInt < graphSize;
+ entryPoints[idx++] = entryPointOrdInt;
+ }
+ return new SeededHnswGraphSearcher(delegate, entryPoints);
+ }
+
+ SeededHnswGraphSearcher(AbstractHnswGraphSearcher delegate, int[] seedOrds) {
+ this.delegate = delegate;
+ this.seedOrds = seedOrds;
+ }
+
+ @Override
+ void searchLevel(
+ KnnCollector results,
+ RandomVectorScorer scorer,
+ int level,
+ int[] eps,
+ HnswGraph graph,
+ Bits acceptOrds)
+ throws IOException {
+ delegate.searchLevel(results, scorer, level, eps, graph, acceptOrds);
+ }
+
+ @Override
+ int[] findBestEntryPoint(RandomVectorScorer scorer, HnswGraph graph, KnnCollector collector) {
+ return seedOrds;
+ }
+}
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 efe9500635a..1e8caeff49f 100644
--- a/lucene/core/src/test/org/apache/lucene/search/BaseKnnVectorQueryTestCase.java
+++ b/lucene/core/src/test/org/apache/lucene/search/BaseKnnVectorQueryTestCase.java
@@ -887,7 +887,7 @@ abstract class BaseKnnVectorQueryTestCase extends LuceneTestCase {
TimeLimitingKnnCollectorManager noTimeoutManager =
new TimeLimitingKnnCollectorManager(delegate, null);
KnnCollector noTimeoutCollector =
- noTimeoutManager.newCollector(Integer.MAX_VALUE, searcher.leafContexts.get(0));
+ noTimeoutManager.newCollector(Integer.MAX_VALUE, null, searcher.leafContexts.get(0));
// Check that a normal collector is created without timeout
assertFalse(
@@ -904,7 +904,7 @@ abstract class BaseKnnVectorQueryTestCase extends LuceneTestCase {
TimeLimitingKnnCollectorManager timeoutManager =
new TimeLimitingKnnCollectorManager(delegate, () -> true);
KnnCollector timeoutCollector =
- timeoutManager.newCollector(Integer.MAX_VALUE, searcher.leafContexts.get(0));
+ timeoutManager.newCollector(Integer.MAX_VALUE, null, searcher.leafContexts.get(0));
// Check that a time limiting collector is created, which returns partial results
assertFalse(timeoutCollector instanceof TopKnnCollector);
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 d0fb8c95e03..49ed29fcc08 100644
--- a/lucene/core/src/test/org/apache/lucene/search/TestSeededKnnByteVectorQuery.java
+++ b/lucene/core/src/test/org/apache/lucene/search/TestSeededKnnByteVectorQuery.java
@@ -19,6 +19,7 @@ package org.apache.lucene.search;
import static org.apache.lucene.search.TestKnnByteVectorQuery.floatToBytes;
import java.io.IOException;
+import java.util.concurrent.atomic.AtomicInteger;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.document.IntPoint;
@@ -28,11 +29,13 @@ import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.index.LeafReaderContext;
-import org.apache.lucene.index.QueryTimeout;
import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.search.knn.KnnCollectorManager;
+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.TestUtil;
+import org.apache.lucene.util.Bits;
import org.apache.lucene.util.TestVectorUtil;
public class TestSeededKnnByteVectorQuery extends BaseKnnVectorQueryTestCase {
@@ -41,12 +44,16 @@ public class TestSeededKnnByteVectorQuery extends BaseKnnVectorQueryTestCase {
@Override
AbstractKnnVectorQuery getKnnVectorQuery(String field, float[] query, int k, Query queryFilter) {
- return new SeededKnnByteVectorQuery(field, floatToBytes(query), k, queryFilter, MATCH_NONE);
+ KnnByteVectorQuery knnByteVectorQuery =
+ new KnnByteVectorQuery(field, floatToBytes(query), k, queryFilter);
+ return SeededKnnVectorQuery.fromByteQuery(knnByteVectorQuery, MATCH_NONE);
}
@Override
AbstractKnnVectorQuery getThrowingKnnVectorQuery(String field, float[] vec, int k, Query query) {
- return new ThrowingKnnVectorQuery(field, floatToBytes(vec), k, query, MATCH_NONE);
+ KnnByteVectorQuery knnByteVectorQuery =
+ new TestKnnByteVectorQuery.ThrowingKnnVectorQuery(field, floatToBytes(vec), k, query);
+ return SeededKnnVectorQuery.fromByteQuery(knnByteVectorQuery, MATCH_NONE);
}
@Override
@@ -71,6 +78,46 @@ public class TestSeededKnnByteVectorQuery extends BaseKnnVectorQueryTestCase {
return new KnnByteVectorField(name, floatToBytes(vector), VectorSimilarityFunction.EUCLIDEAN);
}
+ public void testSeedWithTimeout() throws IOException {
+ int numDocs = atLeast(50);
+ int dimension = atLeast(5);
+ int numIters = atLeast(5);
+ try (Directory d = newDirectoryForTest()) {
+ IndexWriterConfig iwc = new IndexWriterConfig().setCodec(TestUtil.getDefaultCodec());
+ RandomIndexWriter w = new RandomIndexWriter(random(), d, iwc);
+ for (int i = 0; i < numDocs; i++) {
+ Document doc = new Document();
+ doc.add(getKnnVectorField("field", randomVector(dimension)));
+ doc.add(new NumericDocValuesField("tag", i));
+ doc.add(new IntPoint("tag", i));
+ w.addDocument(doc);
+ }
+ w.close();
+
+ try (IndexReader reader = DirectoryReader.open(d)) {
+ IndexSearcher searcher = newSearcher(reader);
+ searcher.setTimeout(() -> true);
+ int k = random().nextInt(80) + 1;
+ for (int i = 0; i < numIters; i++) {
+ // All documents as seeds
+ Query seed =
+ random().nextBoolean()
+ ? IntPoint.newRangeQuery("tag", 1, 6)
+ : new MatchAllDocsQuery();
+ Query filter = random().nextBoolean() ? null : new MatchAllDocsQuery();
+ KnnByteVectorQuery byteVectorQuery =
+ new KnnByteVectorQuery("field", floatToBytes(randomVector(dimension)), k, filter);
+ Query knnQuery = SeededKnnVectorQuery.fromByteQuery(byteVectorQuery, seed);
+ assertEquals(0, searcher.count(knnQuery));
+ // No seed documents -- falls back on full approx search
+ seed = new MatchNoDocsQuery();
+ knnQuery = SeededKnnVectorQuery.fromByteQuery(byteVectorQuery, seed);
+ assertEquals(0, searcher.count(knnQuery));
+ }
+ }
+ }
+ }
+
/** Tests with random vectors and a random seed. Uses RandomIndexWriter. */
public void testRandomWithSeed() throws IOException {
int numDocs = 1000;
@@ -101,6 +148,12 @@ public class TestSeededKnnByteVectorQuery extends BaseKnnVectorQueryTestCase {
try (IndexReader reader = DirectoryReader.open(d)) {
IndexSearcher searcher = newSearcher(reader);
for (int i = 0; i < numIters; i++) {
+ // verify timeout collector wrapping is used
+ if (random().nextBoolean()) {
+ searcher.setTimeout(() -> false);
+ } else {
+ searcher.setTimeout(null);
+ }
int k = random().nextInt(80) + 1;
int n = random().nextInt(100) + 1;
// we may get fewer results than requested if there are deletions, but this test doesn't
@@ -108,12 +161,15 @@ public class TestSeededKnnByteVectorQuery extends BaseKnnVectorQueryTestCase {
assert reader.hasDeletions() == false;
// All documents as seeds
+ AtomicInteger seedCalls = new AtomicInteger();
Query seed1 = new MatchAllDocsQuery();
Query filter = random().nextBoolean() ? null : new MatchAllDocsQuery();
- SeededKnnByteVectorQuery query =
- new SeededKnnByteVectorQuery(
- "field", floatToBytes(randomVector(dimension)), k, filter, seed1);
+ KnnByteVectorQuery byteVectorQuery =
+ new KnnByteVectorQuery("field", floatToBytes(randomVector(dimension)), k, filter);
+ AssertingSeededKnnVectorQuery query =
+ new AssertingSeededKnnVectorQuery(byteVectorQuery, seed1, null, seedCalls);
TopDocs results = searcher.search(query, n);
+ assertEquals(seedCalls.get(), 1);
int expected = Math.min(Math.min(n, k), numDocsWithVector);
assertEquals(expected, results.scoreDocs.length);
@@ -127,10 +183,19 @@ public class TestSeededKnnByteVectorQuery extends BaseKnnVectorQueryTestCase {
// Restrictive seed query -- 6 documents
Query seed2 = IntPoint.newRangeQuery("tag", 1, 6);
+ int seedCount =
+ searcher.count(
+ new BooleanQuery.Builder()
+ .add(seed2, BooleanClause.Occur.MUST)
+ .add(new FieldExistsQuery("field"), BooleanClause.Occur.MUST)
+ .build());
+ byteVectorQuery =
+ new KnnByteVectorQuery("field", floatToBytes(randomVector(dimension)), k, null);
query =
- new SeededKnnByteVectorQuery(
- "field", floatToBytes(randomVector(dimension)), k, null, seed2);
+ new AssertingSeededKnnVectorQuery(
+ byteVectorQuery, seed2, null, seedCount > 0 ? seedCalls : null);
results = searcher.search(query, n);
+ assertEquals(seedCalls.get(), seedCount > 0 ? 2 : 1);
expected = Math.min(Math.min(n, k), reader.numDocs());
assertEquals(expected, results.scoreDocs.length);
assertTrue(results.totalHits.value() >= results.scoreDocs.length);
@@ -143,9 +208,7 @@ public class TestSeededKnnByteVectorQuery extends BaseKnnVectorQueryTestCase {
// No seed documents -- falls back on full approx search
Query seed3 = new MatchNoDocsQuery();
- query =
- new SeededKnnByteVectorQuery(
- "field", floatToBytes(randomVector(dimension)), k, null, seed3);
+ query = new AssertingSeededKnnVectorQuery(byteVectorQuery, seed3, null, null);
results = searcher.search(query, n);
expected = Math.min(Math.min(n, k), reader.numDocs());
assertEquals(expected, results.scoreDocs.length);
@@ -161,45 +224,100 @@ public class TestSeededKnnByteVectorQuery extends BaseKnnVectorQueryTestCase {
}
}
- private static class ThrowingKnnVectorQuery extends SeededKnnByteVectorQuery {
+ static class AssertingSeededKnnVectorQuery extends SeededKnnVectorQuery {
+ private final AtomicInteger seedCalls;
- public ThrowingKnnVectorQuery(String field, byte[] target, int k, Query filter, Query seed) {
- super(field, target, k, filter, seed);
- }
-
- private ThrowingKnnVectorQuery(
- String field, byte[] target, int k, Query filter, Weight seedWeight) {
- super(field, target, k, filter, seedWeight);
+ public AssertingSeededKnnVectorQuery(
+ AbstractKnnVectorQuery query, Query seed, Weight seedWeight, AtomicInteger seedCalls) {
+ super(query, seed, seedWeight);
+ this.seedCalls = seedCalls;
}
@Override
- // This is test only and we need to overwrite the inner rewrite to throw
public Query rewrite(IndexSearcher indexSearcher) throws IOException {
if (seedWeight != null) {
return super.rewrite(indexSearcher);
}
- BooleanQuery.Builder booleanSeedQueryBuilder =
- new BooleanQuery.Builder()
- .add(seed, BooleanClause.Occur.MUST)
- .add(new FieldExistsQuery(field), BooleanClause.Occur.FILTER);
- if (filter != null) {
- booleanSeedQueryBuilder.add(filter, BooleanClause.Occur.FILTER);
- }
- Query seedRewritten = indexSearcher.rewrite(booleanSeedQueryBuilder.build());
- Weight seedWeight = indexSearcher.createWeight(seedRewritten, ScoreMode.TOP_SCORES, 1f);
- return new ThrowingKnnVectorQuery(field, target, k, filter, seedWeight)
- .rewrite(indexSearcher);
+ AssertingSeededKnnVectorQuery rewritten =
+ new AssertingSeededKnnVectorQuery(
+ delegate, seed, createSeedWeight(indexSearcher), seedCalls);
+ return rewritten.rewrite(indexSearcher);
}
@Override
- protected TopDocs exactSearch(
- LeafReaderContext context, DocIdSetIterator acceptIterator, QueryTimeout queryTimeout) {
- throw new UnsupportedOperationException("exact search is not supported");
+ protected TopDocs approximateSearch(
+ LeafReaderContext context,
+ Bits acceptDocs,
+ int visitedLimit,
+ KnnCollectorManager knnCollectorManager)
+ throws IOException {
+ return delegate.approximateSearch(
+ context,
+ acceptDocs,
+ visitedLimit,
+ new AssertingSeededCollectorManager(new SeededCollectorManager(knnCollectorManager)));
}
- @Override
- public String toString(String field) {
- return null;
+ class AssertingSeededCollectorManager extends SeededCollectorManager {
+
+ public AssertingSeededCollectorManager(SeededCollectorManager delegate) {
+ super(delegate);
+ }
+
+ @Override
+ public KnnCollector newCollector(
+ int numVisisted, KnnSearchStrategy searchStrategy, LeafReaderContext context)
+ throws IOException {
+ KnnCollector knnCollector =
+ knnCollectorManager.newCollector(numVisisted, searchStrategy, context);
+ if (knnCollector.getSearchStrategy() instanceof KnnSearchStrategy.Seeded seeded) {
+ if (seedCalls == null && seeded.numberOfEntryPoints() > 0) {
+ fail("Expected non-seeded collector but received: " + knnCollector);
+ }
+ return new AssertingKnnCollector(knnCollector);
+ }
+ if (seedCalls != null) {
+ fail("Expected seeded collector but received: " + knnCollector);
+ }
+ return knnCollector;
+ }
+ }
+
+ class AssertingKnnCollector extends KnnCollector.Decorator {
+ public AssertingKnnCollector(KnnCollector collector) {
+ super(collector);
+ }
+
+ @Override
+ public KnnSearchStrategy getSearchStrategy() {
+ KnnSearchStrategy searchStrategy = collector.getSearchStrategy();
+ if (searchStrategy instanceof KnnSearchStrategy.Seeded seeded) {
+ return new AssertingSeededStrategy(seeded);
+ }
+ return searchStrategy;
+ }
+
+ class AssertingSeededStrategy extends KnnSearchStrategy.Seeded {
+ private final KnnSearchStrategy.Seeded seeded;
+
+ public AssertingSeededStrategy(KnnSearchStrategy.Seeded seeded) {
+ super(seeded.entryPoints(), seeded.numberOfEntryPoints());
+ this.seeded = seeded;
+ }
+
+ @Override
+ public int numberOfEntryPoints() {
+ return seeded.numberOfEntryPoints();
+ }
+
+ @Override
+ public DocIdSetIterator entryPoints() {
+ DocIdSetIterator iterator = seeded.entryPoints();
+ assert iterator.cost() > 0;
+ seedCalls.incrementAndGet();
+ return iterator;
+ }
+ }
}
}
}
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 d5630037ef7..c15b720b3d1 100644
--- a/lucene/core/src/test/org/apache/lucene/search/TestSeededKnnFloatVectorQuery.java
+++ b/lucene/core/src/test/org/apache/lucene/search/TestSeededKnnFloatVectorQuery.java
@@ -17,6 +17,7 @@
package org.apache.lucene.search;
import java.io.IOException;
+import java.util.concurrent.atomic.AtomicInteger;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.document.IntPoint;
@@ -25,8 +26,6 @@ import org.apache.lucene.document.NumericDocValuesField;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.IndexWriterConfig;
-import org.apache.lucene.index.LeafReaderContext;
-import org.apache.lucene.index.QueryTimeout;
import org.apache.lucene.index.VectorSimilarityFunction;
import org.apache.lucene.store.Directory;
import org.apache.lucene.tests.index.RandomIndexWriter;
@@ -37,13 +36,16 @@ public class TestSeededKnnFloatVectorQuery extends BaseKnnVectorQueryTestCase {
private static final Query MATCH_NONE = new MatchNoDocsQuery();
@Override
- KnnFloatVectorQuery getKnnVectorQuery(String field, float[] query, int k, Query queryFilter) {
- return new SeededKnnFloatVectorQuery(field, query, k, queryFilter, MATCH_NONE);
+ AbstractKnnVectorQuery getKnnVectorQuery(String field, float[] query, int k, Query queryFilter) {
+ KnnFloatVectorQuery knnQuery = new KnnFloatVectorQuery(field, query, k, queryFilter);
+ return SeededKnnVectorQuery.fromFloatQuery(knnQuery, MATCH_NONE);
}
@Override
AbstractKnnVectorQuery getThrowingKnnVectorQuery(String field, float[] vec, int k, Query query) {
- return new ThrowingKnnVectorQuery(field, vec, k, query, MATCH_NONE);
+ KnnFloatVectorQuery knnQuery =
+ new TestKnnFloatVectorQuery.ThrowingKnnVectorQuery(field, vec, k, query);
+ return SeededKnnVectorQuery.fromFloatQuery(knnQuery, MATCH_NONE);
}
@Override
@@ -62,6 +64,46 @@ public class TestSeededKnnFloatVectorQuery extends BaseKnnVectorQueryTestCase {
return new KnnFloatVectorField(name, vector);
}
+ public void testSeedWithTimeout() throws IOException {
+ int numDocs = atLeast(50);
+ int dimension = atLeast(5);
+ int numIters = atLeast(5);
+ try (Directory d = newDirectoryForTest()) {
+ IndexWriterConfig iwc = new IndexWriterConfig().setCodec(TestUtil.getDefaultCodec());
+ RandomIndexWriter w = new RandomIndexWriter(random(), d, iwc);
+ for (int i = 0; i < numDocs; i++) {
+ Document doc = new Document();
+ doc.add(getKnnVectorField("field", randomVector(dimension)));
+ doc.add(new NumericDocValuesField("tag", i));
+ doc.add(new IntPoint("tag", i));
+ w.addDocument(doc);
+ }
+ w.close();
+
+ try (IndexReader reader = DirectoryReader.open(d)) {
+ IndexSearcher searcher = newSearcher(reader);
+ searcher.setTimeout(() -> true);
+ int k = random().nextInt(80) + 1;
+ for (int i = 0; i < numIters; i++) {
+ // All documents as seeds
+ Query seed =
+ random().nextBoolean()
+ ? IntPoint.newRangeQuery("tag", 1, 6)
+ : new MatchAllDocsQuery();
+ Query filter = random().nextBoolean() ? null : new MatchAllDocsQuery();
+ KnnFloatVectorQuery knnFloatVectorQuery =
+ new KnnFloatVectorQuery("field", randomVector(dimension), k, filter);
+ Query knnQuery = SeededKnnVectorQuery.fromFloatQuery(knnFloatVectorQuery, seed);
+ assertEquals(0, searcher.count(knnQuery));
+ // No seed documents -- falls back on full approx search
+ seed = new MatchNoDocsQuery();
+ knnQuery = SeededKnnVectorQuery.fromFloatQuery(knnFloatVectorQuery, seed);
+ assertEquals(0, searcher.count(knnQuery));
+ }
+ }
+ }
+ }
+
/** Tests with random vectors and a random seed. Uses RandomIndexWriter. */
public void testRandomWithSeed() throws IOException {
int numDocs = 1000;
@@ -92,6 +134,12 @@ public class TestSeededKnnFloatVectorQuery extends BaseKnnVectorQueryTestCase {
try (IndexReader reader = DirectoryReader.open(d)) {
IndexSearcher searcher = newSearcher(reader);
for (int i = 0; i < numIters; i++) {
+ // verify timeout collector wrapping is used
+ if (random().nextBoolean()) {
+ searcher.setTimeout(() -> false);
+ } else {
+ searcher.setTimeout(null);
+ }
int k = random().nextInt(80) + 1;
int n = random().nextInt(100) + 1;
// we may get fewer results than requested if there are deletions, but this test doesn't
@@ -99,11 +147,16 @@ public class TestSeededKnnFloatVectorQuery extends BaseKnnVectorQueryTestCase {
assert reader.hasDeletions() == false;
// All documents as seeds
+ AtomicInteger seedCalls = new AtomicInteger();
Query seed1 = new MatchAllDocsQuery();
Query filter = random().nextBoolean() ? null : new MatchAllDocsQuery();
- AbstractKnnVectorQuery query =
- new SeededKnnFloatVectorQuery("field", randomVector(dimension), k, filter, seed1);
+ KnnFloatVectorQuery knnFloatVectorQuery =
+ new KnnFloatVectorQuery("field", randomVector(dimension), k, filter);
+ TestSeededKnnByteVectorQuery.AssertingSeededKnnVectorQuery query =
+ new TestSeededKnnByteVectorQuery.AssertingSeededKnnVectorQuery(
+ knnFloatVectorQuery, seed1, null, seedCalls);
TopDocs results = searcher.search(query, n);
+ assertEquals(seedCalls.get(), 1);
int expected = Math.min(Math.min(n, k), numDocsWithVector);
assertEquals(expected, results.scoreDocs.length);
@@ -117,8 +170,17 @@ public class TestSeededKnnFloatVectorQuery extends BaseKnnVectorQueryTestCase {
// Restrictive seed query -- 6 documents
Query seed2 = IntPoint.newRangeQuery("tag", 1, 6);
- query = new SeededKnnFloatVectorQuery("field", randomVector(dimension), k, null, seed2);
+ int seedCount =
+ searcher.count(
+ new BooleanQuery.Builder()
+ .add(seed2, BooleanClause.Occur.MUST)
+ .add(new FieldExistsQuery("field"), BooleanClause.Occur.MUST)
+ .build());
+ query =
+ new TestSeededKnnByteVectorQuery.AssertingSeededKnnVectorQuery(
+ knnFloatVectorQuery, seed2, null, seedCount > 0 ? seedCalls : null);
results = searcher.search(query, n);
+ assertEquals(seedCalls.get(), seedCount > 0 ? 2 : 1);
expected = Math.min(Math.min(n, k), reader.numDocs());
assertEquals(expected, results.scoreDocs.length);
assertTrue(results.totalHits.value() >= results.scoreDocs.length);
@@ -131,7 +193,9 @@ public class TestSeededKnnFloatVectorQuery extends BaseKnnVectorQueryTestCase {
// No seed documents -- falls back on full approx search
Query seed3 = new MatchNoDocsQuery();
- query = new SeededKnnFloatVectorQuery("field", randomVector(dimension), k, null, seed3);
+ query =
+ new TestSeededKnnByteVectorQuery.AssertingSeededKnnVectorQuery(
+ knnFloatVectorQuery, seed3, null, null);
results = searcher.search(query, n);
expected = Math.min(Math.min(n, k), reader.numDocs());
assertEquals(expected, results.scoreDocs.length);
@@ -146,46 +210,4 @@ public class TestSeededKnnFloatVectorQuery extends BaseKnnVectorQueryTestCase {
}
}
}
-
- private static class ThrowingKnnVectorQuery extends SeededKnnFloatVectorQuery {
-
- private ThrowingKnnVectorQuery(String field, float[] target, int k, Query filter, Query seed) {
- super(field, target, k, filter, seed);
- }
-
- private ThrowingKnnVectorQuery(
- String field, float[] target, int k, Query filter, Weight seedWeight) {
- super(field, target, k, filter, seedWeight);
- }
-
- @Override
- // This is test only and we need to overwrite the inner rewrite to throw
- public Query rewrite(IndexSearcher indexSearcher) throws IOException {
- if (seedWeight != null) {
- return super.rewrite(indexSearcher);
- }
- BooleanQuery.Builder booleanSeedQueryBuilder =
- new BooleanQuery.Builder()
- .add(seed, BooleanClause.Occur.MUST)
- .add(new FieldExistsQuery(field), BooleanClause.Occur.FILTER);
- if (filter != null) {
- booleanSeedQueryBuilder.add(filter, BooleanClause.Occur.FILTER);
- }
- Query seedRewritten = indexSearcher.rewrite(booleanSeedQueryBuilder.build());
- Weight seedWeight = indexSearcher.createWeight(seedRewritten, ScoreMode.TOP_SCORES, 1f);
- return new ThrowingKnnVectorQuery(field, target, k, filter, seedWeight)
- .rewrite(indexSearcher);
- }
-
- @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 null;
- }
- }
}
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 45cb8b9c88f..a0b0f5d32ac 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
@@ -144,7 +144,7 @@ public class DiversifyingChildrenByteKnnVectorQuery extends KnnByteVectorQuery {
KnnCollectorManager knnCollectorManager)
throws IOException {
ByteVectorValues.checkField(context.reader(), field);
- KnnCollector collector = knnCollectorManager.newCollector(visitedLimit, context);
+ KnnCollector collector = knnCollectorManager.newCollector(visitedLimit, null, 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 9c44a2f7856..3b05b392654 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
@@ -143,7 +143,7 @@ public class DiversifyingChildrenFloatKnnVectorQuery extends KnnFloatVectorQuery
KnnCollectorManager knnCollectorManager)
throws IOException {
FloatVectorValues.checkField(context.reader(), field);
- KnnCollector collector = knnCollectorManager.newCollector(visitedLimit, context);
+ KnnCollector collector = knnCollectorManager.newCollector(visitedLimit, null, context);
if (collector == null) {
return NO_RESULTS;
}
diff --git a/lucene/join/src/java/org/apache/lucene/search/join/DiversifyingNearestChildrenKnnCollector.java b/lucene/join/src/java/org/apache/lucene/search/join/DiversifyingNearestChildrenKnnCollector.java
index 15227a50342..53b7cfae92c 100644
--- a/lucene/join/src/java/org/apache/lucene/search/join/DiversifyingNearestChildrenKnnCollector.java
+++ b/lucene/join/src/java/org/apache/lucene/search/join/DiversifyingNearestChildrenKnnCollector.java
@@ -22,6 +22,7 @@ import org.apache.lucene.search.AbstractKnnCollector;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.search.TotalHits;
+import org.apache.lucene.search.knn.KnnSearchStrategy;
import org.apache.lucene.util.ArrayUtil;
import org.apache.lucene.util.BitSet;
@@ -42,7 +43,20 @@ class DiversifyingNearestChildrenKnnCollector extends AbstractKnnCollector {
* @param parentBitSet The leaf parent bitset
*/
public DiversifyingNearestChildrenKnnCollector(int k, int visitLimit, BitSet parentBitSet) {
- super(k, visitLimit);
+ this(k, visitLimit, null, parentBitSet);
+ }
+
+ /**
+ * Create a new object for joining nearest child kNN documents with a parent bitset
+ *
+ * @param k The number of joined parent documents to collect
+ * @param visitLimit how many child vectors can be visited
+ * @param searchStrategy The search strategy to use
+ * @param parentBitSet The leaf parent bitset
+ */
+ public DiversifyingNearestChildrenKnnCollector(
+ int k, int visitLimit, KnnSearchStrategy searchStrategy, BitSet parentBitSet) {
+ super(k, visitLimit, searchStrategy);
this.parentBitSet = parentBitSet;
this.heap = new NodeIdCachingHeap(k);
}
diff --git a/lucene/join/src/java/org/apache/lucene/search/join/DiversifyingNearestChildrenKnnCollectorManager.java b/lucene/join/src/java/org/apache/lucene/search/join/DiversifyingNearestChildrenKnnCollectorManager.java
index 6450af97d0d..db0e38eafe8 100644
--- a/lucene/join/src/java/org/apache/lucene/search/join/DiversifyingNearestChildrenKnnCollectorManager.java
+++ b/lucene/join/src/java/org/apache/lucene/search/join/DiversifyingNearestChildrenKnnCollectorManager.java
@@ -22,6 +22,7 @@ import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.KnnCollector;
import org.apache.lucene.search.knn.KnnCollectorManager;
+import org.apache.lucene.search.knn.KnnSearchStrategy;
import org.apache.lucene.search.knn.MultiLeafKnnCollector;
import org.apache.lucene.util.BitSet;
import org.apache.lucene.util.hnsw.BlockingFloatHeap;
@@ -59,18 +60,22 @@ public class DiversifyingNearestChildrenKnnCollectorManager implements KnnCollec
* @param context the leaf reader context
*/
@Override
- public KnnCollector newCollector(int visitedLimit, LeafReaderContext context) throws IOException {
+ public KnnCollector newCollector(
+ int visitedLimit, KnnSearchStrategy searchStrategy, LeafReaderContext context)
+ throws IOException {
BitSet parentBitSet = parentsFilter.getBitSet(context);
if (parentBitSet == null) {
return null;
}
if (globalScoreQueue == null) {
- return new DiversifyingNearestChildrenKnnCollector(k, visitedLimit, parentBitSet);
+ return new DiversifyingNearestChildrenKnnCollector(
+ k, visitedLimit, searchStrategy, parentBitSet);
} else {
return new MultiLeafKnnCollector(
k,
globalScoreQueue,
- new DiversifyingNearestChildrenKnnCollector(k, visitedLimit, parentBitSet));
+ new DiversifyingNearestChildrenKnnCollector(
+ k, visitedLimit, searchStrategy, parentBitSet));
}
}
}
(lucene) 01/03: Adjust knn merge stability testing (#14172)
Posted by be...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
benwtrent pushed a commit to branch branch_10x
in repository https://gitbox.apache.org/repos/asf/lucene.git
commit 1baccbd77b91be1b1cd4c4b87be031a823f074a8
Author: Benjamin Trent <be...@gmail.com>
AuthorDate: Wed Jan 29 11:10:23 2025 -0500
Adjust knn merge stability testing (#14172)
---
.../apache/lucene/tests/index/BaseIndexFileFormatTestCase.java | 3 +++
.../apache/lucene/tests/index/BaseKnnVectorsFormatTestCase.java | 8 ++++++++
2 files changed, 11 insertions(+)
diff --git a/lucene/test-framework/src/java/org/apache/lucene/tests/index/BaseIndexFileFormatTestCase.java b/lucene/test-framework/src/java/org/apache/lucene/tests/index/BaseIndexFileFormatTestCase.java
index 297c1b777f5..c2aa7ff0e4d 100644
--- a/lucene/test-framework/src/java/org/apache/lucene/tests/index/BaseIndexFileFormatTestCase.java
+++ b/lucene/test-framework/src/java/org/apache/lucene/tests/index/BaseIndexFileFormatTestCase.java
@@ -275,6 +275,9 @@ public abstract class BaseIndexFileFormatTestCase extends LuceneTestCase {
new IndexWriterConfig(new MockAnalyzer(random()))
.setUseCompoundFile(false)
.setMergePolicy(mp);
+ if (VERBOSE) {
+ cfg.setInfoStream(System.out);
+ }
IndexWriter w = new IndexWriter(dir, cfg);
final int numDocs = atLeast(500);
for (int i = 0; i < numDocs; ++i) {
diff --git a/lucene/test-framework/src/java/org/apache/lucene/tests/index/BaseKnnVectorsFormatTestCase.java b/lucene/test-framework/src/java/org/apache/lucene/tests/index/BaseKnnVectorsFormatTestCase.java
index e42d3e18981..2bbeca0eee9 100644
--- a/lucene/test-framework/src/java/org/apache/lucene/tests/index/BaseKnnVectorsFormatTestCase.java
+++ b/lucene/test-framework/src/java/org/apache/lucene/tests/index/BaseKnnVectorsFormatTestCase.java
@@ -114,6 +114,14 @@ public abstract class BaseKnnVectorsFormatTestCase extends BaseIndexFileFormatTe
}
}
+ @Override
+ protected boolean mergeIsStable() {
+ // suppress this test from base class: merges for knn graphs are not stable due to connected
+ // components
+ // logic
+ return false;
+ }
+
private int getVectorsMaxDimensions(String fieldName) {
return Codec.getDefault().knnVectorsFormat().getMaxDimensions(fieldName);
}
(lucene) 02/03: Merge remote-tracking branch 'refs/remotes/origin/branch_10x' into branch_10x
Posted by be...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
benwtrent pushed a commit to branch branch_10x
in repository https://gitbox.apache.org/repos/asf/lucene.git
commit 823f2c865cf8a841fe2184c24f36950c8332a565
Merge: 1baccbd77b9 9965ec46814
Author: Benjamin Trent <43...@users.noreply.github.com>
AuthorDate: Mon Feb 3 11:07:18 2025 -0500
Merge remote-tracking branch 'refs/remotes/origin/branch_10x' into branch_10x
gradle/testing/randomization.gradle | 4 +++-
lucene/CHANGES.txt | 4 ++++
.../codecs/lucene101/Lucene101PostingsReader.java | 20 +++++++----------
.../org/apache/lucene/index/LogMergePolicy.java | 25 ++++++++++++++++++++++
.../apache/lucene/search/MaxScoreBulkScorer.java | 3 ++-
.../apache/lucene/index/TestLogMergePolicy.java | 7 +++++-
6 files changed, 48 insertions(+), 15 deletions(-)