Add a HNSW collector that exits early when nearest neighbor queue saturates

lucene/CHANGES.txt

@@ -139,6 +139,8 @@ New Features
 
 * GITHUB#14412: Allow skip cache factor to be updated dynamically. (Sagar Upadhyaya)
 
+* GITHUB#14094: New KNN query that early terminates when HNSW nearest neighbor queue saturates. (Tommaso Teofili)
+
 Improvements
 ---------------------
 

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

@@ -0,0 +1,101 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.lucene.search;
+
+import org.apache.lucene.search.knn.KnnSearchStrategy;
+
+/**
+ * A {@link KnnCollector.Decorator} that early exits when nearest neighbor queue keeps saturating
+ * beyond a 'patience' parameter. This records the rate of collection of new nearest neighbors in
+ * the {@code delegate} KnnCollector queue, at each HNSW node candidate visit. Once it saturates for
+ * a number of consecutive node visits (e.g., the patience parameter), this early terminates.
+ *
+ * @lucene.experimental
+ */
+public class HnswQueueSaturationCollector extends KnnCollector.Decorator {
+
+  private final KnnCollector delegate;
+  private final double saturationThreshold;
+  private final int patience;
+  private boolean patienceFinished;
+  private int countSaturated;
+  private int previousQueueSize;
+  private int currentQueueSize;
+
+  HnswQueueSaturationCollector(KnnCollector delegate, double saturationThreshold, int patience) {
+    super(delegate);
+    this.delegate = delegate;
+    this.previousQueueSize = 0;
+    this.currentQueueSize = 0;
+    this.countSaturated = 0;
+    this.patienceFinished = false;
+    this.saturationThreshold = saturationThreshold;
+    this.patience = patience;
+  }
+
+  @Override
+  public boolean earlyTerminated() {
+    return delegate.earlyTerminated() || patienceFinished;
+  }
+
+  @Override
+  public boolean collect(int docId, float similarity) {
+    boolean collect = delegate.collect(docId, similarity);
+    if (collect) {
+      currentQueueSize++;
+    }
+    return collect;
+  }
+
+  @Override
+  public TopDocs topDocs() {
+    TopDocs topDocs;
+    if (patienceFinished && delegate.earlyTerminated() == false) {
+      // this avoids re-running exact search in the filtered scenario when patience is exhausted
+      TopDocs delegateDocs = delegate.topDocs();
+      TotalHits totalHits =
+          new TotalHits(delegateDocs.totalHits.value(), TotalHits.Relation.EQUAL_TO);
+      topDocs = new TopDocs(totalHits, delegateDocs.scoreDocs);
+    } else {
+      topDocs = delegate.topDocs();
+    }
+    return topDocs;
+  }
+
+  public void nextCandidate() {
+    double queueSaturation =
+        (double) Math.min(currentQueueSize, previousQueueSize) / currentQueueSize;
+    previousQueueSize = currentQueueSize;
+    if (queueSaturation >= saturationThreshold) {
+      countSaturated++;
+    } else {
+      countSaturated = 0;
+    }
+    if (countSaturated > patience) {
+      patienceFinished = true;
+    }
+  }
+
+  @Override
+  public KnnSearchStrategy getSearchStrategy() {
+    KnnSearchStrategy delegateStrategy = delegate.getSearchStrategy();
+    assert delegateStrategy instanceof KnnSearchStrategy.Hnsw;
+    return new KnnSearchStrategy.Patience(
+        this, ((KnnSearchStrategy.Hnsw) delegateStrategy).filteredSearchThreshold());
+  }
+}

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

@@ -0,0 +1,237 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.search;
+
+import java.io.IOException;
+import java.util.Objects;
+import org.apache.lucene.index.FieldInfo;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.QueryTimeout;
+import org.apache.lucene.search.knn.KnnCollectorManager;
+import org.apache.lucene.search.knn.KnnSearchStrategy;
+import org.apache.lucene.util.Bits;
+
+/**
+ * This is a version of knn vector query that exits early when HNSW queue saturates over a {@code
+ * #saturationThreshold} for more than {@code #patience} times.
+ *
+ * <p>See <a
+ * href="https://cs.uwaterloo.ca/~jimmylin/publications/Teofili_Lin_ECIR2025.pdf">"Patience in
+ * Proximity: A Simple Early Termination Strategy for HNSW Graph Traversal in Approximate k-Nearest
+ * Neighbor Search"</a> (Teofili and Lin). In ECIR '25: Proceedings of the 47th European Conference
+ * on Information Retrieval.
+ *
+ * @lucene.experimental
+ */
+public class PatienceKnnVectorQuery extends AbstractKnnVectorQuery {
+
+  private static final double DEFAULT_SATURATION_THRESHOLD = 0.995d;
+
+  private final int patience;
+  private final double saturationThreshold;
+
+  final AbstractKnnVectorQuery delegate;
+
+  /**
+   * Construct a new PatienceKnnVectorQuery instance for a float vector field
+   *
+   * @param knnQuery the knn query to be seeded
+   * @param saturationThreshold the early exit saturation threshold
+   * @param patience the patience parameter
+   * @return a new PatienceKnnVectorQuery instance
+   * @lucene.experimental
+   */
+  public static PatienceKnnVectorQuery fromFloatQuery(
+      KnnFloatVectorQuery knnQuery, double saturationThreshold, int patience) {
+    return new PatienceKnnVectorQuery(knnQuery, saturationThreshold, patience);
+  }
+
+  /**
+   * Construct a new PatienceKnnVectorQuery instance for a float vector field
+   *
+   * @param knnQuery the knn query to be seeded
+   * @return a new PatienceKnnVectorQuery instance
+   * @lucene.experimental
+   */
+  public static PatienceKnnVectorQuery fromFloatQuery(KnnFloatVectorQuery knnQuery) {
+    return new PatienceKnnVectorQuery(
+        knnQuery, DEFAULT_SATURATION_THRESHOLD, defaultPatience(knnQuery));
+  }
+
+  /**
+   * Construct a new PatienceKnnVectorQuery instance for a byte vector field
+   *
+   * @param knnQuery the knn query to be seeded
+   * @param saturationThreshold the early exit saturation threshold
+   * @param patience the patience parameter
+   * @return a new PatienceKnnVectorQuery instance
+   * @lucene.experimental
+   */
+  public static PatienceKnnVectorQuery fromByteQuery(
+      KnnByteVectorQuery knnQuery, double saturationThreshold, int patience) {
+    return new PatienceKnnVectorQuery(knnQuery, saturationThreshold, patience);
+  }
+
+  /**
+   * Construct a new PatienceKnnVectorQuery instance for a byte vector field
+   *
+   * @param knnQuery the knn query to be seeded
+   * @return a new PatienceKnnVectorQuery instance
+   * @lucene.experimental
+   */
+  public static PatienceKnnVectorQuery fromByteQuery(KnnByteVectorQuery knnQuery) {
+    return new PatienceKnnVectorQuery(
+        knnQuery, DEFAULT_SATURATION_THRESHOLD, defaultPatience(knnQuery));
+  }
+
+  /**
+   * Construct a new PatienceKnnVectorQuery instance for seeded vector field
+   *
+   * @param knnQuery the knn query to be seeded
+   * @param saturationThreshold the early exit saturation threshold
+   * @param patience the patience parameter
+   * @return a new PatienceKnnVectorQuery instance
+   * @lucene.experimental
+   */
+  public static PatienceKnnVectorQuery fromSeededQuery(
+      SeededKnnVectorQuery knnQuery, double saturationThreshold, int patience) {
+    return new PatienceKnnVectorQuery(knnQuery, saturationThreshold, patience);
+  }
+
+  /**
+   * Construct a new PatienceKnnVectorQuery instance for seeded vector field
+   *
+   * @param knnQuery the knn query to be seeded
+   * @return a new PatienceKnnVectorQuery instance
+   * @lucene.experimental
+   */
+  public static PatienceKnnVectorQuery fromSeededQuery(SeededKnnVectorQuery knnQuery) {
+    return new PatienceKnnVectorQuery(
+        knnQuery, DEFAULT_SATURATION_THRESHOLD, defaultPatience(knnQuery));
+  }
+
+  PatienceKnnVectorQuery(
+      AbstractKnnVectorQuery knnQuery, double saturationThreshold, int patience) {
+    super(knnQuery.field, knnQuery.k, knnQuery.filter, knnQuery.searchStrategy);
+    this.delegate = knnQuery;
+    this.saturationThreshold = saturationThreshold;
+    this.patience = patience;
+  }
+
+  private static int defaultPatience(AbstractKnnVectorQuery delegate) {
+    return Math.max(7, (int) (delegate.k * 0.3));
+  }
+
+  @Override
+  public String toString(String field) {
+    return "PatienceKnnVectorQuery{"
+        + "saturationThreshold="
+        + saturationThreshold
+        + ", patience="
+        + patience
+        + ", delegate="
+        + delegate
+        + '}';
+  }
+
+  @Override
+  protected KnnCollectorManager getKnnCollectorManager(int k, IndexSearcher searcher) {
+    return delegate.getKnnCollectorManager(k, searcher);
+  }
+
+  @Override
+  protected TopDocs approximateSearch(
+      LeafReaderContext context,
+      Bits acceptDocs,
+      int visitedLimit,
+      KnnCollectorManager knnCollectorManager)
+      throws IOException {
+    return delegate.approximateSearch(
+        context, acceptDocs, visitedLimit, new PatienceCollectorManager(knnCollectorManager));
+  }
+
+  @Override
+  protected TopDocs exactSearch(
+      LeafReaderContext context, DocIdSetIterator acceptIterator, QueryTimeout queryTimeout)
+      throws IOException {
+    return delegate.exactSearch(context, acceptIterator, queryTimeout);
+  }
+
+  @Override
+  protected TopDocs mergeLeafResults(TopDocs[] perLeafResults) {
+    return delegate.mergeLeafResults(perLeafResults);
+  }
+
+  @Override
+  public void visit(QueryVisitor visitor) {
+    delegate.visit(visitor);
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) return true;
+    if (o == null || getClass() != o.getClass()) return false;
+    if (!super.equals(o)) return false;
+    PatienceKnnVectorQuery that = (PatienceKnnVectorQuery) o;
+    return saturationThreshold == that.saturationThreshold
+        && patience == that.patience
+        && Objects.equals(delegate, that.delegate);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(super.hashCode(), saturationThreshold, patience, delegate);
+  }
+
+  @Override
+  public String getField() {
+    return delegate.getField();
+  }
+
+  @Override
+  public int getK() {
+    return delegate.getK();
+  }
+
+  @Override
+  public Query getFilter() {
+    return delegate.getFilter();
+  }
+
+  @Override
+  VectorScorer createVectorScorer(LeafReaderContext context, FieldInfo fi) throws IOException {
+    return delegate.createVectorScorer(context, fi);
+  }
+
+  class PatienceCollectorManager implements KnnCollectorManager {
+    final KnnCollectorManager knnCollectorManager;
+
+    PatienceCollectorManager(KnnCollectorManager knnCollectorManager) {
+      this.knnCollectorManager = knnCollectorManager;
+    }
+
+    @Override
+    public KnnCollector newCollector(
+        int visitLimit, KnnSearchStrategy searchStrategy, LeafReaderContext ctx)
+        throws IOException {
+      return new HnswQueueSaturationCollector(
+          knnCollectorManager.newCollector(visitLimit, searchStrategy, ctx),
+          saturationThreshold,
+          patience);
+    }
+  }
+}