Address some comments

Author: shubhamvishu

lucene/core/src/java/org/apache/lucene/codecs/lucene102/Lucene102HnswBinaryQuantizedVectorsFormat.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;
 
@@ -64,9 +65,17 @@ public class Lucene102HnswBinaryQuantizedVectorsFormat 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}).
+   */
+  private final int tinySegmentsThreshold;
+
   /** Constructs a format using default graph construction parameters */
   public Lucene102HnswBinaryQuantizedVectorsFormat() {
-    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);
   }
 
   /**
@@ -76,7 +85,20 @@ public Lucene102HnswBinaryQuantizedVectorsFormat() {
    * @param beamWidth the size of the queue maintained during graph construction.
    */
   public Lucene102HnswBinaryQuantizedVectorsFormat(int maxConn, int beamWidth) {
-    this(maxConn, beamWidth, DEFAULT_NUM_MERGE_WORKER, null);
+    this(maxConn, beamWidth, DEFAULT_NUM_MERGE_WORKER, null, HNSW_GRAPH_THRESHOLD);
+  }
+
+  /**
+   * 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 value of k for the expectedVisitedNodes heuristic, used to
+   *     determine the minimum required graph nodes
+   */
+  public Lucene102HnswBinaryQuantizedVectorsFormat(
+      int maxConn, int beamWidth, int tinySegmentsThreshold) {
+    this(maxConn, beamWidth, DEFAULT_NUM_MERGE_WORKER, null, tinySegmentsThreshold);
   }
 
   /**
@@ -91,6 +113,27 @@ public Lucene102HnswBinaryQuantizedVectorsFormat(int maxConn, int beamWidth) {
    */
   public Lucene102HnswBinaryQuantizedVectorsFormat(
       int maxConn, int beamWidth, int numMergeWorkers, ExecutorService mergeExec) {
+    this(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
+   * @param tinySegmentsThreshold the value of k for the expectedVisitedNodes heuristic, used to
+   *     determine the minimum required graph nodes
+   */
+  public Lucene102HnswBinaryQuantizedVectorsFormat(
+      int maxConn,
+      int beamWidth,
+      int numMergeWorkers,
+      ExecutorService mergeExec,
+      int tinySegmentsThreshold) {
     super(NAME);
     if (maxConn <= 0 || maxConn > MAXIMUM_MAX_CONN) {
       throw new IllegalArgumentException(
@@ -107,6 +150,7 @@ public Lucene102HnswBinaryQuantizedVectorsFormat(
               + beamWidth);
     }
     this.maxConn = maxConn;
+    this.tinySegmentsThreshold = tinySegmentsThreshold;
     this.beamWidth = beamWidth;
     if (numMergeWorkers == 1 && mergeExec != null) {
       throw new IllegalArgumentException(
@@ -128,12 +172,14 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
         beamWidth,
         flatVectorsFormat.fieldsWriter(state),
         numMergeWorkers,
-        mergeExec);
+        mergeExec,
+        tinySegmentsThreshold);
   }
 
   @Override
   public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException {
-    return new Lucene99HnswVectorsReader(state, flatVectorsFormat.fieldsReader(state));
+    return new Lucene99HnswVectorsReader(
+        state, flatVectorsFormat.fieldsReader(state), tinySegmentsThreshold);
   }
 
   @Override
@@ -147,6 +193,8 @@ public String toString() {
         + maxConn
         + ", beamWidth="
         + beamWidth
+        + ", tinySegmentsThreshold="
+        + tinySegmentsThreshold
         + ", flatVectorFormat="
         + flatVectorsFormat
         + ")";

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

@@ -20,6 +20,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;
 
@@ -66,12 +67,11 @@ public class Lucene99HnswScalarQuantizedVectorsFormat extends KnnVectorsFormat {
   private final TaskExecutor mergeExec;
 
   /**
-   * Whether to bypass HNSW graph building for tiny segments (below {@link
-   * Lucene99HnswVectorsFormat#HNSW_GRAPH_THRESHOLD}). When enabled, segments with fewer than the
-   * threshold number of vectors will store only flat vectors with quantization, significantly
-   * improving indexing performance for workloads with frequent flushes.
+   * 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}).
    */
-  private final boolean bypassTinySegments;
+  private final int tinySegmentsThreshold;
 
   /** Constructs a format using default graph construction parameters with 7 bit quantization */
   public Lucene99HnswScalarQuantizedVectorsFormat() {
@@ -83,7 +83,7 @@ public Lucene99HnswScalarQuantizedVectorsFormat() {
         false,
         null,
         null,
-        false);
+        HNSW_GRAPH_THRESHOLD);
   }
 
   /**
@@ -93,19 +93,20 @@ public Lucene99HnswScalarQuantizedVectorsFormat() {
    * @param beamWidth the size of the queue maintained during graph construction.
    */
   public Lucene99HnswScalarQuantizedVectorsFormat(int maxConn, int beamWidth) {
-    this(maxConn, beamWidth, DEFAULT_NUM_MERGE_WORKER, 7, false, null, null, false);
+    this(maxConn, beamWidth, DEFAULT_NUM_MERGE_WORKER, 7, false, null, null, HNSW_GRAPH_THRESHOLD);
   }
 
   /**
    * Constructs a format using the given graph construction parameters with 7 bit 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 bypassTinySegments whether to bypass HNSW graph building for tiny segments
+   * @param tinySegmentsThreshold the value of k for the expectedVisitedNodes heuristic, used to
+   *     determine the minimum required graph nodes
    */
   public Lucene99HnswScalarQuantizedVectorsFormat(
-      int maxConn, int beamWidth, boolean bypassTinySegments) {
-    this(maxConn, beamWidth, DEFAULT_NUM_MERGE_WORKER, 7, false, null, null, bypassTinySegments);
+      int maxConn, int beamWidth, int tinySegmentsThreshold) {
+    this(maxConn, beamWidth, DEFAULT_NUM_MERGE_WORKER, 7, false, null, null, tinySegmentsThreshold);
   }
 
   /**
@@ -135,7 +136,15 @@ public Lucene99HnswScalarQuantizedVectorsFormat(
       boolean compress,
       Float confidenceInterval,
       ExecutorService mergeExec) {
-    this(maxConn, beamWidth, numMergeWorkers, bits, compress, confidenceInterval, mergeExec, false);
+    this(
+        maxConn,
+        beamWidth,
+        numMergeWorkers,
+        bits,
+        compress,
+        confidenceInterval,
+        mergeExec,
+        HNSW_GRAPH_THRESHOLD);
   }
 
   /**
@@ -156,9 +165,8 @@ public Lucene99HnswScalarQuantizedVectorsFormat(
    *     accurate pair.
    * @param mergeExec the {@link ExecutorService} that will be used by ALL vector writers that are
    *     generated by this format to do the merge
-   * @param bypassTinySegments whether to bypass HNSW graph building for tiny segments (below {@link
-   *     Lucene99HnswVectorsFormat#HNSW_GRAPH_THRESHOLD} vectors). When enabled, improves indexing
-   *     performance for workloads with frequent flushes.
+   * @param tinySegmentsThreshold the value of k for the expectedVisitedNodes heuristic, used to
+   *     determine the minimum required graph nodes
    */
   public Lucene99HnswScalarQuantizedVectorsFormat(
       int maxConn,
@@ -168,7 +176,7 @@ public Lucene99HnswScalarQuantizedVectorsFormat(
       boolean compress,
       Float confidenceInterval,
       ExecutorService mergeExec,
-      boolean bypassTinySegments) {
+      int tinySegmentsThreshold) {
     super(NAME);
     if (maxConn <= 0 || maxConn > MAXIMUM_MAX_CONN) {
       throw new IllegalArgumentException(
@@ -186,7 +194,7 @@ public Lucene99HnswScalarQuantizedVectorsFormat(
     }
     this.maxConn = maxConn;
     this.beamWidth = beamWidth;
-    this.bypassTinySegments = bypassTinySegments;
+    this.tinySegmentsThreshold = tinySegmentsThreshold;
     if (numMergeWorkers == 1 && mergeExec != null) {
       throw new IllegalArgumentException(
           "No executor service is needed as we'll use single thread to merge");
@@ -210,12 +218,13 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
         flatVectorsFormat.fieldsWriter(state),
         numMergeWorkers,
         mergeExec,
-        bypassTinySegments);
+        tinySegmentsThreshold);
   }
 
   @Override
   public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException {
-    return new Lucene99HnswVectorsReader(state, flatVectorsFormat.fieldsReader(state));
+    return new Lucene99HnswVectorsReader(
+        state, flatVectorsFormat.fieldsReader(state), tinySegmentsThreshold);
   }
 
   @Override
@@ -229,8 +238,8 @@ public String toString() {
         + maxConn
         + ", beamWidth="
         + beamWidth
-        + ", bypassTinySegments="
-        + bypassTinySegments
+        + ", tinySegmentsThreshold="
+        + tinySegmentsThreshold
         + ", flatVectorFormat="
         + flatVectorsFormat
         + ")";

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

@@ -116,11 +116,16 @@ public final class Lucene99HnswVectorsFormat extends KnnVectorsFormat {
   public static final int DEFAULT_NUM_MERGE_WORKER = 1;
 
   /**
-   * Threshold below which HNSW graph building is bypassed for tiny segments. Segments with fewer
-   * vectors will use flat storage only, improving indexing performance when having frequent
-   * flushes.
+   * Threshold which dynamically uses the HnswGraphSearcher#expectedVisitedNodes heuristic to find
+   * the threshold below which HNSW graph building is bypassed. flushes. It is in terms of k for a
+   * graph i.e. number of docs to match the query. So,
+   *
+   * <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 = 10_000;
+  public static final int HNSW_GRAPH_THRESHOLD = 10;
 
   static final int DIRECT_MONOTONIC_BLOCK_SHIFT = 16;
 
@@ -145,15 +150,16 @@ public final class Lucene99HnswVectorsFormat extends KnnVectorsFormat {
   private final TaskExecutor mergeExec;
 
   /**
-   * Whether to bypass HNSW graph building for tiny segments (below {@link #HNSW_GRAPH_THRESHOLD}).
-   * When enabled, segments with fewer than the threshold number of vectors will store only flat
-   * vectors, significantly improving indexing performance for workloads with frequent flushes.
+   * 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}).
    */
-  private final boolean bypassTinySegments;
+  private final int tinySegmentsThreshold;
 
   /** Constructs a format using default graph construction parameters */
   public Lucene99HnswVectorsFormat() {
-    this(DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH, DEFAULT_NUM_MERGE_WORKER, null, false);
+    this(
+        DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH, DEFAULT_NUM_MERGE_WORKER, null, HNSW_GRAPH_THRESHOLD);
   }
 
   /**
@@ -163,18 +169,19 @@ 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, false);
+    this(maxConn, beamWidth, DEFAULT_NUM_MERGE_WORKER, null, HNSW_GRAPH_THRESHOLD);
   }
 
   /**
    * 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 bypassTinySegments whether to bypass HNSW graph building for tiny segments
+   * @param tinySegmentsThreshold the value of k for the expectedVisitedNodes heuristic, used to
+   *     determine the minimum required graph nodes
    */
-  public Lucene99HnswVectorsFormat(int maxConn, int beamWidth, boolean bypassTinySegments) {
-    this(maxConn, beamWidth, DEFAULT_NUM_MERGE_WORKER, null, bypassTinySegments);
+  public Lucene99HnswVectorsFormat(int maxConn, int beamWidth, int tinySegmentsThreshold) {
+    this(maxConn, beamWidth, DEFAULT_NUM_MERGE_WORKER, null, tinySegmentsThreshold);
   }
 
   /**
@@ -190,7 +197,7 @@ public Lucene99HnswVectorsFormat(int maxConn, int beamWidth, boolean bypassTinyS
    */
   public Lucene99HnswVectorsFormat(
       int maxConn, int beamWidth, int numMergeWorkers, ExecutorService mergeExec) {
-    this(maxConn, beamWidth, numMergeWorkers, mergeExec, false);
+    this(maxConn, beamWidth, numMergeWorkers, mergeExec, HNSW_GRAPH_THRESHOLD);
   }
 
   /**
@@ -203,16 +210,15 @@ 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 bypassTinySegments whether to bypass HNSW graph building for tiny segments (below {@link
-   *     #HNSW_GRAPH_THRESHOLD} vectors). When enabled, improves indexing performance for workloads
-   *     with frequent flushes.
+   * @param tinySegmentsThreshold the value of k for the expectedVisitedNodes heuristic, used to
+   *     determine the minimum required graph nodes
    */
   public Lucene99HnswVectorsFormat(
       int maxConn,
       int beamWidth,
       int numMergeWorkers,
       ExecutorService mergeExec,
-      boolean bypassTinySegments) {
+      int tinySegmentsThreshold) {
     super("Lucene99HnswVectorsFormat");
     if (maxConn <= 0 || maxConn > MAXIMUM_MAX_CONN) {
       throw new IllegalArgumentException(
@@ -230,7 +236,7 @@ public Lucene99HnswVectorsFormat(
     }
     this.maxConn = maxConn;
     this.beamWidth = beamWidth;
-    this.bypassTinySegments = bypassTinySegments;
+    this.tinySegmentsThreshold = tinySegmentsThreshold;
     if (numMergeWorkers == 1 && mergeExec != null) {
       throw new IllegalArgumentException(
           "No executor service is needed as we'll use single thread to merge");
@@ -252,13 +258,13 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
         flatVectorsFormat.fieldsWriter(state),
         numMergeWorkers,
         mergeExec,
-        bypassTinySegments);
+        tinySegmentsThreshold);
   }
 
   @Override
   public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException {
     return new Lucene99HnswVectorsReader(
-        state, flatVectorsFormat.fieldsReader(state), bypassTinySegments);
+        state, flatVectorsFormat.fieldsReader(state), tinySegmentsThreshold);
   }
 
   @Override
@@ -272,8 +278,8 @@ public String toString() {
         + maxConn
         + ", beamWidth="
         + beamWidth
-        + ", bypassTinySegments="
-        + bypassTinySegments
+        + ", tinySegmentsThreshold="
+        + tinySegmentsThreshold
         + ", flatVectorFormat="
         + flatVectorsFormat
         + ")";