Bypass HNSW graph building for tiny segments (#14963)

This change avoids creating a HNSW graph if the segment is small (here we have taken the thresholdfor number of vectors  as `10000` based on the conversation [here](#13447 (comment)) for now).

Some of the points I'm not sure how we would want to go about :
- All the tests passes currently since the option to enable the optimization is `false` by default but setting it to `true` reveals some failing unit tests which inherently assumes that the HNSW graph is created and KNN search is triggered (do we have some idea of how to bypass those in some good clean way?)
- I understand we might want to always keep this optimization on (also less invasive change), but for now in this PR, I made it configurable and enabled it on the KNN format - just to be cautious (wasn't sure if it would not affect back-compact in some unknown way), but happy to make it as default behaviour

Author: shubhamvishu

lucene/CHANGES.txt

@@ -75,6 +75,8 @@ Optimizations
 * GITHUB#15343: Ensure that `AcceptDocs#cost()` only ever calls `BitSets#cardinality()`
   once per instance to avoid redundant computation. (Ben Trent)
 
+* GITHUB#14963: Bypass HNSW graph building for tiny segments. (Shubham Chaudhary, Ben Trent)
+
 Bug Fixes
 ---------------------
 * GITHUB#14161: PointInSetQuery's constructor now throws IllegalArgumentException

lucene/backward-codecs/src/test/org/apache/lucene/backward_codecs/lucene99/Lucene99RWV0HnswScalarQuantizationVectorsFormat.java

@@ -47,7 +47,8 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
         Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH,
         flatVectorsFormat.fieldsWriter(state),
         1,
-        null);
+        null,
+        0);
   }
 
   static class Lucene99RWScalarQuantizedFormat extends Lucene99ScalarQuantizedVectorsFormat {

lucene/backward-codecs/src/test/org/apache/lucene/backward_codecs/lucene99/Lucene99RWV1HnswScalarQuantizationVectorsFormat.java

@@ -57,6 +57,7 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
         Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH,
         flatVectorsFormat.fieldsWriter(state),
         1,
-        null);
+        null,
+        0);
   }
 }

lucene/core/src/java/org/apache/lucene/codecs/lucene102/Lucene102HnswBinaryQuantizedVectorsFormat.java

@@ -128,7 +128,8 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
         beamWidth,
         flatVectorsFormat.fieldsWriter(state),
         numMergeWorkers,
-        mergeExec);
+        mergeExec,
+        0);
   }
 
   @Override

lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104HnswScalarQuantizedVectorsFormat.java

@@ -19,6 +19,7 @@
 import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH;
 import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN;
 import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_NUM_MERGE_WORKER;
+import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.HNSW_GRAPH_THRESHOLD;
 import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.MAXIMUM_BEAM_WIDTH;
 import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.MAXIMUM_MAX_CONN;
 
@@ -60,6 +61,20 @@ public class Lucene104HnswScalarQuantizedVectorsFormat extends KnnVectorsFormat
   /** The format for storing, reading, merging vectors on disk */
   private final Lucene104ScalarQuantizedVectorsFormat flatVectorsFormat;
 
+  /**
+   * The threshold to use to bypass HNSW graph building for tiny segments in terms of k for a graph
+   * i.e. number of docs to match the query (default is {@link
+   * Lucene99HnswVectorsFormat#HNSW_GRAPH_THRESHOLD}).
+   *
+   * <ul>
+   *   <li>0 indicates that the graph is always built.
+   *   <li>greater than 0 indicates that the graph needs a certain number of nodes before it starts
+   *       building. See {@link Lucene99HnswVectorsFormat#HNSW_GRAPH_THRESHOLD} for details.
+   *   <li>Negative values aren't allowed.
+   * </ul>
+   */
+  private final int tinySegmentsThreshold;
+
   private final int numMergeWorkers;
   private final TaskExecutor mergeExec;
 
@@ -70,7 +85,8 @@ public Lucene104HnswScalarQuantizedVectorsFormat() {
         DEFAULT_MAX_CONN,
         DEFAULT_BEAM_WIDTH,
         DEFAULT_NUM_MERGE_WORKER,
-        null);
+        null,
+        HNSW_GRAPH_THRESHOLD);
   }
 
   /**
@@ -80,7 +96,13 @@ public Lucene104HnswScalarQuantizedVectorsFormat() {
    * @param beamWidth the size of the queue maintained during graph construction.
    */
   public Lucene104HnswScalarQuantizedVectorsFormat(int maxConn, int beamWidth) {
-    this(ScalarEncoding.UNSIGNED_BYTE, maxConn, beamWidth, DEFAULT_NUM_MERGE_WORKER, null);
+    this(
+        ScalarEncoding.UNSIGNED_BYTE,
+        maxConn,
+        beamWidth,
+        DEFAULT_NUM_MERGE_WORKER,
+        null,
+        HNSW_GRAPH_THRESHOLD);
   }
 
   /**
@@ -99,6 +121,26 @@ public Lucene104HnswScalarQuantizedVectorsFormat(
       int beamWidth,
       int numMergeWorkers,
       ExecutorService mergeExec) {
+    this(encoding, maxConn, beamWidth, numMergeWorkers, mergeExec, HNSW_GRAPH_THRESHOLD);
+  }
+
+  /**
+   * Constructs a format using the given graph construction parameters and scalar quantization.
+   *
+   * @param maxConn the maximum number of connections to a node in the HNSW graph
+   * @param beamWidth the size of the queue maintained during graph construction.
+   * @param numMergeWorkers number of workers (threads) that will be used when doing merge. If
+   *     larger than 1, a non-null {@link ExecutorService} must be passed as mergeExec
+   * @param mergeExec the {@link ExecutorService} that will be used by ALL vector writers that are
+   *     generated by this format to do the merge
+   */
+  public Lucene104HnswScalarQuantizedVectorsFormat(
+      ScalarEncoding encoding,
+      int maxConn,
+      int beamWidth,
+      int numMergeWorkers,
+      ExecutorService mergeExec,
+      int tinySegmentsThreshold) {
     super(NAME);
     flatVectorsFormat = new Lucene104ScalarQuantizedVectorsFormat(encoding);
     if (maxConn <= 0 || maxConn > MAXIMUM_MAX_CONN) {
@@ -117,6 +159,7 @@ public Lucene104HnswScalarQuantizedVectorsFormat(
     }
     this.maxConn = maxConn;
     this.beamWidth = beamWidth;
+    this.tinySegmentsThreshold = tinySegmentsThreshold;
     if (numMergeWorkers == 1 && mergeExec != null) {
       throw new IllegalArgumentException(
           "No executor service is needed as we'll use single thread to merge");
@@ -137,7 +180,8 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
         beamWidth,
         flatVectorsFormat.fieldsWriter(state),
         numMergeWorkers,
-        mergeExec);
+        mergeExec,
+        tinySegmentsThreshold);
   }
 
   @Override
@@ -156,6 +200,8 @@ public String toString() {
         + maxConn
         + ", beamWidth="
         + beamWidth
+        + ", tinySegmentsThreshold="
+        + tinySegmentsThreshold
         + ", flatVectorFormat="
         + flatVectorsFormat
         + ")";

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

@@ -116,6 +116,23 @@ public final class Lucene99HnswVectorsFormat extends KnnVectorsFormat {
   /** Default to use single thread merge */
   public static final int DEFAULT_NUM_MERGE_WORKER = 1;
 
+  /**
+   * Minimum estimated search effort (in terms of expected visited nodes) required before building
+   * an HNSW graph for a segment.
+   *
+   * <p>This threshold is compared against the value produced by {@link
+   * org.apache.lucene.util.hnsw.HnswGraphSearcher#expectedVisitedNodes(int, int)}, which estimates
+   * how many nodes would be visited during a vector search based on the current graph size and
+   * {@code k} (neighbours to find).
+   *
+   * <p>If the estimated number of visited nodes falls below this threshold, HNSW graph construction
+   * is skipped for that segment - typically for small flushes or low document count segments -
+   * since the overhead of building the graph would outweigh its search benefits.
+   *
+   * <p>Default: {@code 100}
+   */
+  public static final int HNSW_GRAPH_THRESHOLD = 100;
+
   static final int DIRECT_MONOTONIC_BLOCK_SHIFT = 16;
 
   /**
@@ -138,11 +155,30 @@ public final class Lucene99HnswVectorsFormat extends KnnVectorsFormat {
   private final int numMergeWorkers;
   private final TaskExecutor mergeExec;
 
+  /**
+   * The threshold to use to bypass HNSW graph building for tiny segments in terms of k for a graph
+   * i.e. number of docs to match the query (default is {@link
+   * Lucene99HnswVectorsFormat#HNSW_GRAPH_THRESHOLD}).
+   *
+   * <ul>
+   *   <li>0 indicates that the graph is always built.
+   *   <li>0 indicates that the graph needs certain or more nodes before it starts building.
+   *   <li>Negative values aren't allowed.
+   * </ul>
+   */
+  private final int tinySegmentsThreshold;
+
   private final int writeVersion;
 
   /** Constructs a format using default graph construction parameters */
   public Lucene99HnswVectorsFormat() {
-    this(DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH, DEFAULT_NUM_MERGE_WORKER, null);
+    this(
+        DEFAULT_MAX_CONN,
+        DEFAULT_BEAM_WIDTH,
+        DEFAULT_NUM_MERGE_WORKER,
+        null,
+        HNSW_GRAPH_THRESHOLD,
+        VERSION_CURRENT);
   }
 
   /**
@@ -152,7 +188,20 @@ public Lucene99HnswVectorsFormat() {
    * @param beamWidth the size of the queue maintained during graph construction.
    */
   public Lucene99HnswVectorsFormat(int maxConn, int beamWidth) {
-    this(maxConn, beamWidth, DEFAULT_NUM_MERGE_WORKER, null);
+    this(maxConn, beamWidth, DEFAULT_NUM_MERGE_WORKER, null, HNSW_GRAPH_THRESHOLD, VERSION_CURRENT);
+  }
+
+  /**
+   * Constructs a format using the given graph construction parameters.
+   *
+   * @param maxConn the maximum number of connections to a node in the HNSW graph
+   * @param beamWidth the size of the queue maintained during graph construction.
+   * @param tinySegmentsThreshold the expected number of vector operations to return k nearest
+   *     neighbors of the current graph size
+   */
+  public Lucene99HnswVectorsFormat(int maxConn, int beamWidth, int tinySegmentsThreshold) {
+    this(
+        maxConn, beamWidth, DEFAULT_NUM_MERGE_WORKER, null, tinySegmentsThreshold, VERSION_CURRENT);
   }
 
   /**
@@ -168,7 +217,7 @@ public Lucene99HnswVectorsFormat(int maxConn, int beamWidth) {
    */
   public Lucene99HnswVectorsFormat(
       int maxConn, int beamWidth, int numMergeWorkers, ExecutorService mergeExec) {
-    this(maxConn, beamWidth, numMergeWorkers, mergeExec, VERSION_CURRENT);
+    this(maxConn, beamWidth, numMergeWorkers, mergeExec, HNSW_GRAPH_THRESHOLD, VERSION_CURRENT);
   }
 
   /**
@@ -182,13 +231,38 @@ public Lucene99HnswVectorsFormat(
    * @param mergeExec the {@link ExecutorService} that will be used by ALL vector writers that are
    *     generated by this format to do the merge. If null, the configured {@link
    *     MergeScheduler#getIntraMergeExecutor(MergePolicy.OneMerge)} is used.
+   * @param tinySegmentsThreshold the expected number of vector operations to return k nearest
+   *     neighbors of the current graph size
+   */
+  public Lucene99HnswVectorsFormat(
+      int maxConn,
+      int beamWidth,
+      int numMergeWorkers,
+      ExecutorService mergeExec,
+      int tinySegmentsThreshold) {
+    this(maxConn, beamWidth, numMergeWorkers, mergeExec, tinySegmentsThreshold, VERSION_CURRENT);
+  }
+
+  /**
+   * Constructs a format using the given graph construction parameters.
+   *
+   * @param maxConn the maximum number of connections to a node in the HNSW graph
+   * @param beamWidth the size of the queue maintained during graph construction.
+   * @param numMergeWorkers number of workers (threads) that will be used when doing merge. If
+   *     larger than 1, a non-null {@link ExecutorService} must be passed as mergeExec
+   * @param mergeExec the {@link ExecutorService} that will be used by ALL vector writers that are
+   *     generated by this format to do the merge. If null, the configured {@link
+   *     MergeScheduler#getIntraMergeExecutor(MergePolicy.OneMerge)} is used.
+   * @param tinySegmentsThreshold the expected number of vector operations to return k nearest
+   *     neighbors of the current graph size
    * @param writeVersion the version used for the writer to encode docID's (VarInt=0, GroupVarInt=1)
    */
   Lucene99HnswVectorsFormat(
       int maxConn,
       int beamWidth,
       int numMergeWorkers,
       ExecutorService mergeExec,
+      int tinySegmentsThreshold,
       int writeVersion) {
     super("Lucene99HnswVectorsFormat");
     if (maxConn <= 0 || maxConn > MAXIMUM_MAX_CONN) {
@@ -207,6 +281,7 @@ public Lucene99HnswVectorsFormat(
     }
     this.maxConn = maxConn;
     this.beamWidth = beamWidth;
+    this.tinySegmentsThreshold = tinySegmentsThreshold;
     this.writeVersion = writeVersion;
     if (numMergeWorkers == 1 && mergeExec != null) {
       throw new IllegalArgumentException(
@@ -229,6 +304,7 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
         flatVectorsFormat.fieldsWriter(state),
         numMergeWorkers,
         mergeExec,
+        tinySegmentsThreshold,
         writeVersion);
   }
 
@@ -248,6 +324,8 @@ public String toString() {
         + maxConn
         + ", beamWidth="
         + beamWidth
+        + ", tinySegmentsThreshold="
+        + tinySegmentsThreshold
         + ", flatVectorFormat="
         + flatVectorsFormat
         + ")";

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

@@ -338,18 +338,18 @@ private void search(
     final RandomVectorScorer scorer = scorerSupplier.get();
     final KnnCollector collector =
         new OrdinalTranslatedKnnCollector(knnCollector, scorer::ordToDoc);
-    HnswGraph graph = getGraph(fieldEntry);
     // Take into account if quantized? E.g. some scorer cost?
     // Use approximate cardinality as this is good enough, but ensure we don't exceed the graph
     // size as that is illogical
+    HnswGraph graph = getGraph(fieldEntry);
     int filteredDocCount = Math.min(acceptDocs.cost(), graph.size());
     Bits accepted = acceptDocs.bits();
     final Bits acceptedOrds = scorer.getAcceptOrds(accepted);
     int numVectors = scorer.maxOrd();
     boolean doHnsw = knnCollector.k() < numVectors;
     // The approximate number of vectors that would be visited if we did not filter
     int unfilteredVisit = HnswGraphSearcher.expectedVisitedNodes(knnCollector.k(), graph.size());
-    if (unfilteredVisit >= filteredDocCount) {
+    if (unfilteredVisit >= filteredDocCount || graph.size() == 0) {
       doHnsw = false;
     }
     if (doHnsw) {
@@ -405,6 +405,9 @@ public HnswGraph getGraph(String field) throws IOException {
   }
 
   private HnswGraph getGraph(FieldEntry entry) throws IOException {
+    if (entry.vectorIndexLength == 0) {
+      return HnswGraph.EMPTY;
+    }
     return new OffHeapHnswGraph(entry, vectorIndex);
   }