Adjusts Seeded knn searches to clean up user and internal interfaces (apache#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.

Author: benwtrent

lucene/CHANGES.txt

@@ -49,8 +49,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).
 
 

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 final int k() {
 
   @Override
   public abstract TopDocs topDocs();
+
+  @Override
+  public KnnSearchStrategy getSearchStrategy() {
+    return searchStrategy;
+  }
 }

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;

lucene/core/src/java/org/apache/lucene/search/ByteVectorSimilarityQuery.java

@@ -113,7 +113,7 @@ protected TopDocs approximateSearch(
       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();
   }

lucene/core/src/java/org/apache/lucene/search/FloatVectorSimilarityQuery.java

@@ -115,7 +115,7 @@ protected TopDocs approximateSearch(
       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();
   }

lucene/core/src/java/org/apache/lucene/search/KnnByteVectorQuery.java

@@ -83,7 +83,7 @@ protected TopDocs approximateSearch(
       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) {

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,14 +88,21 @@ 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.
    *
    * @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 float minCompetitiveSimilarity() {
     public TopDocs topDocs() {
       return collector.topDocs();
     }
+
+    @Override
+    public KnnSearchStrategy getSearchStrategy() {
+      return collector.getSearchStrategy();
+    }
   }
 }

lucene/core/src/java/org/apache/lucene/search/KnnFloatVectorQuery.java

@@ -84,7 +84,7 @@ protected TopDocs approximateSearch(
       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) {