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

lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsReader.java

@@ -36,6 +36,7 @@
 import org.apache.lucene.index.VectorEncoding;
 import org.apache.lucene.index.VectorSimilarityFunction;
 import org.apache.lucene.internal.hppc.IntObjectHashMap;
+import org.apache.lucene.search.HnswQueueSaturationCollector;
 import org.apache.lucene.search.KnnCollector;
 import org.apache.lucene.store.ChecksumIndexInput;
 import org.apache.lucene.store.DataInput;
@@ -314,10 +315,16 @@ private void search(
       return;
     }
     final RandomVectorScorer scorer = scorerSupplier.get();
-    final KnnCollector collector =
-        new OrdinalTranslatedKnnCollector(knnCollector, scorer::ordToDoc);
     final Bits acceptedOrds = scorer.getAcceptOrds(acceptDocs);
     if (knnCollector.k() < scorer.maxOrd()) {
+      final KnnCollector collector;
+      OrdinalTranslatedKnnCollector ordinalTranslatedKnnCollector =
+          new OrdinalTranslatedKnnCollector(knnCollector, scorer::ordToDoc);
+      if (scorer.maxOrd() > 1000) {
+        collector = new HnswQueueSaturationCollector(ordinalTranslatedKnnCollector);
+      } else {
+        collector = ordinalTranslatedKnnCollector;
+      }
       HnswGraphSearcher.search(scorer, collector, getGraph(fieldEntry), acceptedOrds);
     } else {
       // if k is larger than the number of vectors, we can just iterate over all vectors

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

@@ -0,0 +1,24 @@
+/*
+ * 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;
+
+/** {@link KnnCollector} that exposes methods to hook into specific parts of the HNSW algorithm. */
+public interface HnswKnnCollector extends KnnCollector {
+
+  /** Indicates exploration of the next HNSW candidate graph node. */
+  void nextCandidate();
+}

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

@@ -0,0 +1,119 @@
+/*
+ * 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;
+
+/**
+ * A {@link HnswKnnCollector} that early exits when nearest neighbor queue keeps saturating beyond a
+ * 'patience' parameter.
+ */
+public class HnswQueueSaturationCollector implements HnswKnnCollector {
+
+  private static final double DEFAULT_SATURATION_THRESHOLD = 0.995d;
+
+  private final KnnCollector delegate;
+  private double saturationThreshold;
+  private int patience;
+  private boolean patienceFinished;
+  private int countSaturated;
+  private int previousQueueSize;
+  private int currentQueueSize;
+
+  public HnswQueueSaturationCollector(
+      KnnCollector delegate, double saturationThreshold, int patience) {
+    this.delegate = delegate;
+    this.previousQueueSize = 0;
+    this.currentQueueSize = 0;
+    this.countSaturated = 0;
+    this.patienceFinished = false;
+    this.saturationThreshold = saturationThreshold;
+    this.patience = patience;
+  }
+
+  public HnswQueueSaturationCollector(KnnCollector delegate) {
+    this.delegate = delegate;
+    this.previousQueueSize = 0;
+    this.currentQueueSize = 0;
+    this.countSaturated = 0;
+    this.patienceFinished = false;
+    this.saturationThreshold = DEFAULT_SATURATION_THRESHOLD;
+    this.patience = defaultPatience();
+  }
+
+  private int defaultPatience() {
+    return Math.max(7, (int) (k() * 0.3));
+  }
+
+  @Override
+  public boolean earlyTerminated() {
+    return delegate.earlyTerminated() || patienceFinished;
+  }
+
+  @Override
+  public void incVisitedCount(int count) {
+    delegate.incVisitedCount(count);
+  }
+
+  @Override
+  public long visitedCount() {
+    return delegate.visitedCount();
+  }
+
+  @Override
+  public long visitLimit() {
+    return delegate.visitLimit();
+  }
+
+  @Override
+  public int k() {
+    return delegate.k();
+  }
+
+  @Override
+  public boolean collect(int docId, float similarity) {
+    boolean collect = delegate.collect(docId, similarity);
+    if (collect) {
+      currentQueueSize++;
+    }
+    return collect;
+  }
+
+  @Override
+  public float minCompetitiveSimilarity() {
+    return delegate.minCompetitiveSimilarity();
+  }
+
+  @Override
+  public TopDocs topDocs() {
+    return delegate.topDocs();
+  }
+
+  @Override
+  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;
+    }
+  }
+}

lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphSearcher.java

@@ -20,6 +20,7 @@
 import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
 
 import java.io.IOException;
+import org.apache.lucene.search.HnswKnnCollector;
 import org.apache.lucene.search.KnnCollector;
 import org.apache.lucene.search.TopKnnCollector;
 import org.apache.lucene.util.BitSet;
@@ -245,6 +246,9 @@ void searchLevel(
           }
         }
       }
+      if (results instanceof HnswKnnCollector hnswKnnCollector) {
+        hnswKnnCollector.nextCandidate();
+      }
     }
   }
 

lucene/core/src/test/org/apache/lucene/search/HnswQueueSaturationCollectorTest.java

@@ -0,0 +1,43 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.search;
+
+import java.util.Random;
+import org.apache.lucene.tests.util.LuceneTestCase;
+import org.junit.Test;
+
+/** Tests for {@link HnswQueueSaturationCollector} */
+public class HnswQueueSaturationCollectorTest extends LuceneTestCase {
+
+  @Test
+  public void testDelegate() {
+    Random random = random();
+    KnnCollector delegate = new TopKnnCollector(random.nextInt(100), random.nextInt(1000));
+    HnswQueueSaturationCollector queueSaturationCollector =
+        new HnswQueueSaturationCollector(delegate);
+    for (int i = 0; i < random.nextInt(100); i++) {
+      queueSaturationCollector.collect(random.nextInt(1000), random.nextFloat(1.0f));
+    }
+    assertEquals(delegate.k(), queueSaturationCollector.k());
+    assertEquals(delegate.visitedCount(), queueSaturationCollector.visitedCount());
+    assertEquals(delegate.visitLimit(), queueSaturationCollector.visitLimit());
+    assertEquals(
+        delegate.minCompetitiveSimilarity(),
+        queueSaturationCollector.minCompetitiveSimilarity(),
+        1e-3);
+  }
+}