Add two new "Seeded" Knn queries for seeded vector search

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

@@ -46,7 +46,7 @@ public class KnnByteVectorQuery extends AbstractKnnVectorQuery {
 
   private static final TopDocs NO_RESULTS = TopDocsCollector.EMPTY_TOPDOCS;
 
-  private final byte[] target;
+  protected final byte[] target;
 
   /**
    * Find the <code>k</code> nearest documents to the target vector according to the vectors in the

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

@@ -85,4 +85,58 @@ public interface KnnCollector {
    * @return The collected top documents
    */
   TopDocs topDocs();
+
+  /**
+   * 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 KnnCollector collector;
+
+    public Decorator(KnnCollector collector) {
+      this.collector = collector;
+    }
+
+    @Override
+    public boolean earlyTerminated() {
+      return collector.earlyTerminated();
+    }
+
+    @Override
+    public void incVisitedCount(int count) {
+      collector.incVisitedCount(count);
+    }
+
+    @Override
+    public long visitedCount() {
+      return collector.visitedCount();
+    }
+
+    @Override
+    public long visitLimit() {
+      return collector.visitLimit();
+    }
+
+    @Override
+    public int k() {
+      return collector.k();
+    }
+
+    @Override
+    public boolean collect(int docId, float similarity) {
+      return collector.collect(docId, similarity);
+    }
+
+    @Override
+    public float minCompetitiveSimilarity() {
+      return collector.minCompetitiveSimilarity();
+    }
+
+    @Override
+    public TopDocs topDocs() {
+      return collector.topDocs();
+    }
+  }
 }

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

@@ -47,7 +47,7 @@ public class KnnFloatVectorQuery extends AbstractKnnVectorQuery {
 
   private static final TopDocs NO_RESULTS = TopDocsCollector.EMPTY_TOPDOCS;
 
-  private final float[] target;
+  protected final float[] target;
 
   /**
    * Find the <code>k</code> nearest documents to the target vector according to the vectors in the

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

@@ -0,0 +1,85 @@
+/*
+ * 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.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
+ *
+ * @lucene.experimental
+ */
+public class SeededKnnByteVectorQuery extends KnnByteVectorQuery {
+  private final Query seed;
+  private final Weight seedWeight;
+
+  /**
+   * Construct a new SeededKnnFloatVectorQuery 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 -> leaf.getFloatVectorValues(field));
+  }
+}

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

@@ -0,0 +1,85 @@
+/*
+ * 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.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
+ *
+ * @lucene.experimental
+ */
+public class SeededKnnFloatVectorQuery extends KnnFloatVectorQuery {
+  private final Query seed;
+  private 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 -> leaf.getFloatVectorValues(field));
+  }
+}

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

@@ -45,51 +45,19 @@ public KnnCollector newCollector(int visitedLimit, LeafReaderContext context) th
     return new TimeLimitingKnnCollector(collector);
   }
 
-  class TimeLimitingKnnCollector implements KnnCollector {
-    private final KnnCollector collector;
-
-    TimeLimitingKnnCollector(KnnCollector collector) {
-      this.collector = collector;
+  class TimeLimitingKnnCollector extends KnnCollector.Decorator {
+    public TimeLimitingKnnCollector(KnnCollector collector) {
+      super(collector);
     }
 
     @Override
     public boolean earlyTerminated() {
-      return queryTimeout.shouldExit() || collector.earlyTerminated();
-    }
-
-    @Override
-    public void incVisitedCount(int count) {
-      collector.incVisitedCount(count);
-    }
-
-    @Override
-    public long visitedCount() {
-      return collector.visitedCount();
-    }
-
-    @Override
-    public long visitLimit() {
-      return collector.visitLimit();
-    }
-
-    @Override
-    public int k() {
-      return collector.k();
-    }
-
-    @Override
-    public boolean collect(int docId, float similarity) {
-      return collector.collect(docId, similarity);
-    }
-
-    @Override
-    public float minCompetitiveSimilarity() {
-      return collector.minCompetitiveSimilarity();
+      return queryTimeout.shouldExit() || super.earlyTerminated();
     }
 
     @Override
     public TopDocs topDocs() {
-      TopDocs docs = collector.topDocs();
+      TopDocs docs = super.topDocs();
 
       // Mark results as partial if timeout is met
       TotalHits.Relation relation =

lucene/core/src/java/org/apache/lucene/search/knn/EntryPointProvider.java

@@ -0,0 +1,25 @@
+/*
+ * 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();
+}

lucene/core/src/java/org/apache/lucene/search/knn/SeededKnnCollector.java

@@ -0,0 +1,34 @@
+/*
+ * 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;
+
+public class SeededKnnCollector extends KnnCollector.Decorator implements EntryPointProvider {
+  private final DocIdSetIterator entryPoints;
+
+  public SeededKnnCollector(KnnCollector collector, DocIdSetIterator entryPoints) {
+    super(collector);
+    this.entryPoints = entryPoints;
+  }
+
+  @Override
+  public DocIdSetIterator entryPoints() {
+    return entryPoints;
+  }
+}