Sometimes use FixedBitSet when doing HNSW searches (apache#14836)

For smaller graphs, the overhead cost of a SparseFixedBitSet shows up in the performance metrics.

This adjusts the bitset creation logic to be more similar to how we utilize [Sparse]FixedBitSet elsewhere.

The expectedVisitedNodes is both empirical and intuitive. This could possibly be refined given the number of connections within the graph, but I think this is "good enough" for now.

Author: benwtrent

lucene/CHANGES.txt

@@ -90,6 +90,9 @@ Optimizations
 
 * GITHUB#14774: Use IntArrayList/IntHashSet to replace usages of List/Set of Integer. (Zhang Chao)
 
+* GITHUB#14836: Sometimes use FixedBitSet when doing HNSW searches. This slightly improves HNSW
+  search performance on smaller graphs. (Ben Trent)
+
 Bug Fixes
 ---------------------
 * GITHUB#14654: ValueSource.fromDoubleValuesSource(dvs).getSortField() would throw errors when

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

@@ -337,7 +337,7 @@ private void search(
     // Take into account if quantized? E.g. some scorer cost?
     int filteredDocCount = 0;
     // The approximate number of vectors that would be visited if we did not filter
-    int unfilteredVisit = (int) (Math.log(graph.size()) * knnCollector.k());
+    int unfilteredVisit = HnswGraphSearcher.expectedVisitedNodes(knnCollector.k(), graph.size());
     if (acceptDocs instanceof BitSet bitSet) {
       // Use approximate cardinality as this is good enough, but ensure we don't exceed the graph
       // size as that is illogical

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.DocIdSetIterator;
 import org.apache.lucene.search.KnnCollector;
 import org.apache.lucene.search.TopKnnCollector;
 import org.apache.lucene.search.knn.KnnSearchStrategy;
@@ -41,6 +42,32 @@ public class HnswGraphSearcher extends AbstractHnswGraphSearcher {
 
   protected BitSet visited;
 
+  /**
+   * HNSW search is roughly logarithmic. This doesn't take maxConn into account, but it is a pretty
+   * good approximation.
+   *
+   * @param k neighbors to find
+   * @param graphSize size of the graph
+   * @return expected number of visited nodes
+   */
+  public static int expectedVisitedNodes(int k, int graphSize) {
+    return (int) (Math.log(graphSize) * k);
+  }
+
+  /**
+   * Follows similar logic to {@link FixedBitSet#of(DocIdSetIterator, int)} to determine the best
+   * bit set given the expected number of visited nodes vs total graph size.
+   *
+   * @param k neighbors to find
+   * @param graphSize size of the graph
+   * @return a bit set appropriate for the expected number of visited nodes
+   */
+  static BitSet createBitSet(int k, int graphSize) {
+    return expectedVisitedNodes(k, graphSize) < (graphSize >>> 7)
+        ? new SparseFixedBitSet(graphSize)
+        : new FixedBitSet(graphSize);
+  }
+
   /**
    * Creates a new graph searcher.
    *
@@ -117,7 +144,7 @@ public static void search(
       innerSearcher =
           new HnswGraphSearcher(
               new NeighborQueue(knnCollector.k(), true),
-              new SparseFixedBitSet(getGraphSize(graph)));
+              createBitSet(knnCollector.k(), getGraphSize(graph)));
     }
     // Then, check if we the search strategy is seeded
     final AbstractHnswGraphSearcher graphSearcher;