Optimize cuvs-lucene CAGRA_HNSW index build; add example

java/cuvs-lucene/examples/README.md

@@ -32,3 +32,17 @@ To run the Index and Search on GPU example do:
 ```sh
 mvn clean install && java -Djava.util.logging.config.file=src/main/resources/logging.properties -cp target/examples-26.10.0-jar-with-merged-services.jar com.nvidia.cuvs.lucene.examples.IndexAndSearchonGPUExample
 ```
+
+To run the optimized CAGRA-HNSW build example (reference pattern for efficiently building an
+accelerated HNSW index from a large `.fbin` with every ingest-side knob on — open the file once and
+stream sequential prefetched chunks that overlap the disk read with indexing, hold at most two chunks
+in memory, reuse a single vector array, size a native flat buffer per segment, auto-select the CAGRA
+graph-build algorithm, and optionally partition into K segments built sequentially or overlapped) do:
+
+```sh
+mvn clean install && java -Djava.util.logging.config.file=src/main/resources/logging.properties -cp target/examples-26.10.0-jar-with-merged-services.jar com.nvidia.cuvs.lucene.examples.OptimizedCagraHnswBuildExample
+```
+
+With no arguments it generates and indexes a small demo `.fbin` as a single segment; pass a real file,
+chunk size, segment count, and overlap flag as
+`... OptimizedCagraHnswBuildExample <path-to.fbin> <chunkSizeMB> <numSegments> <overlap:true|false>`.

java/cuvs-lucene/examples/pom.xml

@@ -92,6 +92,11 @@
         <artifactId>lucene-backward-codecs</artifactId>
         <version>10.2.0</version>
     </dependency>
+    <dependency>
+        <groupId>org.apache.lucene</groupId>
+        <artifactId>lucene-misc</artifactId>
+        <version>10.2.0</version>
+    </dependency>
     <dependency>
         <groupId>commons-io</groupId>
         <artifactId>commons-io</artifactId>

java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java

@@ -59,13 +59,14 @@ public static enum Strategy {
   public static final int DEFAULT_MAX_CONN = 32;
   public static final int DEFAULT_BEAM_WIDTH = 32;
   public static final CagraGraphBuildAlgo DEFAULT_CAGRA_GRAPH_BUILD_ALGO =
-      CagraGraphBuildAlgo.NN_DESCENT;
+      CagraGraphBuildAlgo.AUTO_SELECT;
   public static final int DEFAULT_NUM_MERGE_WORKERS = 1;
   public static final Strategy DEFAULT_STRATEGY = Strategy.HEURISTIC;
   public static final CuvsDistanceType DEFAULT_CUVS_DISTANCE_TYPE = CuvsDistanceType.L2Expanded;
   public static final int DEFAULT_NN_DESCENT_NUM_ITERATIONS = 20;
   public static final HnswHeuristicType DEFAULT_HNSW_HEURISTIC_TYPE =
       HnswHeuristicType.SAME_GRAPH_FOOTPRINT;
+  public static final int DEFAULT_NUM_INPUT_VECTORS = 0;
 
   public static final Supplier<CuVSIvfPqParams> DEFAULT_IVF_PQ_PARAMS =
       () -> {
@@ -91,24 +92,31 @@ public static enum Strategy {
   private final CuvsDistanceType cuvsDistanceType;
   private final int nnDescentNumIterations;
   private final HnswHeuristicType hnswHeuristicType;
+  private final int numInputVectors;
 
   /**
    * Constructs an instance of {@link AcceleratedHNSWParams} with specific parameter values.
    *
    * @param writerThreads Number of cuVS writer threads to use.
    * @param intermediateGraphDegree The intermediate graph degree while building the CAGRA index.
-   * @param graphdegree The graph degree to use while building the CAGRA index.
+   *     Only consulted under the {@link Strategy#CUSTOM} strategy.
+   * @param graphdegree The graph degree to use while building the CAGRA index. Only consulted
+   *     under the {@link Strategy#CUSTOM} strategy.
    * @param hnswLayers The number of HNSW layers to build in the HNSW index.
    * @param maxConn The max connection parameter used when building HNSW index with the fallback mechanism.
    * @param beamWidth The beam width parameter used when building HNSW index with the fallback mechanism.
-   * @param cagraGraphBuildAlgo The CAGRA graph build algorithm to use [NN_DESCENT, IVF_PQ].
+   * @param cagraGraphBuildAlgo The CAGRA graph build algorithm to use [NN_DESCENT, IVF_PQ]. Only
+   *     consulted under the {@link Strategy#CUSTOM} strategy.
    * @param cuVSIvfPqParams An instance of CuVSIvfPqParams containing IVF_PQ specific parameters.
+   *     Only consulted under the {@link Strategy#CUSTOM} strategy.
    * @param numMergeWorkers The number of merge workers to use with the fallback mechanism.
    * @param mergeExec The instance of {@link ExecutorService} to use with the fallback mechanism.
    * @param strategy either HEURISTIC [Default] that delegates the CAGRA build parameters to cuVS (derived from the HNSW-equivalent maxConn and beamWidth) or CUSTOM that uses the parameters passed through this class.
    * @param cuvsDistanceType the cuvsDistanceType. The default option is L2Expanded.
    * @param nnDescentNumIterations the number of Iterations to run if building with NN_DESCENT.
+   *     Only consulted under the {@link Strategy#CUSTOM} strategy.
    * @param hnswHeuristicType the heuristic cuVS applies when deriving the CAGRA build parameters from maxConn and beamWidth under the HEURISTIC strategy.
+   * @param numInputVectors exact number of vectors to be indexed, used to pre-size the native flat buffer (0 = disabled).
    */
   private AcceleratedHNSWParams(
       int writerThreads,
@@ -124,7 +132,8 @@ private AcceleratedHNSWParams(
       Strategy strategy,
       CuvsDistanceType cuvsDistanceType,
       int nnDescentNumIterations,
-      HnswHeuristicType hnswHeuristicType) {
+      HnswHeuristicType hnswHeuristicType,
+      int numInputVectors) {
     super();
     this.writerThreads = writerThreads;
     this.intermediateGraphDegree = intermediateGraphDegree;
@@ -140,6 +149,7 @@ private AcceleratedHNSWParams(
     this.cuvsDistanceType = cuvsDistanceType;
     this.nnDescentNumIterations = nnDescentNumIterations;
     this.hnswHeuristicType = hnswHeuristicType;
+    this.numInputVectors = numInputVectors;
   }
 
   /**
@@ -152,7 +162,7 @@ public int getWriterThreads() {
   }
 
   /**
-   * Get the intermediate graph degree
+   * Get the intermediate graph degree. Only consulted under the {@link Strategy#CUSTOM} strategy.
    *
    * @return the graph degree parameter
    */
@@ -161,7 +171,7 @@ public int getIntermediateGraphDegree() {
   }
 
   /**
-   * Get the graph degree
+   * Get the graph degree. Only consulted under the {@link Strategy#CUSTOM} strategy.
    *
    * @return the graph degree parameter
    */
@@ -197,7 +207,8 @@ public int getBeamWidth() {
   }
 
   /**
-   * Get the CAGRA graph build algorithm
+   * Get the CAGRA graph build algorithm. Only consulted under the {@link Strategy#CUSTOM}
+   * strategy; under {@link Strategy#HEURISTIC} the algorithm is chosen by cuVS.
    *
    * @return the CAGRA graph build algorithm
    */
@@ -206,7 +217,8 @@ public CagraGraphBuildAlgo getCagraGraphBuildAlgo() {
   }
 
   /**
-   * Get the instance of {@link CuVSIvfPqParams}
+   * Get the instance of {@link CuVSIvfPqParams}. Only consulted under the {@link
+   * Strategy#CUSTOM} strategy.
    *
    * @return the instance of {@link CuVSIvfPqParams}
    */
@@ -254,7 +266,8 @@ public CuvsDistanceType getCuvsDistanceType() {
   }
 
   /**
-   * get the number of Iterations to run if building with NN_DESCENT
+   * get the number of Iterations to run if building with NN_DESCENT. Only consulted under the
+   * {@link Strategy#CUSTOM} strategy.
    *
    * @return the number of iterations for NN_DESCENT
    */
@@ -272,6 +285,17 @@ public HnswHeuristicType getHnswHeuristicType() {
     return hnswHeuristicType;
   }
 
+  /**
+   * Get the number of input vectors used to pre-size the native flat buffer. A value of
+   * {@value DEFAULT_NUM_INPUT_VECTORS} means unset (the writer uses the default heap-buffered
+   * flat path).
+   *
+   * @return the number of vectors to be indexed, or 0 if unset
+   */
+  public int getNumInputVectors() {
+    return numInputVectors;
+  }
+
   @Override
   public String toString() {
     return "AcceleratedHNSWParams [writerThreads="
@@ -302,6 +326,8 @@ public String toString() {
         + nnDescentNumIterations
         + ", hnswHeuristicType="
         + hnswHeuristicType
+        + ", numInputVectors="
+        + numInputVectors
         + "]";
   }
 
@@ -324,6 +350,7 @@ public static class Builder {
     private CuvsDistanceType cuvsDistanceType = DEFAULT_CUVS_DISTANCE_TYPE;
     private int nnDescentNumIterations = DEFAULT_NN_DESCENT_NUM_ITERATIONS;
     private HnswHeuristicType hnswHeuristicType = DEFAULT_HNSW_HEURISTIC_TYPE;
+    private int numInputVectors = DEFAULT_NUM_INPUT_VECTORS;
 
     /**
      * Set the number of cuVS writer threads while building the index
@@ -339,7 +366,8 @@ public Builder withWriterThreads(int writerThreads) {
     }
 
     /**
-     * Set the intermediate graph degree to use while building CAGRA index
+     * Set the intermediate graph degree to use while building CAGRA index. Only consulted under
+     * the {@link Strategy#CUSTOM} strategy.
      * Valid range - Minimum: {@value MIN_INT_GRAPH_DEG}, Maximum: {@value MAX_INT_GRAPH_DEG}
      * Default value - {@value DEFAULT_INT_GRAPH_DEGREE}
      *
@@ -352,7 +380,8 @@ public Builder withIntermediateGraphDegree(int intermediateGraphDegree) {
     }
 
     /**
-     * Set the graph degree to use while building CAGRA index
+     * Set the graph degree to use while building CAGRA index. Only consulted under the {@link
+     * Strategy#CUSTOM} strategy.
      * Valid range - Minimum: {@value MIN_GRAPH_DEG}, Maximum: {@value MAX_GRAPH_DEG}
      * Default value - {@value DEFAULT_GRAPH_DEGREE}
      *
@@ -404,8 +433,9 @@ public Builder withBeamWidth(int beamWidth) {
     }
 
     /**
-     * Set the CAGRA graph build algorithm to use
-     * Default value - NN_DESCENT
+     * Set the CAGRA graph build algorithm to use. Only consulted under the {@link
+     * Strategy#CUSTOM} strategy; under {@link Strategy#HEURISTIC} the algorithm is chosen by cuVS.
+     * Default value - AUTO_SELECT
      *
      * @param cagraGraphBuildAlgo
      * @return instance of {@link Builder}
@@ -416,7 +446,8 @@ public Builder withCagraGraphBuildAlgo(CagraGraphBuildAlgo cagraGraphBuildAlgo)
     }
 
     /**
-     * Set the instance of {@link CuVSIvfPqParams}
+     * Set the instance of {@link CuVSIvfPqParams}. Only consulted under the {@link
+     * Strategy#CUSTOM} strategy.
      *
      * @param cuVSIvfPqParams
      * @return instance of {@link Builder}
@@ -479,7 +510,8 @@ public Builder withCuvsDistanceType(CuvsDistanceType cuvsDistanceType) {
     }
 
     /**
-     * Set the number of Iterations to run if building with NN_DESCENT
+     * Set the number of Iterations to run if building with NN_DESCENT. Only consulted under the
+     * {@link Strategy#CUSTOM} strategy.
      *
      * Valid range - Minimum: {@value MIN_NN_DESCENT_NUM_ITERATIONS}, Maximum: {@value MAX_NN_DESCENT_NUM_ITERATIONS}
      * Default value - {@value DEFAULT_NN_DESCENT_NUM_ITERATIONS}
@@ -507,6 +539,40 @@ public Builder withHnswHeuristicType(HnswHeuristicType hnswHeuristicType) {
       return this;
     }
 
+    /**
+     * Set the exact number of vectors to be indexed, used to pre-allocate a single contiguous
+     * native flat buffer (avoiding the on-heap {@code List<float[]>} and the extra host-matrix
+     * copy). The native buffer is sized for exactly this many rows, so the value MUST equal the
+     * number of vectors actually added; the writer fails fast otherwise. Only supported for the
+     * unsorted single-segment CAGRA_HNSW build (no merges). Not yet supported for the
+     * binary/scalar quantized writers. A value of {@value DEFAULT_NUM_INPUT_VECTORS} (the
+     * default) disables it and uses the default heap-buffered flat path.
+     *
+     * <p><b>Requires the caller to own {@code IndexWriterConfig}'s flush policy.</b> The exact-count
+     * guarantee above only holds if auto-flush is disabled — raise {@code setMaxBufferedDocs} above
+     * the batch size and set {@code setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH)} — and
+     * merges/index sorts stay off ({@code NoMergePolicy}, no index sort) for the life of the batch.
+     * A platform that triggers its own auto-flush based on RAM pressure or doc count (e.g. Solr,
+     * Elasticsearch) will hit the fail-fast check above under real load, even if a small test batch
+     * happened to fit under the default RAM buffer and never triggered it. Use the default
+     * heap-buffered path (leave this unset) unless the caller has that level of control.
+     *
+     * <p><b>Applies identically to every vector field.</b> This count is a single value shared by
+     * the whole segment, not set per field: if the segment has more than one vector field, each one
+     * is checked against the same {@code numInputVectors} independently, so every vector field must
+     * have a value on exactly this many documents. There is no support for a sparse/optional vector
+     * field alongside a fully-populated one, or for multiple vector fields with different
+     * cardinalities — every vector field must be populated on the same {@code numInputVectors}
+     * documents.
+     *
+     * @param numInputVectors the exact number of vectors to be indexed, or 0 to disable
+     * @return instance of {@link Builder}
+     */
+    public Builder withNumInputVectors(int numInputVectors) {
+      this.numInputVectors = numInputVectors;
+      return this;
+    }
+
     /**
      * Validates the input parameters.
      *
@@ -591,6 +657,9 @@ private void validate() throws IllegalArgumentException {
                 + MAX_NN_DESCENT_NUM_ITERATIONS
                 + "]");
       }
+      if (numInputVectors < 0) {
+        throw new IllegalArgumentException("numInputVectors cannot be negative.");
+      }
     }
 
     /**
@@ -620,7 +689,8 @@ public AcceleratedHNSWParams build() {
           strategy,
           cuvsDistanceType,
           nnDescentNumIterations,
-          hnswHeuristicType);
+          hnswHeuristicType,
+          numInputVectors);
     }
   }
 }