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 =
       () -> {
@@ -85,12 +86,14 @@ public static enum Strategy {
   private final int beamWidth;
   private final CagraGraphBuildAlgo cagraGraphBuildAlgo;
   private final CuVSIvfPqParams cuVSIvfPqParams;
+  private final boolean cuVSIvfPqParamsExplicit;
   private final int numMergeWorkers;
   private final ExecutorService mergeExec;
   private final Strategy 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.
@@ -103,12 +106,16 @@ public static enum Strategy {
    * @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 cuVSIvfPqParams An instance of CuVSIvfPqParams containing IVF_PQ specific parameters.
+   * @param cuVSIvfPqParamsExplicit whether cuVSIvfPqParams was set explicitly by the caller, as
+   *     opposed to defaulted; consulted under HEURISTIC with an explicit IVF_PQ override so a
+   *     caller-supplied value is honored instead of silently replaced by the auto-tuned one.
    * @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.
    * @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,
@@ -119,12 +126,14 @@ private AcceleratedHNSWParams(
       int beamWidth,
       CagraGraphBuildAlgo cagraGraphBuildAlgo,
       CuVSIvfPqParams cuVSIvfPqParams,
+      boolean cuVSIvfPqParamsExplicit,
       int numMergeWorkers,
       ExecutorService mergeExec,
       Strategy strategy,
       CuvsDistanceType cuvsDistanceType,
       int nnDescentNumIterations,
-      HnswHeuristicType hnswHeuristicType) {
+      HnswHeuristicType hnswHeuristicType,
+      int numInputVectors) {
     super();
     this.writerThreads = writerThreads;
     this.intermediateGraphDegree = intermediateGraphDegree;
@@ -134,12 +143,14 @@ private AcceleratedHNSWParams(
     this.beamWidth = beamWidth;
     this.cagraGraphBuildAlgo = cagraGraphBuildAlgo;
     this.cuVSIvfPqParams = cuVSIvfPqParams;
+    this.cuVSIvfPqParamsExplicit = cuVSIvfPqParamsExplicit;
     this.numMergeWorkers = numMergeWorkers;
     this.mergeExec = mergeExec;
     this.strategy = strategy;
     this.cuvsDistanceType = cuvsDistanceType;
     this.nnDescentNumIterations = nnDescentNumIterations;
     this.hnswHeuristicType = hnswHeuristicType;
+    this.numInputVectors = numInputVectors;
   }
 
   /**
@@ -214,6 +225,16 @@ public CuVSIvfPqParams getCuVSIvfPqParams() {
     return cuVSIvfPqParams;
   }
 
+  /**
+   * Whether {@link #getCuVSIvfPqParams()} was set explicitly via {@link
+   * Builder#withCuVSIvfPqParams(CuVSIvfPqParams)}, as opposed to defaulted.
+   *
+   * @return true if the caller explicitly set cuVSIvfPqParams
+   */
+  public boolean isCuVSIvfPqParamsExplicit() {
+    return cuVSIvfPqParamsExplicit;
+  }
+
   /**
    * Get the number of merge workers set to be used in the fallback mechanism
    *
@@ -272,6 +293,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 +334,8 @@ public String toString() {
         + nnDescentNumIterations
         + ", hnswHeuristicType="
         + hnswHeuristicType
+        + ", numInputVectors="
+        + numInputVectors
         + "]";
   }
 
@@ -319,11 +353,13 @@ public static class Builder {
     private CagraGraphBuildAlgo cagraGraphBuildAlgo = DEFAULT_CAGRA_GRAPH_BUILD_ALGO;
     private int numMergeWorkers = DEFAULT_NUM_MERGE_WORKERS;
     private CuVSIvfPqParams cuVSIvfPqParams = null;
+    private boolean cuVSIvfPqParamsExplicit = false;
     private ExecutorService mergeExec = null;
     private Strategy strategy = DEFAULT_STRATEGY;
     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
@@ -423,6 +459,7 @@ public Builder withCagraGraphBuildAlgo(CagraGraphBuildAlgo cagraGraphBuildAlgo)
      */
     public Builder withCuVSIvfPqParams(CuVSIvfPqParams cuVSIvfPqParams) {
       this.cuVSIvfPqParams = cuVSIvfPqParams;
+      this.cuVSIvfPqParamsExplicit = true;
       return this;
     }
 
@@ -507,6 +544,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 +662,9 @@ private void validate() throws IllegalArgumentException {
                 + MAX_NN_DESCENT_NUM_ITERATIONS
                 + "]");
       }
+      if (numInputVectors < 0) {
+        throw new IllegalArgumentException("numInputVectors cannot be negative.");
+      }
     }
 
     /**
@@ -615,12 +689,14 @@ public AcceleratedHNSWParams build() {
           beamWidth,
           cagraGraphBuildAlgo,
           cuVSIvfPqParams,
+          cuVSIvfPqParamsExplicit,
           numMergeWorkers,
           mergeExec,
           strategy,
           cuvsDistanceType,
           nnDescentNumIterations,
-          hnswHeuristicType);
+          hnswHeuristicType,
+          numInputVectors);
     }
   }
 }