Bypass HNSW graph building for tiny segments

lucene/CHANGES.txt

@@ -200,6 +200,8 @@ Optimizations
 
 * GITHUB#15261: Implement longValues for MultiFieldNormValues to speedup CombinedQuery (Ge Song)
 
+* GITHUB#14963: Bypass HNSW graph building for tiny segments. (Shubham Chaudhary, Benjamin 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,19 @@ 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>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 numMergeWorkers;
   private final TaskExecutor mergeExec;
 
@@ -70,7 +84,8 @@ public Lucene104HnswScalarQuantizedVectorsFormat() {
         DEFAULT_MAX_CONN,
         DEFAULT_BEAM_WIDTH,
         DEFAULT_NUM_MERGE_WORKER,
-        null);
+        null,
+        HNSW_GRAPH_THRESHOLD);
   }
 
   /**
@@ -80,7 +95,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);
   }
 
   /**
@@ -100,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) {
@@ -118,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");
@@ -138,7 +180,8 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
         beamWidth,
         flatVectorsFormat.fieldsWriter(state),
         numMergeWorkers,
-        mergeExec);
+        mergeExec,
+        tinySegmentsThreshold);
   }
 
   @Override
@@ -157,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,18 @@ public final class Lucene99HnswVectorsFormat extends KnnVectorsFormat {
   /** Default to use single thread merge */
   public static final int DEFAULT_NUM_MERGE_WORKER = 1;
 
+  /**
+   * Threshold which HnswGraphSearcher#expectedVisitedNodes uses as k to determine when HNSW graph
+   * building is bypassed (useful in case if frequent flushes). It is in terms of k for a graph i.e.
+   * number of docs to match for the query. So having a graph only helps if,
+   *
+   * <pre> k &lt;&lt; size / log(size) </pre>
+   *
+   * i.e. k is at least 1 order less than size / log(size) where size if the number of nodes in the
+   * graph
+   */
+  public static final int HNSW_GRAPH_THRESHOLD = 100;
+
   static final int DIRECT_MONOTONIC_BLOCK_SHIFT = 16;
 
   /**
@@ -138,11 +150,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 +183,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 +212,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 +226,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 +276,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 +299,7 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
         flatVectorsFormat.fieldsWriter(state),
         numMergeWorkers,
         mergeExec,
+        tinySegmentsThreshold,
         writeVersion);
   }
 
@@ -248,6 +319,8 @@ public String toString() {
         + maxConn
         + ", beamWidth="
         + beamWidth
+        + ", tinySegmentsThreshold="
+        + tinySegmentsThreshold
         + ", flatVectorFormat="
         + flatVectorsFormat
         + ")";

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

@@ -332,18 +332,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) {
@@ -399,6 +399,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);
   }