diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3e35c05767..d3dec88c14 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -97,7 +97,7 @@ repos: name: spotless-fmt entry: ci/checks/run_spotless.sh pass_filenames: false - files: ^java/(cuvs-java|cuvs-lucene)/([^/]+/)?src/.*\.java$ + files: ^(java/(cuvs-java|cuvs-lucene)/([^/]+/)?|examples/java/)src/.*\.java$ exclude: .*/panama/.* language: script verbose: true diff --git a/ci/checks/run_spotless.sh b/ci/checks/run_spotless.sh index 08bf023d20..3bfac2a9a2 100755 --- a/ci/checks/run_spotless.sh +++ b/ci/checks/run_spotless.sh @@ -15,11 +15,11 @@ set -euo pipefail # Keep these in sync with the spotless-fmt hook's 'files'/'exclude' entries in # .pre-commit-config.yaml. -JAVA_SRC_PATTERN='^java/(cuvs-java|cuvs-lucene)/([^/]+/)?src/.*\.java$' +JAVA_SRC_PATTERN='^(java/(cuvs-java|cuvs-lucene)/([^/]+/)?|examples/java/)src/.*\.java$' JAVA_SRC_EXCLUDE='.*/panama/.*' java_sources_modified() { - git status --porcelain --untracked-files=all -- java/cuvs-java java/cuvs-lucene | + git status --porcelain --untracked-files=all -- java/cuvs-java java/cuvs-lucene examples/java | cut -c4- | grep -Ev "${JAVA_SRC_EXCLUDE}" | grep -Eq "${JAVA_SRC_PATTERN}" @@ -42,7 +42,7 @@ POMS=( java/cuvs-java/pom.xml java/cuvs-lucene/pom.xml java/cuvs-lucene/bench/pom.xml - java/cuvs-lucene/examples/pom.xml + examples/java/pom.xml ) for pom in "${POMS[@]}"; do diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh index 03cbb95f92..d13365f8f3 100755 --- a/ci/release/update-version.sh +++ b/ci/release/update-version.sh @@ -171,7 +171,7 @@ done NEXT_FULL_JAVA_TAG="${NEXT_SHORT_TAG}.${PATCH_PEP440}" sed_runner "s/VERSION=\".*\"/VERSION=\"${NEXT_FULL_JAVA_TAG}\"/g" java/build.sh sed_runner "s/VERSION=\".*\"/VERSION=\"${NEXT_FULL_JAVA_TAG}\"/g" java/cuvs-lucene/build.sh -for FILE in java/*/pom.xml java/cuvs-lucene/bench/pom.xml java/cuvs-lucene/examples/pom.xml; do +for FILE in java/*/pom.xml java/cuvs-lucene/bench/pom.xml examples/java/pom.xml; do sed_runner "/.*/s//${NEXT_FULL_JAVA_TAG}<\/version>/g" "${FILE}" done @@ -183,4 +183,4 @@ sed_runner "s|/[[:digit:]]\{2\}\.[[:digit:]]\{2\}\.[[:digit:]]\{1,2\}/|/${NEXT_F # title contains a release number, and that reference must not be rewritten. sed_runner "s|[[:digit:]]\{2\}\.[[:digit:]]\{2\}\.[[:digit:]]\{1,2\}|${NEXT_FULL_JAVA_TAG}|g" java/cuvs-lucene/README.md -sed_runner "s|target/examples-[\.0-9]*-jar|target/examples-${NEXT_FULL_JAVA_TAG}-jar|g" java/cuvs-lucene/examples/README.md +sed_runner "s|target/examples-[\.0-9]*-jar|target/examples-${NEXT_FULL_JAVA_TAG}-jar|g" examples/java/README.md diff --git a/examples/java/README.md b/examples/java/README.md new file mode 100644 index 0000000000..f7ae5d86bd --- /dev/null +++ b/examples/java/README.md @@ -0,0 +1,48 @@ +# Examples + +This maven project contains basic examples that showcase how `cuvs-lucene` can be used. + +## Prerequisites + +- The [`cuvs-lucene` prerequisites](../../java/cuvs-lucene/README.md#prerequisites) + +## Steps + +First build `cuvs-lucene` and install it into your local Maven repository, as described in +[Building from source](../../java/cuvs-lucene/README.md#building-from-source). From the cuVS repository root: + +```sh +./build.sh libcuvs java lucene +``` + +Then return to this directory: + +```sh +cd examples/java +``` + +To run Accelerated HNSW 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.AcceleratedHnswExample +``` + +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 `. diff --git a/java/cuvs-lucene/examples/pom.xml b/examples/java/pom.xml similarity index 100% rename from java/cuvs-lucene/examples/pom.xml rename to examples/java/pom.xml diff --git a/java/cuvs-lucene/examples/src/main/assembly/jar-with-merged-services.xml b/examples/java/src/main/assembly/jar-with-merged-services.xml similarity index 100% rename from java/cuvs-lucene/examples/src/main/assembly/jar-with-merged-services.xml rename to examples/java/src/main/assembly/jar-with-merged-services.xml diff --git a/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/AcceleratedHnswExample.java b/examples/java/src/main/java/com/nvidia/cuvs/lucene/examples/AcceleratedHnswExample.java similarity index 100% rename from java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/AcceleratedHnswExample.java rename to examples/java/src/main/java/com/nvidia/cuvs/lucene/examples/AcceleratedHnswExample.java diff --git a/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java b/examples/java/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java similarity index 100% rename from java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java rename to examples/java/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java diff --git a/examples/java/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java b/examples/java/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java new file mode 100644 index 0000000000..04dc7d900f --- /dev/null +++ b/examples/java/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java @@ -0,0 +1,263 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene.examples; + +import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; +import com.nvidia.cuvs.lucene.AcceleratedHNSWParams; +import com.nvidia.cuvs.lucene.CagraHnswBulkIndexWriter; +import com.nvidia.cuvs.lucene.FbinVectorSource; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.channels.FileChannel; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.nio.file.StandardOpenOption; +import java.util.Random; +import java.util.UUID; +import java.util.logging.Logger; +import org.apache.commons.io.FileUtils; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; + +/** + * Reference usage of {@link CagraHnswBulkIndexWriter}: builds an accelerated HNSW index (whose + * graph is built on the GPU with CAGRA) from a {@code .fbin} vector file on local disk. + * + *

Demonstrates both ways to use {@link CagraHnswBulkIndexWriter}: + * + *

    + *
  • {@link #main} — the one-shot convenience path. All of the bulk-build mechanics — + * prefetched streaming reads, native flat buffering, the {@code IndexWriterConfig} tuning + * that guarantees a single unmerged segment per slice, K-segment partitioning, and combining + * the result by hardlink — are owned by {@link CagraHnswBulkIndexWriter} itself; see its + * Javadoc for how each of those works and the tradeoffs of {@code numSegments} and {@code + * overlap}. This example only wires up what's genuinely application-specific: where the + * vectors come from ({@link FbinVectorSource}, or your own {@link + * com.nvidia.cuvs.lucene.VectorSource} for a different data source), the graph-build quality + * knobs ({@link AcceleratedHNSWParams}), and — via a {@link + * CagraHnswBulkIndexWriter.FieldCallback} — any per-vector metadata to attach. + *
  • {@link #runManualExample} — the manual, direct-instance path: construct a {@link + * CagraHnswBulkIndexWriter} yourself and drive {@code addDocument}/{@code close} exactly + * like a plain Lucene {@code IndexWriter}, building each {@link Document} (metadata included) + * yourself instead of going through a callback. + *
+ * + *

Usage: {@code OptimizedCagraHnswBuildExample [] [] [] + * []}. With no arguments a small demo {@code .fbin} is generated and indexed as a + * single segment. + */ +public class OptimizedCagraHnswBuildExample { + + private static final Logger log = + Logger.getLogger(OptimizedCagraHnswBuildExample.class.getName()); + private static final String ID_FIELD = "id"; + private static final String CATEGORY_FIELD = "category"; + private static final String VECTOR_FIELD = "vector_field"; + + public static void main(String[] args) throws Exception { + int chunkSizeMB = args.length >= 2 ? Integer.parseInt(args[1]) : 32; + int numSegments = args.length >= 3 ? Math.max(1, Integer.parseInt(args[2])) : 1; + boolean overlap = args.length >= 4 && Boolean.parseBoolean(args[3]); + Path indexDirPath = Paths.get(UUID.randomUUID().toString()); + + Path fbinPath; + boolean generated = false; + if (args.length >= 1) { + fbinPath = Paths.get(args[0]); + } else { + fbinPath = Paths.get("demo-" + UUID.randomUUID() + ".fbin"); + writeDemoFbin(fbinPath, 5000, 32, new Random(222)); + generated = true; + log.info("No .fbin provided; generated a demo file at " + fbinPath); + } + + try { + int dim; + try (FbinVectorSource probe = new FbinVectorSource(fbinPath, 1)) { + dim = probe.dimensions(); + } + + CagraHnswBulkIndexWriter.Config config = + CagraHnswBulkIndexWriter.Config.builder() + .field(VECTOR_FIELD, dim, VectorSimilarityFunction.EUCLIDEAN) + .idField(ID_FIELD) + .graphBuild( + new AcceleratedHNSWParams.Builder() + // HEURISTIC lets cuVS pick the build algorithm and auto-tune its parameters + // based on maxConn and beamWidth below. + .withStrategy(AcceleratedHNSWParams.Strategy.HEURISTIC) + // Primary recall/graph-size knobs. Higher values improve recall at the cost + // of a larger graph and longer build. Match to your dataset and recall + // target. + .withMaxConn(32) + .withBeamWidth(32) + // Must match the distance metric used when querying the index. + .withCuvsDistanceType(CuvsDistanceType.L2Expanded) + // Starting point: one thread per logical CPU. Profile and tune for your + // hardware. + .withWriterThreads(Runtime.getRuntime().availableProcessors()) + .build()) + .segments(numSegments, overlap) + .targetDirectory(indexDirPath) + .build(); + + log.info( + "Indexing " + + fbinPath + + " (" + + dim + + "-dim) into " + + numSegments + + " segment(s), " + + (overlap && numSegments > 1 ? "overlapped" : "sequential") + + " build, " + + chunkSizeMB + + " MB prefetched chunks"); + + // FieldCallback lets the one-shot path attach metadata per vector: indexFbin/build build the + // id+vector fields internally (they own the loop), so this is how a caller reaches the + // Document to add anything else -- here, an illustrative "even"/"odd" category by id. + CagraHnswBulkIndexWriter.indexFbin( + fbinPath, + config, + (doc, id) -> + doc.add( + new StringField(CATEGORY_FIELD, id % 2 == 0 ? "even" : "odd", Field.Store.YES)), + chunkSizeMB); + log.info("Index build complete: " + indexDirPath); + + runSampleSearch(indexDirPath, fbinPath, 5); + } finally { + FileUtils.deleteDirectory(indexDirPath.toFile()); + if (generated) { + Files.deleteIfExists(fbinPath); + } + } + + runManualExample(); + } + + /** Runs one k-NN query using the first vector in the file to show the index is searchable. */ + private static void runSampleSearch(Path indexDirPath, Path fbinPath, int topK) throws Exception { + float[] queryVector; + try (FbinVectorSource reader = new FbinVectorSource(fbinPath, 1)) { + queryVector = reader.get(0); + } + try (Directory dir = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(dir)) { + IndexSearcher searcher = new IndexSearcher(reader); + TopDocs results = + searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK), topK); + log.info("Sample search returned " + results.scoreDocs.length + " hits:"); + for (int i = 0; i < results.scoreDocs.length; i++) { + ScoreDoc sd = results.scoreDocs[i]; + Document hit = searcher.storedFields().document(sd.doc); + log.info( + " rank " + + (i + 1) + + ": id=" + + hit.get(ID_FIELD) + + " category=" + + hit.get(CATEGORY_FIELD) + + " score=" + + sd.score); + } + } + } + + /** + * Short demonstration of the manual, direct-instance API: {@link CagraHnswBulkIndexWriter} is + * constructed directly and driven with {@code addDocument}/{@code close}, the same shape as a + * plain Lucene {@code IndexWriter} — the caller builds each {@link Document} itself, including + * whatever metadata it wants, with no callback needed since it already owns the loop. Unlike the + * one-shot path above, this only ever builds a single segment; K-segment partitioning and + * overlap are only available via {@link CagraHnswBulkIndexWriter#indexFbin}/{@link + * CagraHnswBulkIndexWriter#build}. + */ + private static void runManualExample() throws Exception { + int numDocs = 200; + int dim = 16; + Random random = new Random(7); + Path manualIndexDirPath = Paths.get("manual-" + UUID.randomUUID()); + + try { + CagraHnswBulkIndexWriter.Config config = + CagraHnswBulkIndexWriter.Config.builder() + .field(VECTOR_FIELD, dim, VectorSimilarityFunction.EUCLIDEAN) + .graphBuild(new AcceleratedHNSWParams.Builder().build()) + .build(); + + float[][] vectors = new float[numDocs][dim]; + try (Directory dir = FSDirectory.open(manualIndexDirPath); + CagraHnswBulkIndexWriter writer = + new CagraHnswBulkIndexWriter(dir, new IndexWriterConfig(), config, numDocs)) { + for (int i = 0; i < numDocs; i++) { + for (int j = 0; j < dim; j++) { + vectors[i][j] = random.nextFloat() * 100; + } + Document doc = new Document(); + doc.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); + doc.add(new StringField(CATEGORY_FIELD, i % 2 == 0 ? "even" : "odd", Field.Store.YES)); + doc.add( + new KnnFloatVectorField( + VECTOR_FIELD, vectors[i], VectorSimilarityFunction.EUCLIDEAN)); + writer.addDocument(doc); // same call shape as a plain IndexWriter + } + } // close() runs the single native-buffered flush (the GPU CAGRA build happens here) + + try (Directory dir = FSDirectory.open(manualIndexDirPath); + DirectoryReader reader = DirectoryReader.open(dir)) { + IndexSearcher searcher = new IndexSearcher(reader); + TopDocs results = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, vectors[0], 1), 1); + Document hit = searcher.storedFields().document(results.scoreDocs[0].doc); + log.info( + "Manual example: nearest neighbor of vector 0 is id=" + + hit.get(ID_FIELD) + + " category=" + + hit.get(CATEGORY_FIELD)); + } + } finally { + FileUtils.deleteDirectory(manualIndexDirPath.toFile()); + } + } + + /** Writes a small random {@code .fbin} so the example is runnable without external data. */ + private static void writeDemoFbin(Path path, int numVectors, int dim, Random random) + throws IOException { + ByteBuffer buf = + ByteBuffer.allocate(8 + numVectors * dim * Float.BYTES).order(ByteOrder.LITTLE_ENDIAN); + buf.putInt(numVectors); // .fbin header: [num_vectors int32][dimension int32] + buf.putInt(dim); + for (int i = 0; i < numVectors; i++) { + for (int j = 0; j < dim; j++) { + buf.putFloat(random.nextFloat() * 100); + } + } + buf.flip(); + try (FileChannel ch = + FileChannel.open( + path, + StandardOpenOption.CREATE, + StandardOpenOption.WRITE, + StandardOpenOption.TRUNCATE_EXISTING)) { + while (buf.hasRemaining()) { + ch.write(buf); + } + } + } +} diff --git a/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/Utils.java b/examples/java/src/main/java/com/nvidia/cuvs/lucene/examples/Utils.java similarity index 100% rename from java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/Utils.java rename to examples/java/src/main/java/com/nvidia/cuvs/lucene/examples/Utils.java diff --git a/java/cuvs-lucene/examples/src/main/resources/logging.properties b/examples/java/src/main/resources/logging.properties similarity index 100% rename from java/cuvs-lucene/examples/src/main/resources/logging.properties rename to examples/java/src/main/resources/logging.properties diff --git a/java/cuvs-lucene/examples/README.md b/java/cuvs-lucene/examples/README.md deleted file mode 100644 index 19013675b6..0000000000 --- a/java/cuvs-lucene/examples/README.md +++ /dev/null @@ -1,34 +0,0 @@ -# Examples - -This maven project contains basic examples that showcase how `cuvs-lucene` can be used. - -## Prerequisites - -- The [`cuvs-lucene` prerequisites](../README.md#prerequisites) - -## Steps - -First build `cuvs-lucene` and install it into your local Maven repository, as described in -[Building from source](../README.md#building-from-source). From the cuVS repository root: - -```sh -./build.sh libcuvs java lucene -``` - -Then return to this directory: - -```sh -cd java/cuvs-lucene/examples -``` - -To run Accelerated HNSW 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.AcceleratedHnswExample -``` - -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 -``` diff --git a/java/cuvs-lucene/pom.xml b/java/cuvs-lucene/pom.xml index 48dc6d7889..b797bf56a4 100644 --- a/java/cuvs-lucene/pom.xml +++ b/java/cuvs-lucene/pom.xml @@ -123,6 +123,11 @@ SPDX-License-Identifier: Apache-2.0 lucene-backward-codecs 10.2.0 + + org.apache.lucene + lucene-misc + 10.2.0 + org.apache.lucene lucene-test-framework diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java index a5f164b70b..139a28c0e9 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java +++ b/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 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,18 @@ 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). Not settable via the public {@link Builder} — only {@link CagraHnswBulkIndexWriter} + * sets this, on the internal, per-slice {@link AcceleratedHNSWParams} instance it builds itself. + * + * @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 +327,8 @@ public String toString() { + nnDescentNumIterations + ", hnswHeuristicType=" + hnswHeuristicType + + ", numInputVectors=" + + numInputVectors + "]"; } @@ -324,6 +351,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 +367,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 +381,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 +434,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 +447,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 +511,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 +540,32 @@ public Builder withHnswHeuristicType(HnswHeuristicType hnswHeuristicType) { return this; } + /** + * Sets the exact number of vectors to be indexed, used to pre-allocate a single contiguous + * native flat buffer (avoiding the on-heap {@code List} 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. + * + *

Not public API. The guarantee above only holds if the caller fully owns {@code + * IndexWriterConfig}'s flush policy for the life of the batch — auto-flush disabled, no + * merges, no index sort — which most platforms embedding Lucene (Solr, Elasticsearch, + * OpenSearch) cannot promise, since they manage their own {@code IndexWriter} lifecycle. Rather + * than expose a knob a generic codec user could misconfigure and only find out at runtime, this + * is package-private and set only by {@link CagraHnswBulkIndexWriter}, which is the sole caller + * that constructs and owns the {@code IndexWriter} itself and can therefore guarantee the + * invariant. Use {@link CagraHnswBulkIndexWriter} for bulk building from local data; construct this + * codec directly (without this knob) for the general Lucene extension point. + * + * @param numInputVectors the exact number of vectors to be indexed, or 0 to disable + * @return instance of {@link Builder} + */ + Builder withNumInputVectors(int numInputVectors) { + this.numInputVectors = numInputVectors; + return this; + } + /** * Validates the input parameters. * @@ -591,6 +650,9 @@ private void validate() throws IllegalArgumentException { + MAX_NN_DESCENT_NUM_ITERATIONS + "]"); } + if (numInputVectors < 0) { + throw new IllegalArgumentException("numInputVectors cannot be negative."); + } } /** @@ -620,7 +682,8 @@ public AcceleratedHNSWParams build() { strategy, cuvsDistanceType, nnDescentNumIterations, - hnswHeuristicType); + hnswHeuristicType, + numInputVectors); } } } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index 9c49c07fe0..4bcecd6c81 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -19,8 +19,14 @@ import java.util.Random; import java.util.SortedSet; import java.util.TreeSet; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.store.ByteBuffersDataOutput; +import org.apache.lucene.store.DataOutput; import org.apache.lucene.store.IndexOutput; import org.apache.lucene.util.InfoStream; import org.apache.lucene.util.hnsw.HnswGraph; @@ -69,7 +75,7 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens layerAdjacencies.add(adjacencyMatrix); // Create the single-layer graph - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies, 1); } /** @@ -78,18 +84,25 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens * (its column count). Ceil is used to accommodate odd graph degrees. * Each layer contains 1/M nodes from the previous layer * Creates layers until the highest layer has ≤ M nodes + *

+ * Vectors for higher-layer subsets are read directly from the native matrix + * via {@link CuVSMatrix#getRow(long)} and {@link RowView#toArray(float[])}, + * avoiding any additional heap allocation of the full dataset. Used by both + * the flush and merge paths; the caller provides the vectors as a + * {@link CuVSMatrix}. */ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( FieldInfo fieldInfo, - int size, int dimensions, CuVSMatrix adjacencyListMatrix, - List vectors, + CuVSMatrix vectorDataset, int hnswLayers, CagraIndexParams params, - QuantizationType quantization) + QuantizationType quantization, + int numThreads) throws Throwable { + int size = (int) vectorDataset.size(); int M = Math.ceilDiv((int) adjacencyListMatrix.columns(), 2); // Store all layers data @@ -119,8 +132,7 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( // Select from previous layer nodes int[] prevLayerNodes = layerNodes.get(layerNodes.size() - 1); while (selectedNodesSet.size() < nextLayerSize) { - int idx = random.nextInt(prevLayerNodes.length); - selectedNodesSet.add(prevLayerNodes[idx]); + selectedNodesSet.add(prevLayerNodes[random.nextInt(prevLayerNodes.length)]); } } @@ -131,24 +143,22 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( layerNodes.add(selectedNodes); if (quantization == QuantizationType.NONE) { - // Extract vectors for selected nodes - float[][] selectedVectors = new float[nextLayerSize][]; + // Read only the sampled rows from the native matrix — no full-dataset heap copy + float[][] selectedVectors = new float[nextLayerSize][dimensions]; for (int i = 0; i < nextLayerSize; i++) { - selectedVectors[i] = (float[]) vectors.get(selectedNodes[i]); + vectorDataset.getRow(selectedNodes[i]).toArray(selectedVectors[i]); } // Build CAGRA graph for this layer layerAdjacencies.add( buildCagraGraphForSubset( selectedVectors, selectedNodes, 0, params, dimensions, quantization)); - } else { - - // Extract vectors for selected nodes - int bytesPerVector = (dimensions + 7) / 8; - byte[][] selectedVectors = new byte[nextLayerSize][]; + // Byte width comes from the matrix itself: binary packs 8 dims/byte, scalar is 1 byte/dim. + int bytesPerVector = (int) vectorDataset.columns(); + byte[][] selectedVectors = new byte[nextLayerSize][bytesPerVector]; for (int i = 0; i < nextLayerSize; i++) { - selectedVectors[i] = (byte[]) vectors.get(selectedNodes[i]); + vectorDataset.getRow(selectedNodes[i]).toArray(selectedVectors[i]); } // Build CAGRA graph for this layer @@ -166,7 +176,7 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( } // Create the multi-layer graph with all layers - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies, numThreads); } /** @@ -184,11 +194,9 @@ private static CuVSMatrix buildCagraGraphForSubset( CuVSMatrix subsetDataset; if (quantization == QuantizationType.BINARY) { - subsetDataset = - createByteMatrixFromArray((byte[][]) vectors, bytesPerVector, getCuVSResourcesInstance()); + subsetDataset = createByteMatrixFromArray((byte[][]) vectors, bytesPerVector); } else if (quantization == QuantizationType.SCALAR) { - subsetDataset = - createByteMatrixFromArray((byte[][]) vectors, dimensions, getCuVSResourcesInstance()); + subsetDataset = createByteMatrixFromArray((byte[][]) vectors, dimensions); } else { subsetDataset = CuVSMatrix.ofArray((float[][]) vectors); } @@ -235,56 +243,161 @@ private static CuVSMatrix buildCagraGraphForSubset( * @return a 2D array of offsets * @throws IOException I/O Exceptions */ - public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex) + public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex, int numThreads) throws IOException { // write vectors' neighbors on each level into the vectorIndex file int countOnLevel0 = graph.size(); - int[][] offsets = new int[graph.numLevels()][]; - int[] scratch = new int[graph.maxConn() * 2]; - for (int level = 0; level < graph.numLevels(); level++) { + int numLevels = graph.numLevels(); + int[][] offsets = new int[numLevels][]; + + // Level 0 holds all nodes and dominates serialization cost. Each node's delta/VInt block is + // independent, so encode level 0 in parallel and concatenate the per-thread buffers serially in + // node order, in memory-bounded waves. Higher levels are tiny and stay serial. The on-disk + // bytes + // are identical to the fully-serial path (blocks in node order, offsets = per-node byte + // lengths). + // graph.maxConn() scans every layer-0 adjacency row (O(graph size)); compute it once here + // rather than per level/per task below. + int maxConn = graph.maxConn(); + + int[] level0Nodes = NodesIterator.getSortedNodes(graph.getNodesOnLevel(0)); + offsets[0] = new int[level0Nodes.length]; + if (numThreads > 1 && level0Nodes.length >= PARALLEL_MIN_NODES) { + writeLevel0Parallel( + graph, vectorIndex, level0Nodes, offsets[0], countOnLevel0, maxConn, numThreads); + } else { + writeLevelSerial(graph, vectorIndex, 0, level0Nodes, offsets[0], countOnLevel0, maxConn); + } + + for (int level = 1; level < numLevels; level++) { int[] sortedNodes = NodesIterator.getSortedNodes(graph.getNodesOnLevel(level)); offsets[level] = new int[sortedNodes.length]; - int nodeOffsetId = 0; - - for (int node : sortedNodes) { - // Get node neighbors - NeighborArray neighbors = graph.getNeighbors(level, node); - // Get the size of the neighbor array - int size = neighbors.size(); - // Write size in VInt as the neighbors list is typically small - long offsetStart = vectorIndex.getFilePointer(); - // Get neighbors - int[] nnodes = neighbors.nodes(); - // Sort them - Arrays.sort(nnodes, 0, size); - // Now that we have sorted, do delta encoding to minimize the required bits to store the - // information - int actualSize = 0; - if (size > 0) { - scratch[0] = nnodes[0]; - actualSize = 1; - } - // De-duplication - for (int i = 1; i < size; i++) { - assert nnodes[i] < countOnLevel0 : "node too large: " + nnodes[i] + ">=" + countOnLevel0; - // Sorting step helps here - if (nnodes[i - 1] == nnodes[i]) { + writeLevelSerial( + graph, vectorIndex, level, sortedNodes, offsets[level], countOnLevel0, maxConn); + } + // Return offsets (information written while writing the meta info) + return offsets; + } + + /** Node count below which parallel level-0 serialization is not worth the overhead. */ + private static final int PARALLEL_MIN_NODES = 1 << 16; + + /** Nodes per wave — bounds the transient encode buffer regardless of dataset size. */ + private static final int WAVE_NODES = 1 << 20; + + /** Serially encodes a level's nodes into {@code out}, recording per-node byte lengths. */ + private static void writeLevelSerial( + GPUBuiltHnswGraph graph, + IndexOutput out, + int level, + int[] sortedNodes, + int[] offsets, + int countOnLevel0, + int maxConn) + throws IOException { + int[] scratch = new int[maxConn * 2]; + int idx = 0; + for (int node : sortedNodes) { + long start = out.getFilePointer(); + encodeNode(graph.getNeighbors(level, node), scratch, out, countOnLevel0); + offsets[idx++] = Math.toIntExact(out.getFilePointer() - start); + } + } + + /** + * Encodes level 0 in parallel: within memory-bounded waves, threads encode contiguous node + * sub-ranges into per-thread buffers, which are then concatenated to {@code out} in node order + * (identical layout to the serial path). + */ + private static void writeLevel0Parallel( + GPUBuiltHnswGraph graph, + IndexOutput out, + int[] nodes, + int[] offsets, + int countOnLevel0, + int maxConn, + int numThreads) + throws IOException { + ExecutorService pool = Executors.newFixedThreadPool(numThreads); + try { + int n = nodes.length; + for (int waveStart = 0; waveStart < n; waveStart += WAVE_NODES) { + int waveEnd = Math.min(waveStart + WAVE_NODES, n); + int perThread = (waveEnd - waveStart + numThreads - 1) / numThreads; + + ByteBuffersDataOutput[] buffers = new ByteBuffersDataOutput[numThreads]; + List> futures = new ArrayList<>(numThreads); + for (int t = 0; t < numThreads; t++) { + final int subStart = waveStart + t * perThread; + final int subEnd = Math.min(subStart + perThread, waveEnd); + final int slot = t; + if (subStart >= subEnd) { continue; } - scratch[actualSize++] = nnodes[i] - nnodes[i - 1]; + futures.add( + pool.submit( + () -> { + ByteBuffersDataOutput buffer = new ByteBuffersDataOutput(); + int[] scratch = new int[maxConn * 2]; + for (int i = subStart; i < subEnd; i++) { + long before = buffer.size(); + encodeNode(graph.getNeighbors(0, nodes[i]), scratch, buffer, countOnLevel0); + offsets[i] = Math.toIntExact(buffer.size() - before); + } + buffers[slot] = buffer; + return null; + })); + } + for (Future f : futures) { + f.get(); } - // Write the size after duplicates are removed - vectorIndex.writeVInt(actualSize); - // Write de-duplicated neighbors - for (int i = 0; i < actualSize; i++) { - vectorIndex.writeVInt(scratch[i]); + // Concatenate in thread order (== node order), preserving the serial byte layout. + for (ByteBuffersDataOutput buffer : buffers) { + if (buffer != null) { + buffer.copyTo(out); + } } - offsets[level][nodeOffsetId++] = - Math.toIntExact(vectorIndex.getFilePointer() - offsetStart); } + pool.shutdown(); + } catch (InterruptedException e) { + pool.shutdownNow(); + Thread.currentThread().interrupt(); + throw new IOException("Interrupted during parallel writeGraph", e); + } catch (ExecutionException e) { + pool.shutdownNow(); + throw new IOException("Parallel writeGraph failed", e.getCause()); + } catch (RuntimeException | IOException e) { + pool.shutdownNow(); + throw e; + } + } + + /** + * Sorts, delta-encodes and de-duplicates a node's neighbors and writes the block (VInt size + VInt + * deltas) to {@code out}. Shared by the serial and parallel paths so encoding is identical. + */ + private static void encodeNode( + NeighborArray neighbors, int[] scratch, DataOutput out, int countOnLevel0) + throws IOException { + int size = neighbors == null ? 0 : neighbors.size(); + int actualSize = 0; + if (size > 0) { + int[] nnodes = neighbors.nodes(); + Arrays.sort(nnodes, 0, size); + scratch[0] = nnodes[0]; + actualSize = 1; + for (int i = 1; i < size; i++) { + assert nnodes[i] < countOnLevel0 : "node too large: " + nnodes[i] + ">=" + countOnLevel0; + if (nnodes[i - 1] == nnodes[i]) { + continue; + } + scratch[actualSize++] = nnodes[i] - nnodes[i - 1]; + } + } + out.writeVInt(actualSize); + for (int i = 0; i < actualSize; i++) { + out.writeVInt(scratch[i]); } - // Return offsets (information written while writing the meta info) - return offsets; } /** diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraHnswBulkIndexWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraHnswBulkIndexWriter.java new file mode 100644 index 0000000000..219d7df955 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraHnswBulkIndexWriter.java @@ -0,0 +1,660 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.spi.CuVSProvider; +import java.io.Closeable; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.List; +import java.util.Objects; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.Semaphore; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.stream.Stream; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.IndexableField; +import org.apache.lucene.index.NoMergePolicy; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.misc.store.HardlinkCopyDirectoryWrapper; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; + +/** + * Builds a GPU-accelerated CAGRA/HNSW Lucene index from data the caller already has ready on + * local disk, as few segments as possible, without going through a generically-configurable + * {@link org.apache.lucene.codecs.Codec}. + * + *

This is not a general-purpose Lucene extension point: an instance owns a single + * {@link IndexWriter} and its {@link IndexWriterConfig} so that the invariants the underlying GPU + * writer requires (single unmerged segment, no index sort, exact vector count) are guaranteed by + * this class rather than left to a caller to get right. If you need a general Lucene codec that + * any indexing pipeline (including ones that don't control their own {@code IndexWriter} + * lifecycle, e.g. Solr or Elasticsearch) can register, use {@link Lucene101AcceleratedHNSWCodec} + * directly instead. + * + *

Two ways to use this class: + * + *

    + *
  • Manual, single segment: construct an instance directly, call {@link #addDocument} + * per document exactly like a plain {@link IndexWriter}, then {@link #close}. You own the + * loop and the {@link Document} you build (any fields, not just the vector). + *
  • One-shot, one or many segments: {@link #indexFbin} / {@link #build(VectorSource, + * Config)} own the loop for you — they read vectors from a {@code .fbin} file or {@link + * VectorSource}, optionally split into {@code numSegments} partitions (sequential or + * overlapped), and combine the result. Since they build each row's {@link Document} + * internally, an optional {@link FieldCallback} lets you add extra fields to it. + *
+ * + *

Scope: CAGRA_HNSW (GPU build, CPU search) only. This class builds indexes for {@link + * Lucene101AcceleratedHNSWCodec}. It does not support the GPU-search codec ({@code + * CuVS2510GPUSearchCodec}) — that writer does not (yet) have the native flat-buffering + * optimization this class relies on. Reading an existing index for CAGRA_SEARCH-style GPU search + * is unaffected by this class either way; only bulk-building one is out of scope for now. + */ +public final class CagraHnswBulkIndexWriter implements Closeable { + + private static final int DEFAULT_CHUNK_SIZE_MB = 32; + private static final AtomicBoolean RMM_ENABLED = new AtomicBoolean(false); + + private final IndexWriter writer; + private final int exactVectorCount; + private int documentsAdded; + private boolean closed; + + /** + * Opens a single-segment, native-flat-buffered writer. {@code exactVectorCount} must equal the + * number of {@link #addDocument} calls that will follow — the native buffer is pre-sized to it, + * so {@link #close} rejects a mismatched count rather than let the underlying writer produce a + * corrupt or incomplete segment. + * + *

{@code conf}'s {@code Analyzer}, {@code Similarity}, {@code InfoStream}, and {@code + * OpenMode} are honored; an explicit {@code IndexSort} is rejected ({@link + * IllegalArgumentException}), since native flat buffering does not support index-sorted + * segments. Codec, merge policy, and flush thresholds are always owned by this class regardless + * of what {@code conf} contains — {@link IndexWriterConfig} does not expose whether the caller + * explicitly set those or left them at Lucene's defaults, so there is no reliable way to + * validate-and-reject a caller-supplied value for them the way {@code IndexSort} can be; this + * class simply never reads them from {@code conf}. + * + *

{@code config.targetDirectory()}, {@code config.numSegments()}, {@code + * config.overlapped()}, and {@code config.pipelineDepth()} are not consulted here — {@code + * directory} is passed explicitly, and this constructor always builds exactly one segment. Those + * fields only matter to {@link #indexFbin} / {@link #build(VectorSource, Config)}. + */ + public CagraHnswBulkIndexWriter( + Directory directory, IndexWriterConfig conf, Config config, int exactVectorCount) + throws Exception { + Objects.requireNonNull(directory, "directory"); + Objects.requireNonNull(conf, "conf"); + Objects.requireNonNull(config, "config"); + if (exactVectorCount <= 0) { + throw new IllegalArgumentException("exactVectorCount must be > 0, got " + exactVectorCount); + } + if (conf.getIndexSort() != null) { + throw new IllegalArgumentException( + "CagraHnswBulkIndexWriter does not support an index-sorted segment (native flat" + + " buffering requires an unsorted single-segment build); leave" + + " IndexWriterConfig.indexSort unset"); + } + ensureRMMEnabled(); + this.exactVectorCount = exactVectorCount; + + Codec codec = + new Lucene101AcceleratedHNSWCodec( + sizedForSlice(config.graphBuildParams(), exactVectorCount)); + IndexWriterConfig ownedConf = + new IndexWriterConfig(conf.getAnalyzer()) + .setSimilarity(conf.getSimilarity()) + .setInfoStream(conf.getInfoStream()) + .setCodec(codec) + .setUseCompoundFile(false) + .setMaxBufferedDocs(Math.max(2, exactVectorCount + 1)) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) + .setMergePolicy(NoMergePolicy.INSTANCE) + .setOpenMode(conf.getOpenMode()); + this.writer = new IndexWriter(directory, ownedConf); + } + + /** + * Adds one document, exactly like {@link IndexWriter#addDocument}. {@code doc} may contain any + * fields — the vector field (matching {@link Config#fieldName()}) is routed into the native + * flat buffer automatically by the underlying codec, the same way any {@link + * KnnFloatVectorField} is for any Lucene codec; every other field is indexed normally. + */ + public long addDocument(Iterable doc) throws IOException { + if (closed) { + throw new IllegalStateException("addDocument called after close()"); + } + if (documentsAdded >= exactVectorCount) { + throw new IllegalStateException( + "addDocument called more than exactVectorCount (" + exactVectorCount + ") times"); + } + long seqNo = writer.addDocument(doc); + documentsAdded++; + return seqNo; + } + + /** + * Runs the single native-buffered flush (this is where the GPU CAGRA build happens) and closes + * the underlying writer. There is no separate {@code commit()}: unlike a plain {@link + * IndexWriter}, only one flush is ever valid for this instance, so folding it into {@code + * close()} removes any way to trigger it early with fewer than {@code exactVectorCount} vectors + * added. Throws {@link IllegalStateException} if {@link #addDocument} was called fewer times + * than {@code exactVectorCount} promised. + */ + @Override + public void close() throws IOException { + if (closed) { + return; + } + closed = true; + if (documentsAdded != exactVectorCount) { + writer.close(); // release resources before reporting the mismatch + throw new IllegalStateException( + "expected " + exactVectorCount + " documents, got " + documentsAdded); + } + try (writer) { + writer.commit(); + } + } + + /** + * Callback invoked once per row by {@link #indexFbin} / {@link #build(VectorSource, Config, + * FieldCallback)}, right after the id and vector fields have been added to {@code document} and + * right before it is added to the index — lets the caller attach additional fields (metadata) + * per vector. Not used by the manual, direct-instance API, where the caller already builds the + * whole {@link Document} themselves. + */ + @FunctionalInterface + public interface FieldCallback { + void addFields(Document document, int id) throws IOException; + } + + /** + * Builds an index from {@code source} into {@code config.targetDirectory()}, partitioned into + * {@code config.numSegments()} contiguous slices as {@code source} is consumed front-to-back. + * + *

{@code config.overlapped()} is not supported here: a single {@link VectorSource} is + * forward-only/single-consumer (see its contract) and so cannot be read concurrently by + * multiple slice builders. Use {@link #indexFbin} for the overlapped pipeline, which knows how + * to open independent, per-slice sources over the same underlying file. + */ + public static void build(VectorSource source, Config config) throws Exception { + build(source, config, null); + } + + /** As {@link #build(VectorSource, Config)}, with a {@link FieldCallback} for extra fields. */ + public static void build(VectorSource source, Config config, FieldCallback callback) + throws Exception { + Objects.requireNonNull(source, "source"); + Objects.requireNonNull(config, "config"); + if (config.overlapped()) { + throw new IllegalArgumentException( + "Config.overlapped() is not supported by build(VectorSource, Config): a VectorSource is" + + " forward-only and cannot be read by multiple concurrent slice builders. Use" + + " indexFbin(...) for the overlapped multi-segment build."); + } + if (source.dimensions() != config.dimensions()) { + throw new IllegalArgumentException( + "source.dimensions() (" + + source.dimensions() + + ") does not match Config.dimensions() (" + + config.dimensions() + + ")"); + } + List slices = sliceEvenly(source.size(), config.numSegments()); + buildSequential(source, config, callback, slices); + } + + /** + * Convenience entry point: builds an index directly from a {@code .fbin} file ({@code + * [num_vectors int32][dim int32]} header, then contiguous little-endian float32 rows), using + * {@link FbinVectorSource} with {@link #DEFAULT_CHUNK_SIZE_MB}-sized prefetched chunks. + */ + public static void indexFbin(Path fbinPath, Config config) throws Exception { + indexFbin(fbinPath, config, null, DEFAULT_CHUNK_SIZE_MB); + } + + /** As {@link #indexFbin(Path, Config)}, with a {@link FieldCallback} for extra fields. */ + public static void indexFbin(Path fbinPath, Config config, FieldCallback callback) + throws Exception { + indexFbin(fbinPath, config, callback, DEFAULT_CHUNK_SIZE_MB); + } + + /** As {@link #indexFbin(Path, Config)}, with an explicit prefetch chunk size. */ + public static void indexFbin(Path fbinPath, Config config, int chunkSizeMB) throws Exception { + indexFbin(fbinPath, config, null, chunkSizeMB); + } + + /** + * As {@link #indexFbin(Path, Config)}, with an explicit prefetch chunk size and {@link + * FieldCallback}. When {@code config.numSegments() > 1} and {@code config.overlapped()}, builds + * up to {@code config.pipelineDepth()} segments concurrently, each over its own slice of {@code + * fbinPath}, then combines them by hardlink; otherwise builds sequentially over one shared + * reader. + */ + public static void indexFbin( + Path fbinPath, Config config, FieldCallback callback, int chunkSizeMB) throws Exception { + Objects.requireNonNull(fbinPath, "fbinPath"); + Objects.requireNonNull(config, "config"); + int total; + int dim; + try (FbinVectorSource probe = new FbinVectorSource(fbinPath, 1)) { + total = probe.size(); + dim = probe.dimensions(); + } + if (dim != config.dimensions()) { + throw new IllegalArgumentException( + "fbinPath dimension (" + + dim + + ") does not match Config.dimensions() (" + + config.dimensions() + + ")"); + } + List slices = sliceEvenly(total, config.numSegments()); + if (config.overlapped() && slices.size() > 1) { + buildOverlapped(fbinPath, config, callback, chunkSizeMB, slices, dim); + } else { + try (FbinVectorSource source = new FbinVectorSource(fbinPath, chunkSizeMB)) { + buildSequential(source, config, callback, slices); + } + } + } + + private static void ensureRMMEnabled() { + if (RMM_ENABLED.compareAndSet(false, true)) { + CuVSProvider.provider().enableRMMAsyncMemory(); + } + } + + /** + * Sequential partitioned build: {@code source} is streamed front-to-back across all slices, + * each slice built as a single native-flat segment appended to the same directory (first slice + * {@code CREATE}, later slices {@code APPEND}). Peak host memory is one slice's native buffer. + */ + private static void buildSequential( + VectorSource source, Config config, FieldCallback callback, List slices) + throws Exception { + float[] scratch = new float[config.dimensions()]; + try (Directory dir = FSDirectory.open(config.targetDirectory())) { + for (int p = 0; p < slices.size(); p++) { + int[] slice = slices.get(p); + buildSegment( + dir, source, scratch, config, callback, slice[0], slice[0], slice[1], p == 0, null); + } + } + } + + /** + * Overlapped partitioned build: a bounded pool builds up to {@code config.pipelineDepth()} + * segments at once, each with its OWN {@link FbinVectorSource} over just its slice, so a + * segment's ingest overlaps a prior segment's GPU commit. The GPU build itself is serialized on + * a single permit. The finished per-segment indexes are combined into {@code + * config.targetDirectory()} by hardlinking their files (no bulk copy of the vector data). + */ + private static void buildOverlapped( + Path fbinPath, + Config config, + FieldCallback callback, + int chunkSizeMB, + List slices, + int dim) + throws Exception { + Path targetDir = config.targetDirectory(); + int depth = Math.min(slices.size(), config.pipelineDepth()); + List segDirs = new ArrayList<>(); + for (int p = 0; p < slices.size(); p++) { + segDirs.add(targetDir.resolveSibling(targetDir.getFileName() + "_p" + p)); + } + for (Path segDir : segDirs) { + deleteRecursivelyQuietly(segDir); + } + try { + Semaphore gpuPermit = new Semaphore(1); // serialize the GPU CAGRA build across segments + ExecutorService pool = Executors.newFixedThreadPool(depth); + List> futures = new ArrayList<>(); + for (int p = 0; p < slices.size(); p++) { + int[] slice = slices.get(p); + Path segDir = segDirs.get(p); + futures.add( + pool.submit( + () -> { + float[] scratch = new float[dim]; + try (FbinVectorSource source = + new FbinVectorSource(fbinPath, slice[0], slice[1], chunkSizeMB); + Directory d = FSDirectory.open(segDir)) { + // createNew=true: each segment is a fresh single-segment index in its own + // dir; sourceStart=0 since this source is already windowed to the slice. + buildSegment( + d, source, scratch, config, callback, 0, slice[0], slice[1], true, + gpuPermit); + } + return null; + })); + } + pool.shutdown(); + try { + for (Future f : futures) { + f.get(); // propagate any build failure + } + } catch (Exception e) { + throw new IOException("Overlapped bulk index build failed", e); + } finally { + pool.shutdownNow(); + } + combineByHardlink(targetDir, segDirs); + } finally { + for (Path segDir : segDirs) { + deleteRecursivelyQuietly(segDir); + } + } + } + + /** + * Builds one segment from {@code size} vectors: {@code source.get(sourceStart + i, ...)} for + * {@code i} in {@code [0, size)}, labelled with ids {@code idStart + i} (the vectors' absolute + * position in the overall build, regardless of whether {@code source} itself is windowed). Opens + * a {@link CagraHnswBulkIndexWriter} sized to {@code size} and drives it exactly like the manual + * API. When {@code gpuPermit} is non-null the close (which runs the GPU CAGRA build) is + * serialized on it while other segments' host-side ingest may proceed. + */ + private static void buildSegment( + Directory dir, + VectorSource source, + float[] scratch, + Config config, + FieldCallback callback, + int sourceStart, + int idStart, + int size, + boolean createNew, + Semaphore gpuPermit) + throws Exception { + IndexWriterConfig conf = + new IndexWriterConfig() + .setOpenMode( + createNew ? IndexWriterConfig.OpenMode.CREATE : IndexWriterConfig.OpenMode.APPEND); + CagraHnswBulkIndexWriter writer = new CagraHnswBulkIndexWriter(dir, conf, config, size); + try { + for (int i = 0; i < size; i++) { + source.get(sourceStart + i, scratch); + int id = idStart + i; + Document doc = new Document(); + if (config.idFieldName() != null) { + doc.add(new StringField(config.idFieldName(), Integer.toString(id), Field.Store.YES)); + } + doc.add(new KnnFloatVectorField(config.fieldName(), scratch, config.similarity())); + if (callback != null) { + callback.addFields(doc, id); + } + writer.addDocument(doc); // copies the vector -> 'scratch' is safe to reuse next iteration + } + // The single flush inside close() is where the GPU CAGRA build runs; serialize it if asked. + if (gpuPermit != null) { + gpuPermit.acquire(); + try { + writer.close(); + } finally { + gpuPermit.release(); + } + } else { + writer.close(); + } + } finally { + writer.close(); // no-op if already closed above; ensures cleanup on exception too + } + } + + /** Copies every knob from {@code template} except {@code numInputVectors}, set to {@code size}. */ + private static AcceleratedHNSWParams sizedForSlice(AcceleratedHNSWParams template, int size) { + return new AcceleratedHNSWParams.Builder() + .withWriterThreads(template.getWriterThreads()) + .withIntermediateGraphDegree(template.getIntermediateGraphDegree()) + .withGraphDegree(template.getGraphdegree()) + .withHNSWLayer(template.getHnswLayers()) + .withMaxConn(template.getMaxConn()) + .withBeamWidth(template.getBeamWidth()) + .withCagraGraphBuildAlgo(template.getCagraGraphBuildAlgo()) + .withCuVSIvfPqParams(template.getCuVSIvfPqParams()) + .withNumMergeWorkers(template.getNumMergeWorkers()) + .withMergeExecutorService(template.getMergeExec()) + .withStrategy(template.getStrategy()) + .withCuvsDistanceType(template.getCuvsDistanceType()) + .withNNDescentNumIterations(template.getNNDescentNumIterations()) + .withHnswHeuristicType(template.getHnswHeuristicType()) + .withNumInputVectors(size) + .build(); + } + + /** + * Combines the per-segment indexes into {@code targetDir} by hardlinking their files (same + * filesystem) rather than copying the vector data. {@link HardlinkCopyDirectoryWrapper} falls + * back to a byte copy automatically if the segment dirs and the final dir are on different + * filesystems. + */ + private static void combineByHardlink(Path targetDir, List segDirs) throws IOException { + Directory[] sources = new Directory[segDirs.size()]; + try { + for (int i = 0; i < segDirs.size(); i++) { + sources[i] = FSDirectory.open(segDirs.get(i)); + } + IndexWriterConfig iwc = + new IndexWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE); // keep segments separate + try (Directory target = new HardlinkCopyDirectoryWrapper(FSDirectory.open(targetDir)); + IndexWriter combiner = new IndexWriter(target, iwc)) { + combiner.addIndexes(sources); + } + } finally { + for (Directory s : sources) { + if (s != null) { + s.close(); + } + } + } + } + + /** Splits {@code total} into {@code k} contiguous [start, size] slices, spreading the remainder. */ + private static List sliceEvenly(int total, int k) { + List slices = new ArrayList<>(); + int base = total / k; + int rem = total % k; + int start = 0; + for (int p = 0; p < k; p++) { + int size = base + (p < rem ? 1 : 0); // spread the remainder over the first slices + if (size <= 0) { + continue; + } + slices.add(new int[] {start, size}); + start += size; + } + return slices; + } + + private static void deleteRecursivelyQuietly(Path path) { + if (!Files.exists(path)) { + return; + } + try (Stream walk = Files.walk(path)) { + walk.sorted(Comparator.reverseOrder()) + .forEach( + p -> { + try { + Files.deleteIfExists(p); + } catch (IOException ignored) { + // best-effort cleanup of a temp per-segment dir + } + }); + } catch (IOException ignored) { + // best-effort cleanup of a temp per-segment dir + } + } + + /** Immutable configuration for {@link CagraHnswBulkIndexWriter}. */ + public static final class Config { + private final String fieldName; + private final int dimensions; + private final VectorSimilarityFunction similarity; + private final String idFieldName; + private final AcceleratedHNSWParams graphBuildParams; + private final Path targetDirectory; + private final int numSegments; + private final boolean overlapped; + private final int pipelineDepth; + + private Config(Builder b) { + this.fieldName = b.fieldName; + this.dimensions = b.dimensions; + this.similarity = b.similarity; + this.idFieldName = b.idFieldName; + this.graphBuildParams = b.graphBuildParams; + this.targetDirectory = b.targetDirectory; + this.numSegments = b.numSegments; + this.overlapped = b.overlapped; + this.pipelineDepth = b.pipelineDepth; + } + + public String fieldName() { + return fieldName; + } + + public int dimensions() { + return dimensions; + } + + public VectorSimilarityFunction similarity() { + return similarity; + } + + public String idFieldName() { + return idFieldName; + } + + public AcceleratedHNSWParams graphBuildParams() { + return graphBuildParams; + } + + /** Only consulted by {@link #indexFbin} / {@link #build(VectorSource, Config)}. */ + public Path targetDirectory() { + return targetDirectory; + } + + /** Only consulted by {@link #indexFbin} / {@link #build(VectorSource, Config)}. */ + public int numSegments() { + return numSegments; + } + + /** Only consulted by {@link #indexFbin}. */ + public boolean overlapped() { + return overlapped; + } + + /** Only consulted by {@link #indexFbin}. */ + public int pipelineDepth() { + return pipelineDepth; + } + + public static Builder builder() { + return new Builder(); + } + + /** Builder for {@link Config}. */ + public static final class Builder { + private String fieldName = "vector"; + private int dimensions = -1; + private VectorSimilarityFunction similarity = VectorSimilarityFunction.EUCLIDEAN; + private String idFieldName = "id"; + private AcceleratedHNSWParams graphBuildParams; + private Path targetDirectory; + private int numSegments = 1; + private boolean overlapped = false; + private int pipelineDepth = 2; + + /** Sets the vector field name and dimensionality; required. */ + public Builder field(String fieldName, int dimensions, VectorSimilarityFunction similarity) { + this.fieldName = Objects.requireNonNull(fieldName, "fieldName"); + this.dimensions = dimensions; + this.similarity = Objects.requireNonNull(similarity, "similarity"); + return this; + } + + /** + * Sets the stored id field name (holding each document's absolute position in the build, as + * a string), added automatically by {@link #indexFbin} / {@link #build(VectorSource, + * Config)} before the {@link FieldCallback} runs. Defaults to {@code "id"}; pass {@code + * null} to disable. Not used by the manual, direct-instance API. + */ + public Builder idField(String idFieldName) { + this.idFieldName = idFieldName; + return this; + } + + /** + * Sets the CAGRA/HNSW graph-build parameters; required. {@link + * AcceleratedHNSWParams#getNumInputVectors()} on the supplied instance is ignored — this + * class sizes the native flat buffer to each slice/instance itself. + */ + public Builder graphBuild(AcceleratedHNSWParams graphBuildParams) { + this.graphBuildParams = Objects.requireNonNull(graphBuildParams, "graphBuildParams"); + return this; + } + + /** Sets the directory the final index is written to; required for {@link #indexFbin}/{@link #build}. */ + public Builder targetDirectory(Path targetDirectory) { + this.targetDirectory = Objects.requireNonNull(targetDirectory, "targetDirectory"); + return this; + } + + /** + * Splits the build into {@code numSegments} contiguous slices, each a single native-flat + * segment. Peak host memory scales as {@code 1/numSegments}; the GPU build itself is always + * serialized across slices regardless of this setting. Default 1 (single segment). + * + * @param overlapped when {@code numSegments > 1} and building via {@link #indexFbin}, + * builds up to {@link #pipelineDepth} slices concurrently (ingest of one overlapping the + * GPU commit of another) instead of strictly sequentially. Ignored by {@link + * #build(VectorSource, Config)}, which always builds sequentially — see that method's + * Javadoc. + */ + public Builder segments(int numSegments, boolean overlapped) { + if (numSegments < 1) { + throw new IllegalArgumentException("numSegments must be >= 1, got " + numSegments); + } + this.numSegments = numSegments; + this.overlapped = overlapped; + return this; + } + + /** Max segments built concurrently in overlap mode; peak host memory is this many slice buffers. */ + public Builder pipelineDepth(int pipelineDepth) { + if (pipelineDepth < 1) { + throw new IllegalArgumentException("pipelineDepth must be >= 1, got " + pipelineDepth); + } + this.pipelineDepth = pipelineDepth; + return this; + } + + public Config build() { + if (dimensions <= 0) { + throw new IllegalStateException( + "field(...) must be called with a positive dimension count"); + } + Objects.requireNonNull(graphBuildParams, "graphBuild(...) must be called"); + return new Config(this); + } + } + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java index 8bb46faf61..ef0b82b3cf 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java @@ -75,7 +75,8 @@ public static CagraIndexParams create( AcceleratedHNSWParams acceleratedHNSWParams, long rows, long dimension) { if (acceleratedHNSWParams.getStrategy().equals(AcceleratedHNSWParams.Strategy.HEURISTIC)) { // Delegate the derivation of the graph degrees, build algorithm and its parameters to cuVS, - // expressed in terms of the HNSW-equivalent maxConn/beamWidth. + // expressed in terms of the HNSW-equivalent maxConn/beamWidth. cagraGraphBuildAlgo is not + // consulted here; it only applies under CUSTOM. CagraIndexParams derived = CagraIndexParams.fromHnswParams( rows, @@ -84,9 +85,9 @@ public static CagraIndexParams create( acceleratedHNSWParams.getBeamWidth(), acceleratedHNSWParams.getHnswHeuristicType(), acceleratedHNSWParams.getCuvsDistanceType()); - // TODO: fromHnswParams has no writerThreads argument, so its result carries the cuVS default - // (not a heuristic value). We can rebuild the CagraIndexParams with the caller-supplied - // writerThreads for now but should fix this in cuVS in the future. + // TODO: fromHnswParams has no writerThreads argument, so its result carries the cuVS + // default (not a heuristic value). We can rebuild the CagraIndexParams with the + // caller-supplied writerThreads for now but should fix this in cuVS in the future. return new CagraIndexParams.Builder() .withGraphDegree(derived.getGraphDegree()) .withIntermediateGraphDegree(derived.getIntermediateGraphDegree()) @@ -97,6 +98,7 @@ public static CagraIndexParams create( .withNumWriterThreads(acceleratedHNSWParams.getWriterThreads()) .build(); } + // CUSTOM: forward the caller's algorithm and the parameters it consumes. return new CagraIndexParams.Builder() .withGraphDegree(acceleratedHNSWParams.getGraphdegree()) .withIntermediateGraphDegree(acceleratedHNSWParams.getIntermediateGraphDegree()) diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java index aead683dc4..95bf16ebe9 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java @@ -201,8 +201,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro var cagraIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); try { CuVSMatrix cagraDataset = - Utils.createFloatMatrix( - vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance()); + Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); writeCagraIndex(cagraIndexOutputStream, cagraDataset); } catch (Throwable t) { // Fallback to brute force in a few cases, for now. @@ -223,8 +222,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro if (indexType.isBruteForce()) { var bruteForceIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); CuVSMatrix bruteforceDataset = - Utils.createFloatMatrix( - vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance()); + Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); writeBruteForceIndex(bruteForceIndexOutputStream, bruteforceDataset); bruteForceIndexLength = cuvsIndex.getFilePointer() - bruteForceIndexOffset; diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FbinVectorSource.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FbinVectorSource.java new file mode 100644 index 0000000000..3264ce7934 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FbinVectorSource.java @@ -0,0 +1,213 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.channels.FileChannel; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.BlockingQueue; + +/** + * Prefetching, double-buffered {@link VectorSource} over an uncompressed {@code .fbin} file + * ({@code [num_vectors int32][dim int32]} header, then contiguous little-endian float32 rows). + * + *

Opens the file ONCE. A background thread reads the (optionally sliced) range front-to-back + * into two reusable direct buffers: while the caller consumes the current chunk, the reader fills + * the next one, so the disk read overlaps with the caller's per-vector work. {@link #get(int, + * float[])} unpacks directly into a caller-supplied array (no per-vector allocation). + * + *

Forward-only, single-consumer. {@code get} must be called with non-decreasing indices + * from a single thread, relative to the window this instance was constructed over — index 0 is + * the first vector in that window, not necessarily the first vector in the file. To read only a + * slice of a larger file (e.g. one segment of a partitioned build), construct with a {@code + * [firstVector, count)} range so each instance streams just its own portion of the file. + */ +public final class FbinVectorSource implements VectorSource { + + private static final long HEADER_BYTES = 8; + + private static final class Chunk { + final ByteBuffer buf; + final long start; + final int len; + + Chunk(ByteBuffer buf, long start, int len) { + this.buf = buf; + this.start = start; + this.len = len; + } + } + + /** Sentinel placed on the ready queue once the reader has produced the final chunk. */ + private static final Chunk POISON = new Chunk(null, -1, 0); + + private final FileChannel channel; + private final int dimension; + private final int firstVector; // absolute index (in the file) of the first vector served + private final int windowSize; // number of vectors this instance serves + private final int vectorBytes; + private final int chunkVectors; + + private final BlockingQueue free = new ArrayBlockingQueue<>(2); + private final BlockingQueue ready = new ArrayBlockingQueue<>(2); + private final Thread reader; + private volatile IOException readerError; + + private Chunk current; // consumer-owned; the chunk currently being served + private long nextExpectedRelative; // forward-only guard, relative index + + /** Reads the whole file from index 0. */ + public FbinVectorSource(Path path, int chunkSizeMB) throws IOException { + this(path, 0, -1, chunkSizeMB); + } + + /** + * Reads the contiguous range {@code [firstVector, firstVector + count)} of {@code path}, or to + * end of file if {@code count <= 0}. {@link #get} then serves indices relative to {@code + * firstVector} (0 = {@code firstVector}). + */ + public FbinVectorSource(Path path, int firstVector, int count, int chunkSizeMB) + throws IOException { + this.channel = FileChannel.open(path, StandardOpenOption.READ); + ByteBuffer header = ByteBuffer.allocate((int) HEADER_BYTES).order(ByteOrder.LITTLE_ENDIAN); + readFully(header, 0); + header.flip(); + int numVectors = header.getInt(); + this.dimension = header.getInt(); + this.vectorBytes = dimension * Float.BYTES; + this.firstVector = firstVector; + int endVector = count > 0 ? (int) Math.min((long) firstVector + count, numVectors) : numVectors; + this.windowSize = Math.max(0, endVector - firstVector); + + long chunkBytes = (long) Math.max(1, chunkSizeMB) * 1024 * 1024; + int cap = (Integer.MAX_VALUE - 16) / vectorBytes; // keep chunkVectors * vectorBytes in an int + this.chunkVectors = (int) Math.max(1, Math.min(chunkBytes / vectorBytes, cap)); + + // Two reusable direct buffers: the reader fills one while the consumer drains the other. + for (int i = 0; i < 2; i++) { + free.add( + ByteBuffer.allocateDirect(chunkVectors * vectorBytes).order(ByteOrder.LITTLE_ENDIAN)); + } + + this.reader = new Thread(this::readLoop, "fbin-prefetch-reader"); + this.reader.setDaemon(true); + this.reader.start(); + } + + @Override + public int dimensions() { + return dimension; + } + + @Override + public int size() { + return windowSize; + } + + /** Reader thread: fill chunks front-to-back, blocking on a free buffer between chunks. */ + private void readLoop() { + long next = firstVector; + long endVector = firstVector + windowSize; + try { + while (next < endVector) { + ByteBuffer buf = free.take(); + int toRead = (int) Math.min(chunkVectors, endVector - next); + buf.clear(); + buf.limit(toRead * vectorBytes); + readFully(buf, HEADER_BYTES + next * (long) vectorBytes); + ready.put(new Chunk(buf, next, toRead)); + next += toRead; + } + ready.put(POISON); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); // close() requested; stop quietly + } catch (IOException e) { + readerError = e; + try { + ready.put(POISON); // unblock the consumer so it can observe the error + } catch (InterruptedException ignored) { + Thread.currentThread().interrupt(); + } + } + } + + private void advance() throws IOException { + if (current != null) { + try { + free.put(current.buf); // hand the drained buffer back to the reader + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException("Interrupted returning chunk buffer", e); + } + current = null; + } + Chunk next; + try { + next = ready.take(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException("Interrupted awaiting next chunk", e); + } + if (next == POISON) { + if (readerError != null) { + throw new IOException("Prefetch reader failed", readerError); + } + throw new IOException("No more chunks available (unexpected EOF in prefetch)"); + } + current = next; + } + + @Override + public void get(int index, float[] dst) throws IOException { + if (index < 0 || index >= windowSize) { + throw new IndexOutOfBoundsException( + "Index " + index + " out of bounds [0, " + windowSize + ")"); + } + if (index < nextExpectedRelative) { + throw new UnsupportedOperationException( + "FbinVectorSource requires forward-only sequential access; got index " + + index + + " before previously served index " + + (nextExpectedRelative - 1)); + } + nextExpectedRelative = index + 1; + long absolute = firstVector + index; + while (current == null || absolute >= current.start + current.len) { + advance(); + } + int base = (int) (absolute - current.start) * vectorBytes; + for (int i = 0; i < dimension; i++) { + dst[i] = current.buf.getFloat(base + i * Float.BYTES); + } + } + + /** Convenience allocating variant (e.g. for a one-off query vector). */ + public float[] get(int index) throws IOException { + float[] dst = new float[dimension]; + get(index, dst); + return dst; + } + + private void readFully(ByteBuffer buf, long position) throws IOException { + long pos = position; + while (buf.hasRemaining()) { + int n = channel.read(buf, pos); + if (n < 0) { + throw new IOException("Unexpected EOF reading at position " + pos); + } + pos += n; + } + } + + @Override + public void close() throws IOException { + reader.interrupt(); + channel.close(); + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java index 8e2d9a70e2..dd1c516ebd 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java @@ -8,6 +8,8 @@ import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.quantizeFloatVectorsToBinary; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.quantizeFloatVectorsToScalar; +import com.nvidia.cuvs.CuVSHostMatrix; +import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.QuantizationType; import java.io.IOException; import java.util.List; @@ -23,18 +25,64 @@ public class FieldWriter extends KnnFieldVectorsWriter { RamUsageEstimator.shallowSizeOfInstance(FieldWriter.class); private final FieldInfo fieldInfo; + private final int dimension; private final FlatFieldVectorsWriter flatFieldVectorsWriter; private int lastDocID = -1; private QuantizationType quantizationType; + /** + * Native-buffering state. When {@code numInputVectors > 0} the writer streams incoming vectors + * directly into a native host matrix ({@link CuVSMatrix#hostBuilder}) instead of accumulating a + * {@code List} on the Java heap via {@link #flatFieldVectorsWriter}. That matrix is reused + * as the CAGRA build input, so the full dataset is never held twice (heap list + native copy). + * + *

The matrix is preallocated for exactly {@code numInputVectors} rows, so the hint must equal + * the number of vectors actually added (validated at build time in the caller). + * + *

Not yet supported for quantized fields ({@code quantizationType != QuantizationType.NONE}). + */ + private final int numInputVectors; + + private final boolean nativeBuffering; + private final CuVSMatrix.Builder hostMatrixBuilder; + private final DocsWithFieldSet nativeDocsWithField; + private int nativeCount; + private CuVSHostMatrix builtMatrix; + @SuppressWarnings("unchecked") public FieldWriter( QuantizationType quantizationType, FieldInfo fieldInfo, FlatFieldVectorsWriter flatFieldVectorsWriter) { + this(quantizationType, fieldInfo, flatFieldVectorsWriter, 0); + } + + @SuppressWarnings("unchecked") + public FieldWriter( + QuantizationType quantizationType, + FieldInfo fieldInfo, + FlatFieldVectorsWriter flatFieldVectorsWriter, + int numInputVectors) { this.quantizationType = quantizationType; this.fieldInfo = fieldInfo; + this.dimension = fieldInfo.getVectorDimension(); this.flatFieldVectorsWriter = (FlatFieldVectorsWriter) flatFieldVectorsWriter; + this.numInputVectors = numInputVectors; + this.nativeBuffering = numInputVectors > 0; + if (nativeBuffering) { + if (quantizationType != QuantizationType.NONE) { + throw new UnsupportedOperationException( + "AcceleratedHNSWParams.numInputVectors (native flat buffering) does not support" + + " quantized fields; unset it (0) to use the heap-buffered path"); + } + // Preallocates one contiguous native region of numInputVectors * dimension * 4 bytes. + this.hostMatrixBuilder = + CuVSMatrix.hostBuilder(numInputVectors, dimension, CuVSMatrix.DataType.FLOAT); + this.nativeDocsWithField = new DocsWithFieldSet(); + } else { + this.hostMatrixBuilder = null; + this.nativeDocsWithField = null; + } } @Override @@ -46,7 +94,27 @@ public void addValue(int docID, Object vectorValue) throws IOException { + "\" appears more than once in this document (only one value is allowed per" + " field)"); } - flatFieldVectorsWriter.addValue(docID, (float[]) vectorValue); + if (nativeBuffering) { + if (nativeCount >= numInputVectors) { + throw new IllegalStateException( + "Buffered vectors (" + + (nativeCount + 1) + + ") exceed numInputVectors (" + + numInputVectors + + ") for field \"" + + fieldInfo.name + + "\". This usually means more vectors arrived than the numInputVectors hint" + + " promised (e.g. a merge or a later flush cycle reused this config); see" + + " AcceleratedHNSWParams.Builder#withNumInputVectors."); + } + // hostMatrixBuilder.addVector validates the dimension and performs the native row copy. + hostMatrixBuilder.addVector((float[]) vectorValue); + nativeDocsWithField.add(docID); + nativeCount++; + lastDocID = docID; + } else { + flatFieldVectorsWriter.addValue(docID, (float[]) vectorValue); + } } List getByteVectors() { @@ -68,7 +136,41 @@ FieldInfo fieldInfo() { } DocsWithFieldSet getDocsWithFieldSet() { - return flatFieldVectorsWriter.getDocsWithFieldSet(); + return nativeBuffering ? nativeDocsWithField : flatFieldVectorsWriter.getDocsWithFieldSet(); + } + + /** Whether this writer streams vectors into a native host matrix (hint path). */ + boolean isNativeBuffering() { + return nativeBuffering; + } + + /** + * The native host matrix holding the buffered vectors. Valid only when {@link #isNativeBuffering()} + * is true. The matrix is built once and cached; the caller owns closing it via + * {@link #releaseNativeBuffer()} once the CAGRA build has consumed it. + */ + CuVSHostMatrix getHostMatrix() { + if (builtMatrix == null) { + builtMatrix = hostMatrixBuilder.build(); + } + return builtMatrix; + } + + /** Number of vectors buffered so far (hint path). */ + int getNativeVectorCount() { + return nativeCount; + } + + int dimension() { + return dimension; + } + + /** Closes the native host matrix. Safe to call multiple times; a no-op in the non-native path. */ + void releaseNativeBuffer() { + if (nativeBuffering) { + getHostMatrix().close(); + builtMatrix = null; + } } @Override @@ -78,6 +180,8 @@ public Object copyValue(Object vectorValue) { @Override public long ramBytesUsed() { - return SHALLOW_SIZE + flatFieldVectorsWriter.ramBytesUsed(); + // The native host matrix is off-heap and intentionally excluded from Lucene's heap RAM + // accounting. + return SHALLOW_SIZE + (nativeBuffering ? 0 : flatFieldVectorsWriter.ramBytesUsed()); } } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java index 7e9f888e32..c4f8a3ea8f 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java @@ -6,10 +6,16 @@ import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; +import com.nvidia.cuvs.CuVSDeviceMatrix; +import com.nvidia.cuvs.CuVSHostMatrix; import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.RowView; import java.util.ArrayList; import java.util.List; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; import org.apache.lucene.util.hnsw.HnswGraph; import org.apache.lucene.util.hnsw.NeighborArray; @@ -38,9 +44,14 @@ public class GPUBuiltHnswGraph extends HnswGraph { * @param dimensions the vector dimension * @param layerNodes the nodes on the layer * @param layerAdjacencies adjacency list + * @param numThreads threads to use for materializing the adjacency (1 = serial) */ public GPUBuiltHnswGraph( - int size, int dimensions, List layerNodes, List layerAdjacencies) { + int size, + int dimensions, + List layerNodes, + List layerAdjacencies, + int numThreads) { this.size = size; this.dimensions = dimensions; @@ -50,38 +61,104 @@ public GPUBuiltHnswGraph( // Process Layer 0 (base layer with all nodes) CuVSMatrix layer0Adjacency = layerAdjacencies.get(0); - this.layer0Neighbors = fillNeighborArray(layer0Adjacency, size); + this.layer0Neighbors = fillNeighborArray(layer0Adjacency, size, numThreads); // Process higher layers (1 to numLevels-1) for (int level = 1; level < numLevels; level++) { int[] nodes = layerNodes.get(level); CuVSMatrix adjacency = layerAdjacencies.get(level); this.layerNodes.add(nodes); - this.layerNeighbors.add(fillNeighborArray(adjacency, nodes.length)); + this.layerNeighbors.add(fillNeighborArray(adjacency, nodes.length, numThreads)); } } + /** Node count below which parallel materialization is not worth the thread overhead. */ + private static final int PARALLEL_MIN_NODES = 1 << 16; + /** - * Fills the neighbor array using the adjacency matrix. + * Materializes the adjacency matrix into on-heap {@link NeighborArray}s, one per node. + * + *

The serial path reads the adjacency directly (a device matrix's {@code getRow} is safe + * single-threaded). The parallel path cannot: the CAGRA layer-0 adjacency is a device matrix whose + * {@code getRow} uses a shared, stateful buffered reader that is not safe for concurrent access, so + * it is pulled to host once (a single bulk device->host copy) before materializing disjoint node + * ranges concurrently. Host matrices (the upper layers, built via {@link CuVSMatrix#ofArray}) are + * read directly in both paths. * * @param adjacency instance of adjacency CuVSMatrix * @param size the number of nodes + * @param numThreads threads to use (1, or fewer than {@value #PARALLEL_MIN_NODES} nodes = serial) * @return the NeighborArray */ - private NeighborArray[] fillNeighborArray(CuVSMatrix adjacency, int size) { + private static NeighborArray[] fillNeighborArray(CuVSMatrix adjacency, int size, int numThreads) { NeighborArray[] neighbors = new NeighborArray[size]; - for (int i = 0; i < size; i++) { - RowView rv = adjacency.getRow(i); + if (numThreads <= 1 || size < PARALLEL_MIN_NODES) { + fillNeighborRange(adjacency, neighbors, 0, size); + return neighbors; + } + CuVSMatrix source = adjacency; + CuVSHostMatrix hostCopy = null; + if (adjacency instanceof CuVSDeviceMatrix deviceAdjacency) { + hostCopy = deviceAdjacency.toHost(); + source = hostCopy; + } + try { + fillNeighborArrayParallel(source, neighbors, size, numThreads); + return neighbors; + } finally { + if (hostCopy != null) { + hostCopy.close(); + } + } + } + + /** + * Materializes disjoint node ranges concurrently. Each thread writes its own slots of {@code + * neighbors} and its own {@link NeighborArray} instances, so no synchronization is needed; {@code + * source} must be a host matrix (stateless {@code getRow}). + */ + private static void fillNeighborArrayParallel( + CuVSMatrix source, NeighborArray[] neighbors, int size, int numThreads) { + ExecutorService pool = Executors.newFixedThreadPool(numThreads); + try { + int perThread = (size + numThreads - 1) / numThreads; + List> futures = new ArrayList<>(numThreads); + for (int t = 0; t < numThreads; t++) { + final int start = t * perThread; + final int end = Math.min(start + perThread, size); + if (start >= end) { + break; + } + futures.add(pool.submit(() -> fillNeighborRange(source, neighbors, start, end))); + } + for (Future f : futures) { + f.get(); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Interrupted during parallel HNSW conversion", e); + } catch (ExecutionException e) { + throw new RuntimeException("Failed during parallel HNSW conversion", e.getCause()); + } finally { + pool.shutdown(); + } + } + + /** Fills {@code neighbors[start, end)} from the adjacency rows. */ + private static void fillNeighborRange( + CuVSMatrix source, NeighborArray[] neighbors, int start, int end) { + for (int i = start; i < end; i++) { + RowView rv = source.getRow(i); if (rv != null && rv.size() > 0) { - neighbors[i] = new NeighborArray((int) rv.size(), true); + NeighborArray na = new NeighborArray((int) rv.size(), true); for (int j = 0; j < rv.size(); j++) { - neighbors[i].addInOrder(rv.getAsInt(j), 1.0f - (j * 0.001f)); + na.addInOrder(rv.getAsInt(j), 1.0f - (j * 0.001f)); } + neighbors[i] = na; } else { neighbors[i] = new NeighborArray(0, true); } } - return neighbors; } /** diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java index 61c9b41c3a..050f626657 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java @@ -74,7 +74,10 @@ public Lucene99AcceleratedHNSWVectorsFormat(AcceleratedHNSWParams acceleratedHNS */ @Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { - var flatWriter = FLAT_VECTORS_FORMAT.fieldsWriter(state); + // In hint mode the accelerated writer owns the flat .vec/.vemf files, so the Lucene flat writer + // must not be created (it would open the same outputs). The fallback path below still needs it. + boolean nativeMode = isSupported() && acceleratedHNSWParams.getNumInputVectors() > 0; + var flatWriter = nativeMode ? null : FLAT_VECTORS_FORMAT.fieldsWriter(state); if (isSupported()) { log.log(Level.FINE, "cuVS is supported so using the Lucene99AcceleratedHNSWVectorsWriter"); return new Lucene99AcceleratedHNSWVectorsWriter(state, acceleratedHNSWParams, flatWriter); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index 13edc64975..f168919770 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -16,12 +16,12 @@ import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_NAME; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.closeCuVSResourcesInstance; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.getCuVSResourcesInstance; -import static com.nvidia.cuvs.lucene.Utils.createListFromMergedVectors; import static org.apache.lucene.index.VectorEncoding.FLOAT32; import static org.apache.lucene.util.RamUsageEstimator.shallowSizeOfInstance; import com.nvidia.cuvs.CagraIndex; import com.nvidia.cuvs.CagraIndexParams; +import com.nvidia.cuvs.CuVSHostMatrix; import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.QuantizationType; import java.io.IOException; @@ -34,12 +34,16 @@ import org.apache.lucene.codecs.hnsw.FlatVectorsWriter; import org.apache.lucene.index.DocsWithFieldSet; import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FloatVectorValues; import org.apache.lucene.index.IndexFileNames; +import org.apache.lucene.index.KnnVectorValues; import org.apache.lucene.index.MergeState; import org.apache.lucene.index.SegmentWriteState; import org.apache.lucene.index.Sorter; import org.apache.lucene.index.Sorter.DocMap; +import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.util.Bits; import org.apache.lucene.util.IOUtils; import org.apache.lucene.util.InfoStream; @@ -61,6 +65,18 @@ public class Lucene99AcceleratedHNSWVectorsWriter extends KnnVectorsWriter { private final FlatVectorsWriter flatVectorsWriter; private final List fields = new ArrayList<>(); private final InfoStream infoStream; + + /** + * Hint-path state. When {@code numInputVectors > 0}, vectors are streamed into a native host + * matrix (see {@link FieldWriter}) rather than a heap {@code List}, and the flat + * {@code .vec}/{@code .vemf} files are written by {@link #nativeFlat} instead of by + * {@link #flatVectorsWriter} (which is {@code null} in this mode). Supports only the unsorted + * single-segment flush path; merges and index-sorted flushes are rejected. + */ + private final int numInputVectors; + + private final boolean nativeMode; + private final NativeFlatVectorsWriter nativeFlat; private IndexOutput hnswMeta = null; private IndexOutput hnswVectorIndex = null; private String vemFileName; @@ -93,6 +109,13 @@ public Lucene99AcceleratedHNSWVectorsWriter( this.flatVectorsWriter = flatVectorsWriter; this.infoStream = state.infoStream; this.acceleratedHNSWParams = acceleratedHNSWParams; + this.numInputVectors = acceleratedHNSWParams.getNumInputVectors(); + this.nativeMode = numInputVectors > 0; + if (nativeMode && state.segmentInfo.getIndexSort() != null) { + throw new IllegalArgumentException( + "AcceleratedHNSWParams.numInputVectors (native flat buffering) does not support" + + " index-sorted segments; unset it (0) to use the heap-buffered path"); + } vemFileName = IndexFileNames.segmentFileName( state.segmentInfo.name, state.segmentSuffix, HNSW_META_CODEC_EXT); @@ -114,6 +137,9 @@ public Lucene99AcceleratedHNSWVectorsWriter( VERSION_CURRENT, state.segmentInfo.getId(), state.segmentSuffix); + // In hint mode we own the flat files; the Lucene flat writer must be absent to avoid opening + // the same .vec/.vemf outputs. + nativeFlat = nativeMode ? new NativeFlatVectorsWriter(state) : null; success = true; printInfoStream(infoStream, COMPONENT, "Lucene99AcceleratedHNSWVectorsWriter is initialized"); } finally { @@ -132,6 +158,14 @@ public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException if (encoding != FLOAT32) { throw new IllegalArgumentException("Expected float32, got:" + encoding); } + if (nativeMode) { + // Buffer directly into a native host matrix; return the FieldWriter itself so Lucene routes + // addValue() here rather than to a (nonexistent) Lucene flat field writer. + var cuvsFieldWriter = + new FieldWriter(QuantizationType.NONE, fieldInfo, null, numInputVectors); + fields.add(cuvsFieldWriter); + return cuvsFieldWriter; + } var writer = Objects.requireNonNull(flatVectorsWriter.addField(fieldInfo)); var cuvsFieldWriter = new FieldWriter(QuantizationType.NONE, fieldInfo, writer); fields.add(cuvsFieldWriter); @@ -139,7 +173,8 @@ public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException } /** - * Builds the intermediate CAGRA index and builds and writes the HNSW index. + * Flush/sorting path: builds a host matrix from the heap vectors, then delegates + * to {@link #writeFieldInternal(FieldInfo, CuVSMatrix)}. * * @param fieldInfo instance of FieldInfo that has the field description * @param vectors vectors to index @@ -154,34 +189,55 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro writeSingleVectorGraph(fieldInfo, vectors); return; } - try { - CuVSMatrix dataset = - Utils.createFloatMatrix( - vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance()); + CuVSMatrix dataset = Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); + writeFieldInternal(fieldInfo, dataset); + } + /** + * Builds the intermediate CAGRA index and builds and writes the HNSW index. + * Single implementation used by both the flush and merge paths. The dataset is a + * {@link CuVSMatrix} (host-backed on the merge path) so the full set of vectors is + * never double-materialised on the Java heap. + * + * @param fieldInfo instance of FieldInfo that has the field description + * @param dataset matrix of all vectors to index + * @throws IOException + */ + private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws IOException { + int size = (int) dataset.size(); + if (size == 0) { + writeEmpty(fieldInfo, hnswMeta); + return; + } + if (size < 2) { + float[] buf = new float[fieldInfo.getVectorDimension()]; + dataset.getRow(0).toArray(buf); + writeSingleVectorGraph(fieldInfo, List.of(buf)); + return; + } + try { CagraIndexParams params = CagraIndexParamsFactory.create(acceleratedHNSWParams, dataset.size(), dataset.columns()); - CagraIndex cagraIndex = CagraIndex.newBuilder(getCuVSResourcesInstance()) .withDataset(dataset) .withIndexParams(params) .build(); CuVSMatrix adjacencyListMatrix = cagraIndex.getGraph(); - int size = (int) dataset.size(); int dimensions = fieldInfo.getVectorDimension(); GPUBuiltHnswGraph hnswGraph = createMultiLayerHnswGraph( fieldInfo, - size, dimensions, adjacencyListMatrix, - vectors, + dataset, acceleratedHNSWParams.getHnswLayers(), params, - QuantizationType.NONE); + QuantizationType.NONE, + acceleratedHNSWParams.getWriterThreads()); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + int[][] graphLevelNodeOffsets = + writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; writeMeta( hnswVectorIndex, @@ -203,6 +259,17 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro */ @Override public void flush(int maxDoc, DocMap sortMap) throws IOException { + if (nativeMode) { + if (sortMap != null) { + throw new UnsupportedOperationException( + "AcceleratedHNSWParams.numInputVectors (native flat buffering) does not support" + + " index-sorted segments; unset it (0) to enable the sorted flush path"); + } + for (var field : fields) { + writeFieldNative(field, maxDoc); + } + return; + } flatVectorsWriter.flush(maxDoc, sortMap); for (var field : fields) { if (sortMap == null) { @@ -213,6 +280,37 @@ public void flush(int maxDoc, DocMap sortMap) throws IOException { } } + /** + * Hint-path flush for a single field: writes the flat {@code .vec}/{@code .vemf} from the native + * host matrix, builds the CAGRA/HNSW graph from the same matrix, then releases the matrix. Both + * consumers read the matrix before it is closed. + */ + private void writeFieldNative(FieldWriter fieldData, int maxDoc) throws IOException { + int count = fieldData.getNativeVectorCount(); + if (count != numInputVectors) { + throw new IllegalStateException( + "numInputVectors (" + + numInputVectors + + ") must equal the number of vectors added (" + + count + + ") for field \"" + + fieldData.fieldInfo().name + + "\"; the native host matrix is sized for the hint exactly. This usually means" + + " IndexWriterConfig's auto-flush wasn't disabled (setMaxBufferedDocs /" + + " setRAMBufferSizeMB(DISABLE_AUTO_FLUSH)), so a flush landed before exactly" + + " numInputVectors vectors were added; see AcceleratedHNSWParams.Builder" + + "#withNumInputVectors."); + } + FieldInfo fieldInfo = fieldData.fieldInfo(); + try { + CuVSHostMatrix dataset = fieldData.getHostMatrix(); + nativeFlat.writeField(fieldInfo, dataset, maxDoc, fieldData.getDocsWithFieldSet()); + writeFieldInternal(fieldInfo, dataset); + } finally { + fieldData.releaseNativeBuffer(); + } + } + /** * Builds the index and writes it to the disk. * @@ -256,7 +354,8 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) int dimensions = fieldInfo.getVectorDimension(); GPUBuiltHnswGraph hnswGraph = createSingleVectorHnswGraph(size, dimensions); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + int[][] graphLevelNodeOffsets = + writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; writeMeta( hnswVectorIndex, @@ -273,13 +372,63 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) } /** - * Create combined data set for the merged segment and call writeFieldInternal. + * Streams merged vectors directly into a native host-memory matrix (CuVSHostMatrix) + * without materialising a List on the Java heap, then calls writeFieldInternal. + * This avoids the double-copy OOM (heap list + native matrix simultaneously) that + * occurs when force-merging large segments. */ private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws IOException { try { - List dataset = - createListFromMergedVectors( - KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState)); + // FloatVectorValues#size() on the merged view is the raw sum of every source segment's + // on-disk vector count (MergedVectorValues.MergedFloat32VectorValues computes it once at + // construction from each sub-reader's unfiltered size) -- NOT the number of live + // (non-deleted) vectors the iterator below will actually yield, which is what + // CuVSMatrix.hostBuilder needs since it preallocates a fixed-size native buffer. Using + // size() here under-fills that buffer whenever the merge drops deleted docs, leaving the + // graph built over more rows than were actually populated. + // + // size() IS trustworthy when no segment being merged has any deletions: per-segment vector + // counts already exclude docs without a value for this field (sparse fields are handled at + // the single-segment level, independent of deletions), so the raw sum equals the live count + // in that case and the extra counting pass below can be skipped. + boolean anySegmentHasDeletions = false; + for (Bits liveDocs : mergeState.liveDocs) { + if (liveDocs != null) { + anySegmentHasDeletions = true; + break; + } + } + + int size; + if (anySegmentHasDeletions) { + // Count the live vectors via a throwaway iteration first (mergeFloatVectorValues + // constructs a fresh, independent view each call, so this doesn't disturb the real build + // pass below). + size = 0; + FloatVectorValues counting = + KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); + KnnVectorValues.DocIndexIterator countingIt = counting.iterator(); + for (int doc = countingIt.nextDoc(); + doc != DocIdSetIterator.NO_MORE_DOCS; + doc = countingIt.nextDoc()) { + size++; + } + } else { + size = + KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState) + .size(); + } + + FloatVectorValues mergedVectors = + KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); + int dims = fieldInfo.getVectorDimension(); + CuVSMatrix.Builder builder = + CuVSMatrix.hostBuilder(size, dims, CuVSMatrix.DataType.FLOAT); + KnnVectorValues.DocIndexIterator it = mergedVectors.iterator(); + for (int doc = it.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = it.nextDoc()) { + builder.addVector(mergedVectors.vectorValue(it.index())); + } + CuVSHostMatrix dataset = builder.build(); writeFieldInternal(fieldInfo, dataset); } catch (Throwable t) { Utils.handleThrowable(t); @@ -291,6 +440,11 @@ private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws */ @Override public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOException { + if (nativeMode) { + throw new UnsupportedOperationException( + "AcceleratedHNSWParams.numInputVectors (native flat buffering) supports only the" + + " unsorted single-segment flush path; unset it (0) to enable merges"); + } flatVectorsWriter.mergeOneField(fieldInfo, mergeState); vectorBasedMerge(fieldInfo, mergeState); } @@ -304,7 +458,11 @@ public void finish() throws IOException { throw new IllegalStateException("already finished"); } finished = true; - flatVectorsWriter.finish(); + if (nativeMode) { + nativeFlat.finish(); + } else { + flatVectorsWriter.finish(); + } if (hnswMeta != null) { // write end of fields marker hnswMeta.writeInt(-1); @@ -321,7 +479,7 @@ public void finish() throws IOException { @Override public void close() throws IOException { printInfoStream(infoStream, COMPONENT, "Closing resources"); - IOUtils.close(hnswMeta, hnswVectorIndex, flatVectorsWriter); + IOUtils.close(hnswMeta, hnswVectorIndex, flatVectorsWriter, nativeFlat); closeCuVSResourcesInstance(); } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java index 87907d2cbb..cf6d19f7ca 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java @@ -155,8 +155,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throw int dimensions = fieldInfo.getVectorDimension(); int bytesPerVector = (dimensions + 7) / 8; - CuVSMatrix dataset = - Utils.createByteMatrix(vectors, bytesPerVector, getCuVSResourcesInstance()); + CuVSMatrix dataset = Utils.createByteMatrix(vectors, bytesPerVector); if (dataset.size() < 2) { writeSingleVectorGraph(fieldInfo, vectors); @@ -179,17 +178,18 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throw GPUBuiltHnswGraph hnswGraph = createMultiLayerHnswGraph( fieldInfo, - size, dimensions, adjacencyListMatrix, - vectors, + dataset, acceleratedHNSWParams.getHnswLayers(), params, - QuantizationType.BINARY); + QuantizationType.BINARY, + acceleratedHNSWParams.getWriterThreads()); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + int[][] graphLevelNodeOffsets = + writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; // Write metadata @@ -277,7 +277,8 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + int[][] graphLevelNodeOffsets = + writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index 7141af56ee..aebfbea7c6 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -181,8 +181,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE } // Create CuVSMatrix with BYTE data type (unsigned bytes) - CuVSMatrix dataset = - Utils.createByteMatrix(unsignedVectors, dimensions, getCuVSResourcesInstance()); + CuVSMatrix dataset = Utils.createByteMatrix(unsignedVectors, dimensions); if (dataset.size() < 2) { writeSingleVectorGraph(fieldInfo, unsignedVectors); @@ -204,18 +203,19 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE GPUBuiltHnswGraph hnswGraph = createMultiLayerHnswGraph( fieldInfo, - size, dimensions, adjacencyListMatrix, - unsignedVectors, + dataset, acceleratedHNSWParams.getHnswLayers(), params, - QuantizationType.SCALAR); + QuantizationType.SCALAR, + acceleratedHNSWParams.getWriterThreads()); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + int[][] graphLevelNodeOffsets = + writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; @@ -302,7 +302,8 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + int[][] graphLevelNodeOffsets = + writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; // Write metadata diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatVectorsWriter.java new file mode 100644 index 0000000000..09b4fe3021 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatVectorsWriter.java @@ -0,0 +1,199 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.CuVSHostMatrix; +import java.io.Closeable; +import java.io.IOException; +import java.lang.foreign.MemorySegment; +import java.lang.foreign.ValueLayout; +import java.nio.ByteOrder; +import org.apache.lucene.codecs.CodecUtil; +import org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration; +import org.apache.lucene.index.DocsWithFieldSet; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.IndexFileNames; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.util.IOUtils; + +/** + * Writes the flat vector files ({@code .vec} data + {@code .vemf} meta) directly from a native host + * matrix, byte-for-byte compatible with Lucene's {@code Lucene99FlatVectorsWriter} so the stock + * {@code Lucene99FlatVectorsReader} can read them. + * + *

This is the hint-path counterpart to delegating to Lucene's {@code FlatVectorsWriter}: the + * accelerated writer streams vectors into a {@link CuVSHostMatrix} during indexing (see + * {@link FieldWriter}) and never materialises the full dataset as a {@code List} on the + * Java heap, so the flat file is written here from that native matrix instead. + * + *

Ported code — pinned to the {@code Lucene99} format. The dense float32 layout (format + * constants, header, meta field order, and footer) is transcribed from {@code + * Lucene99FlatVectorsFormat}/{@code Lucene99FlatVectorsWriter} (tag {@code + * releases/lucene/10.2.0}): + * + *

    + *
  • {@code org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat} — format constants: + * https://github.com/apache/lucene/blob/releases/lucene/10.2.0/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99FlatVectorsFormat.java + *
  • {@code org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsWriter} — write sequence: + * https://github.com/apache/lucene/blob/releases/lucene/10.2.0/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99FlatVectorsWriter.java + *
+ * + *

{@code Lucene99} is a frozen, versioned codec name: its constants and byte layout stay fixed + * for the life of the codec, and files written to it stay readable by the stock {@code + * Lucene99FlatVectorsReader} — via {@code lucene-backward-codecs} once a newer default codec ships + * — for the rest of the Lucene 10.x major version. The constants below (codec names, extensions, + * {@code VERSION_CURRENT}) are the fixed handshake the reader string-matches ({@code + * CodecUtil.checkIndexHeader}) and file-extension-matches against, so they stay as transcribed here + * for the life of the {@code Lucene99} codec. + * + *

On a {@code lucene-core} version bump: run {@code TestNativeFlatVectorsWriterRoundTrip} + * to confirm the stock reader on the new classpath still accepts what this class writes; it builds + * a small index and asserts every vector round-trips byte-exact through {@code + * Lucene99FlatVectorsReader}. + * + *

Moving to a newer flat-vector format (e.g. once {@code Lucene99} support is dropped + * ahead of an 11.x major bump): diff the new format's writer against the one linked above and + * re-derive the {@code writeField}/{@code writeMeta} byte sequence (header, meta field order, + * footer) here. This class writes directly from native memory to avoid the per-vector {@code + * FloatVectorValues} indirection Lucene's own writer requires — that's the reason to keep + * hand-porting the format rather than delegating to it. Then update the codec name/extension/version + * constants below and the Lucene version and links in this javadoc. + */ +final class NativeFlatVectorsWriter implements Closeable { + + // Mirrors org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat (10.2.0) so the standard + // Lucene99FlatVectorsReader accepts the header/codec of the files written here. + private static final String META_CODEC_NAME = "Lucene99FlatVectorsFormatMeta"; + private static final String VECTOR_DATA_CODEC_NAME = "Lucene99FlatVectorsFormatData"; + private static final String META_EXTENSION = "vemf"; + private static final String VECTOR_DATA_EXTENSION = "vec"; + private static final int VERSION_CURRENT = 0; + private static final int DIRECT_MONOTONIC_BLOCK_SHIFT = 16; + + // Little-endian float layout matching Lucene's on-disk .vec byte order. Must be UNALIGNED: the + // destination is a heap byte[]-backed MemorySegment whose max alignment is 1 byte, so a 4-byte + // aligned JAVA_FLOAT layout is rejected with "incompatible with alignment constraints". + private static final ValueLayout.OfFloat LE_FLOAT = + ValueLayout.JAVA_FLOAT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN); + + // Byte granularity for a single writeBytes call; bounds the transient encode buffer. + private static final int CHUNK_BYTES = 1 << 18; // 256 KiB + + private final IndexOutput meta; + private final IndexOutput vectorData; + private boolean finished; + + NativeFlatVectorsWriter(SegmentWriteState state) throws IOException { + String metaFileName = + IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, META_EXTENSION); + String vectorDataFileName = + IndexFileNames.segmentFileName( + state.segmentInfo.name, state.segmentSuffix, VECTOR_DATA_EXTENSION); + boolean success = false; + try { + meta = state.directory.createOutput(metaFileName, state.context); + vectorData = state.directory.createOutput(vectorDataFileName, state.context); + CodecUtil.writeIndexHeader( + meta, META_CODEC_NAME, VERSION_CURRENT, state.segmentInfo.getId(), state.segmentSuffix); + CodecUtil.writeIndexHeader( + vectorData, + VECTOR_DATA_CODEC_NAME, + VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + success = true; + } finally { + if (success == false) { + IOUtils.closeWhileHandlingException(this); + } + } + } + + /** + * Writes one dense float32 field: the raw vectors to {@code .vec} and the field metadata (plus the + * ordinal-to-doc mapping) to {@code .vemf}. Vectors are read from {@code matrix} in ordinal order, + * which matches the ascending-docID order in which {@code docsWithField} was populated. + * + * @param field the field being written + * @param matrix the native host matrix holding {@code docsWithField.cardinality()} rows of {@code + * field.getVectorDimension()} floats each + * @param maxDoc the segment's maxDoc, used to build the ordinal-to-doc mapping + * @param docsWithField the set of docs that have a value for this field + */ + void writeField( + FieldInfo field, CuVSHostMatrix matrix, int maxDoc, DocsWithFieldSet docsWithField) + throws IOException { + // Mirrors Lucene99FlatVectorsWriter#writeField (see class-level version pin). + int count = docsWithField.cardinality(); + int dim = field.getVectorDimension(); + long vectorDataOffset = vectorData.alignFilePointer(Float.BYTES); + writeFloat32Vectors(matrix, count, dim); + long vectorDataLength = vectorData.getFilePointer() - vectorDataOffset; + writeMeta(field, maxDoc, count, vectorDataOffset, vectorDataLength, docsWithField); + } + + private void writeFloat32Vectors(CuVSHostMatrix matrix, int count, int dim) throws IOException { + int rowBytes = dim * Float.BYTES; + int chunkRows = Math.max(1, CHUNK_BYTES / rowBytes); + byte[] chunk = new byte[chunkRows * rowBytes]; + MemorySegment chunkSeg = MemorySegment.ofArray(chunk); + float[] rowBuf = new float[dim]; + int r = 0; + for (int ord = 0; ord < count; ord++) { + matrix.getRow(ord).toArray(rowBuf); // native -> heap float[] (bulk) + MemorySegment.copy(rowBuf, 0, chunkSeg, LE_FLOAT, (long) r * rowBytes, dim); // -> LE bytes + if (++r == chunkRows) { + vectorData.writeBytes(chunk, r * rowBytes); + r = 0; + } + } + if (r > 0) { + vectorData.writeBytes(chunk, r * rowBytes); + } + } + + private void writeMeta( + FieldInfo field, + int maxDoc, + int count, + long vectorDataOffset, + long vectorDataLength, + DocsWithFieldSet docsWithField) + throws IOException { + // Mirrors Lucene99FlatVectorsWriter#writeMeta (see class-level version pin); field order is + // load-bearing and must match Lucene's reader. + meta.writeInt(field.number); + meta.writeInt(field.getVectorEncoding().ordinal()); + meta.writeInt(field.getVectorSimilarityFunction().ordinal()); + meta.writeVLong(vectorDataOffset); + meta.writeVLong(vectorDataLength); + meta.writeVInt(field.getVectorDimension()); + meta.writeInt(count); + OrdToDocDISIReaderConfiguration.writeStoredMeta( + DIRECT_MONOTONIC_BLOCK_SHIFT, meta, vectorData, count, maxDoc, docsWithField); + } + + /** Writes the end-of-fields marker and footers. Mirrors {@code Lucene99FlatVectorsWriter.finish}. */ + void finish() throws IOException { + if (finished) { + throw new IllegalStateException("already finished"); + } + finished = true; + if (meta != null) { + meta.writeInt(-1); + CodecUtil.writeFooter(meta); + } + if (vectorData != null) { + CodecUtil.writeFooter(vectorData); + } + } + + @Override + public void close() throws IOException { + IOUtils.close(meta, vectorData); + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java index e4a20d2b4d..0fa96ded3e 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java @@ -43,82 +43,50 @@ static void handleThrowable(Throwable t) throws IOException { } /** - * A method to build a CuVSMatrix from a list of float vectors. + * Builds a host-memory CuVSMatrix from a list of float vectors. * - * Uses CuVSMatrix.Builder to copy vectors directly to device memory - * without creating intermediate heap arrays. + *

Copies vectors directly into a native host matrix via {@link CuVSMatrix#hostBuilder}, + * without creating an intermediate {@code float[][]} on the heap. * * @param data The float vectors - * @param dimensions The number float elements in each vector - * @param resources The CuVS resources for device matrix creation - * @return an instance of CuVSMatrix + * @param dimensions The number of float elements in each vector + * @return a host-memory CuVSMatrix */ - static CuVSMatrix createFloatMatrix(List data, int dimensions, CuVSResources resources) { - // Use Builder pattern to avoid intermediate float[][] allocation - // and copy directly from List to device memory + static CuVSMatrix createFloatMatrix(List data, int dimensions) { CuVSMatrix.Builder builder = - CuVSMatrix.deviceBuilder( - resources, - data.size(), // rows (number of vectors) - dimensions, // columns (vector dimension) - CuVSMatrix.DataType.FLOAT); - - // Add vectors one by one - builder copies directly to device memory + CuVSMatrix.hostBuilder(data.size(), dimensions, CuVSMatrix.DataType.FLOAT); for (float[] vector : data) { builder.addVector(vector); } - return builder.build(); } /** - * A method to build a CuVSMatrix from a list of byte vectors (for binary quantized vectors). - * - * Uses CuVSMatrix.Builder to copy vectors directly to device memory - * without creating intermediate heap arrays. + * Builds a host-memory CuVSMatrix from a list of byte vectors (e.g. quantized vectors). * * @param data The byte vectors (packed bits for binary quantization) * @param bytesPerVector The number of bytes in each vector - * @param resources The CuVS resources for device matrix creation - * @return an instance of CuVSMatrix with BYTE data type + * @return a host-memory CuVSMatrix with BYTE data type */ - static CuVSMatrix createByteMatrix( - List data, int bytesPerVector, CuVSResources resources) { - // Use Builder pattern to avoid intermediate byte[][] allocation - // and copy directly from List to device memory + static CuVSMatrix createByteMatrix(List data, int bytesPerVector) { CuVSMatrix.Builder builder = - CuVSMatrix.deviceBuilder( - resources, - data.size(), // rows (number of vectors) - bytesPerVector, // columns (bytes per vector) - CuVSMatrix.DataType.BYTE); - - // Add vectors one by one - builder copies directly to device memory + CuVSMatrix.hostBuilder(data.size(), bytesPerVector, CuVSMatrix.DataType.BYTE); for (byte[] vector : data) { builder.addVector(vector); } - return builder.build(); } /** - * A method to build a CuVSMatrix from a 2D byte array (for binary quantized vectors). + * Builds a host-memory CuVSMatrix from a 2D byte array (e.g. quantized vectors). * * @param data The 2D byte array (packed bits for binary quantization) * @param bytesPerVector The number of bytes in each vector - * @param resources The CuVS resources for device matrix creation - * @return an instance of CuVSMatrix with BYTE data type + * @return a host-memory CuVSMatrix with BYTE data type */ - static CuVSMatrix createByteMatrixFromArray( - byte[][] data, int bytesPerVector, CuVSResources resources) { + static CuVSMatrix createByteMatrixFromArray(byte[][] data, int bytesPerVector) { CuVSMatrix.Builder builder = - CuVSMatrix.deviceBuilder( - resources, - data.length, // rows (number of vectors) - bytesPerVector, // columns (bytes per vector) - CuVSMatrix.DataType.BYTE); - - // Add vectors one by one - builder copies directly to device memory + CuVSMatrix.hostBuilder(data.length, bytesPerVector, CuVSMatrix.DataType.BYTE); for (byte[] vector : data) { builder.addVector(vector); } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/VectorSource.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/VectorSource.java new file mode 100644 index 0000000000..c8803651ea --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/VectorSource.java @@ -0,0 +1,34 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import java.io.Closeable; +import java.io.IOException; + +/** + * A forward-only, single-consumer source of float vectors for {@link CagraHnswBulkIndexWriter}. + * + *

Implementations own how the underlying data is fetched (a local file, a database cursor, an + * object-store stream, ...); {@link CagraHnswBulkIndexWriter} only depends on this contract, not on + * any particular storage. + * + *

{@link #get(int, float[])} must be called with non-decreasing indices from a single thread. + * This mirrors how a bulk indexer consumes its input (front-to-back, one pass) and lets + * implementations use a simple prefetch/streaming strategy instead of arbitrary random access. + */ +public interface VectorSource extends Closeable { + + /** Number of dimensions of every vector returned by this source. */ + int dimensions(); + + /** Number of vectors available, i.e. the exclusive upper bound on indices passed to {@link #get}. */ + int size(); + + /** + * Fills {@code dst} with the vector at {@code index} (no allocation). {@code index} must be + * greater than or equal to the index passed to the previous call, if any. + */ + void get(int index, float[] dst) throws IOException; +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParamsSurface.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParamsSurface.java new file mode 100644 index 0000000000..927670bcb3 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParamsSurface.java @@ -0,0 +1,33 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static org.junit.Assert.assertFalse; + +import java.lang.reflect.Method; +import java.lang.reflect.Modifier; +import org.junit.Test; + +/** + * Regression guard for the native-flat-buffering surface narrowing: {@code + * AcceleratedHNSWParams.Builder#withNumInputVectors} must stay package-private, reachable only + * from within {@code com.nvidia.cuvs.lucene} (i.e. only by {@link CagraHnswBulkIndexWriter}), not from + * a generic external Lucene codec user. A future change that accidentally re-widens this method to + * {@code public} would reopen exactly the footgun this package's design is meant to close, without + * necessarily being caught by any functional test -- this test exists to catch that specific + * mistake directly. + */ +public class TestAcceleratedHNSWParamsSurface { + + @Test + public void testWithNumInputVectorsIsNotPublic() throws NoSuchMethodException { + Method method = + AcceleratedHNSWParams.Builder.class.getDeclaredMethod("withNumInputVectors", int.class); + assertFalse( + "AcceleratedHNSWParams.Builder#withNumInputVectors must not be public; it should only be" + + " reachable from within com.nvidia.cuvs.lucene (CagraHnswBulkIndexWriter)", + Modifier.isPublic(method.getModifiers())); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBulkIndexWriter.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBulkIndexWriter.java new file mode 100644 index 0000000000..ddf5f337b1 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBulkIndexWriter.java @@ -0,0 +1,378 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Random; +import java.util.UUID; +import org.apache.commons.io.FileUtils; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.Sort; +import org.apache.lucene.search.SortField; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +/** + * Functional coverage for {@link CagraHnswBulkIndexWriter}: the manual/direct-instance {@link + * CagraHnswBulkIndexWriter#addDocument} API (including its safety checks) and the one-shot {@link + * CagraHnswBulkIndexWriter#indexFbin}/{@link CagraHnswBulkIndexWriter#build(VectorSource, Config)} + * convenience entry points (single-segment, K-segment sequential, K-segment overlapped, + * {@link CagraHnswBulkIndexWriter.FieldCallback} metadata, rejection of {@code overlapped} for + * {@code build}). Guard-rail behavior of the underlying native-buffered writer itself is already + * covered by {@link TestNativeFlatBufferingGuardRails}; this class covers this class's own + * orchestration and its own safety checks layered on top. + */ +@SuppressSysoutChecks(bugUrl = "") +public class TestCagraHnswBulkIndexWriter extends LuceneTestCase { + + private static final String ID_FIELD = "id"; + private static final String VECTOR_FIELD = "vector_field"; + private static final String CATEGORY_FIELD = "category"; + + private Random random; + private Path indexDirPath; + private Path fbinPath; + + @Before + public void beforeTest() { + assumeTrue("cuVS not supported", isSupported()); + random = new Random(222); + indexDirPath = Paths.get(UUID.randomUUID().toString()); + fbinPath = Paths.get(UUID.randomUUID() + ".fbin"); + } + + @After + public void afterTest() throws IOException { + if (indexDirPath != null) { + File dir = indexDirPath.toFile(); + if (dir.exists() && dir.isDirectory()) { + FileUtils.deleteDirectory(dir); + } + } + if (fbinPath != null) { + new File(fbinPath.toString()).delete(); + } + } + + @Test + public void testSingleSegmentBuildIsSearchable() throws Exception { + int numDocs = 300; + int dimension = 32; + float[][] dataset = generateDataset(random, numDocs, dimension); + TestUtils.writeFbin(fbinPath, dataset); + + CagraHnswBulkIndexWriter.indexFbin(fbinPath, configFor(dimension, 1, false)); + + assertSearchable(numDocs, dataset, /* expectedSegments= */ 1); + } + + @Test + public void testPartitionedSequentialBuildProducesKSegments() throws Exception { + int numDocs = 400; + int dimension = 24; + int k = 4; + float[][] dataset = generateDataset(random, numDocs, dimension); + TestUtils.writeFbin(fbinPath, dataset); + + CagraHnswBulkIndexWriter.indexFbin(fbinPath, configFor(dimension, k, false)); + + assertSearchable(numDocs, dataset, /* expectedSegments= */ k); + } + + @Test + public void testOverlappedBuildProducesKSegments() throws Exception { + int numDocs = 400; + int dimension = 24; + int k = 4; + float[][] dataset = generateDataset(random, numDocs, dimension); + TestUtils.writeFbin(fbinPath, dataset); + + CagraHnswBulkIndexWriter.indexFbin(fbinPath, configFor(dimension, k, true)); + + assertSearchable(numDocs, dataset, /* expectedSegments= */ k); + } + + @Test + public void testGenericBuildViaVectorSource() throws Exception { + int numDocs = 200; + int dimension = 16; + float[][] dataset = generateDataset(random, numDocs, dimension); + + CagraHnswBulkIndexWriter.build( + new InMemoryVectorSource(dataset), configFor(dimension, 1, false)); + + assertSearchable(numDocs, dataset, /* expectedSegments= */ 1); + } + + @Test + public void testGenericBuildRejectsOverlapped() throws Exception { + int dimension = 8; + float[][] dataset = generateDataset(random, 50, dimension); + CagraHnswBulkIndexWriter.Config config = configFor(dimension, 2, true); + try { + CagraHnswBulkIndexWriter.build(new InMemoryVectorSource(dataset), config); + fail( + "expected IllegalArgumentException: overlapped is not supported by build(VectorSource," + + " Config)"); + } catch (IllegalArgumentException expected) { + // expected + } + } + + @Test + public void testIdFieldCanBeDisabled() throws Exception { + int numDocs = 50; + int dimension = 8; + float[][] dataset = generateDataset(random, numDocs, dimension); + TestUtils.writeFbin(fbinPath, dataset); + + CagraHnswBulkIndexWriter.Config config = + CagraHnswBulkIndexWriter.Config.builder() + .field(VECTOR_FIELD, dimension, EUCLIDEAN) + .idField(null) + .graphBuild(new AcceleratedHNSWParams.Builder().build()) + .segments(1, false) + .targetDirectory(indexDirPath) + .build(); + CagraHnswBulkIndexWriter.indexFbin(fbinPath, config); + + try (Directory dir = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(dir)) { + IndexSearcher searcher = new IndexSearcher(reader); + TopDocs results = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, dataset[0], 5), 5); + assertEquals(5, results.scoreDocs.length); + assertNull( + "id field should not be stored when idField(null) is used", + searcher.storedFields().document(results.scoreDocs[0].doc).get(ID_FIELD)); + } + } + + /** {@link CagraHnswBulkIndexWriter.FieldCallback} lets the one-shot API attach metadata per row. */ + @Test + public void testFieldCallbackAttachesMetadata() throws Exception { + int numDocs = 60; + int dimension = 8; + float[][] dataset = generateDataset(random, numDocs, dimension); + TestUtils.writeFbin(fbinPath, dataset); + + CagraHnswBulkIndexWriter.indexFbin( + fbinPath, + configFor(dimension, 1, false), + (doc, id) -> + doc.add( + new StringField(CATEGORY_FIELD, id % 2 == 0 ? "even" : "odd", Field.Store.YES))); + + try (Directory dir = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(dir)) { + IndexSearcher searcher = new IndexSearcher(reader); + TopDocs results = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, dataset[0], 1), 1); + String category = + searcher.storedFields().document(results.scoreDocs[0].doc).get(CATEGORY_FIELD); + assertEquals("even", category); // id=0 is even + } + } + + /** + * The manual/direct-instance API: caller builds the {@link Document} (including arbitrary extra + * fields) and drives {@link CagraHnswBulkIndexWriter#addDocument} directly, same shape as a + * plain {@link org.apache.lucene.index.IndexWriter}. + */ + @Test + public void testManualAddDocumentWithMetadata() throws Exception { + int numDocs = 40; + int dimension = 12; + float[][] dataset = generateDataset(random, numDocs, dimension); + + CagraHnswBulkIndexWriter.Config config = + CagraHnswBulkIndexWriter.Config.builder() + .field(VECTOR_FIELD, dimension, EUCLIDEAN) + .graphBuild(new AcceleratedHNSWParams.Builder().build()) + .build(); + + try (Directory dir = FSDirectory.open(indexDirPath); + CagraHnswBulkIndexWriter writer = + new CagraHnswBulkIndexWriter(dir, new IndexWriterConfig(), config, numDocs)) { + for (int i = 0; i < numDocs; i++) { + Document doc = new Document(); + doc.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); + doc.add(new StringField(CATEGORY_FIELD, i % 2 == 0 ? "even" : "odd", Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + writer.addDocument(doc); + } + } + + try (Directory dir = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(dir)) { + assertEquals(1, reader.leaves().size()); + IndexSearcher searcher = new IndexSearcher(reader); + TopDocs results = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, dataset[0], 1), 1); + assertEquals( + "even", searcher.storedFields().document(results.scoreDocs[0].doc).get(CATEGORY_FIELD)); + } + } + + @Test + public void testCloseWithTooFewDocumentsThrows() throws Exception { + int dimension = 8; + CagraHnswBulkIndexWriter.Config config = + CagraHnswBulkIndexWriter.Config.builder() + .field(VECTOR_FIELD, dimension, EUCLIDEAN) + .graphBuild(new AcceleratedHNSWParams.Builder().build()) + .build(); + + try (Directory dir = FSDirectory.open(indexDirPath)) { + CagraHnswBulkIndexWriter writer = + new CagraHnswBulkIndexWriter(dir, new IndexWriterConfig(), config, 10); + Document doc = new Document(); + doc.add( + new KnnFloatVectorField( + VECTOR_FIELD, generateDataset(random, 1, dimension)[0], EUCLIDEAN)); + writer.addDocument(doc); // only 1 of the promised 10 + + try { + writer.close(); + fail("expected IllegalStateException: fewer documents added than exactVectorCount"); + } catch (IllegalStateException expected) { + // expected + } + } + } + + @Test + public void testAddDocumentBeyondExactCountThrows() throws Exception { + int dimension = 8; + CagraHnswBulkIndexWriter.Config config = + CagraHnswBulkIndexWriter.Config.builder() + .field(VECTOR_FIELD, dimension, EUCLIDEAN) + .graphBuild(new AcceleratedHNSWParams.Builder().build()) + .build(); + + try (Directory dir = FSDirectory.open(indexDirPath)) { + CagraHnswBulkIndexWriter writer = + new CagraHnswBulkIndexWriter(dir, new IndexWriterConfig(), config, 1); + float[][] vectors = generateDataset(random, 2, dimension); + Document doc1 = new Document(); + doc1.add(new KnnFloatVectorField(VECTOR_FIELD, vectors[0], EUCLIDEAN)); + writer.addDocument(doc1); + + Document doc2 = new Document(); + doc2.add(new KnnFloatVectorField(VECTOR_FIELD, vectors[1], EUCLIDEAN)); + try { + writer.addDocument(doc2); // exactVectorCount was 1 + fail("expected IllegalStateException: more documents added than exactVectorCount"); + } catch (IllegalStateException expected) { + // expected -- the rejected call never reached the underlying writer or incremented the + // count, so it's still exactly 1 and close() below completes normally. + } finally { + writer.close(); + } + } + } + + @Test + public void testConstructorRejectsIndexSort() throws Exception { + int dimension = 8; + CagraHnswBulkIndexWriter.Config config = + CagraHnswBulkIndexWriter.Config.builder() + .field(VECTOR_FIELD, dimension, EUCLIDEAN) + .graphBuild(new AcceleratedHNSWParams.Builder().build()) + .build(); + IndexWriterConfig conf = + new IndexWriterConfig() + .setIndexSort(new Sort(new SortField(ID_FIELD, SortField.Type.STRING))); + + try (Directory dir = FSDirectory.open(indexDirPath)) { + try { + new CagraHnswBulkIndexWriter(dir, conf, config, 10); + fail("expected IllegalArgumentException: index-sorted segments are not supported"); + } catch (IllegalArgumentException expected) { + // expected + } + } + } + + private CagraHnswBulkIndexWriter.Config configFor( + int dimension, int numSegments, boolean overlapped) { + return CagraHnswBulkIndexWriter.Config.builder() + .field(VECTOR_FIELD, dimension, EUCLIDEAN) + .graphBuild(new AcceleratedHNSWParams.Builder().build()) + .segments(numSegments, overlapped) + .targetDirectory(indexDirPath) + .build(); + } + + private void assertSearchable(int numDocs, float[][] dataset, int expectedSegments) + throws Exception { + try (Directory dir = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(dir)) { + assertEquals("unexpected segment count", expectedSegments, reader.leaves().size()); + + IndexSearcher searcher = new IndexSearcher(reader); + int topK = 10; + TopDocs results = + searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, dataset[0], topK), topK); + assertEquals(topK, results.scoreDocs.length); + boolean sawQueryVectorItself = false; + for (var scoreDoc : results.scoreDocs) { + String id = searcher.storedFields().document(scoreDoc.doc).get(ID_FIELD); + int idValue = Integer.parseInt(id); + assertTrue("returned id out of range: " + id, idValue >= 0 && idValue < numDocs); + sawQueryVectorItself |= idValue == 0; + } + // Querying with dataset[0] itself (an exact, zero-distance match) should reliably surface + // id=0 within topK for a graph this small, across all segments the query fans out over. + assertTrue( + "expected id=0 (the query vector itself) within topK results", sawQueryVectorItself); + } + } + + private static final class InMemoryVectorSource implements VectorSource { + private final float[][] dataset; + + InMemoryVectorSource(float[][] dataset) { + this.dataset = dataset; + } + + @Override + public int dimensions() { + return dataset.length == 0 ? 0 : dataset[0].length; + } + + @Override + public int size() { + return dataset.length; + } + + @Override + public void get(int index, float[] dst) { + System.arraycopy(dataset[index], 0, dst, 0, dst.length); + } + + @Override + public void close() { + // nothing to release + } + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexParamsFactory.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexParamsFactory.java index f559ea3808..302629d6c1 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexParamsFactory.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexParamsFactory.java @@ -195,6 +195,43 @@ public void testHnswHeuristicDelegatesToCuVS() { assertEquals(7, cagraParams.getNumWriterThreads()); } + /** + * {@code cagraGraphBuildAlgo} is only consulted under {@code CUSTOM}; an explicit override left + * on the builder while the strategy is {@code HEURISTIC} must not change the delegated-to-cuVS + * result. Requires the native cuVS library. + */ + @Test + public void testHnswHeuristicIgnoresCagraGraphBuildAlgoOverride() { + assumeTrue("cuVS not supported", isSupported()); + + AcceleratedHNSWParams withoutOverride = + new AcceleratedHNSWParams.Builder() + .withStrategy(AcceleratedHNSWParams.Strategy.HEURISTIC) + .withMaxConn(16) + .withBeamWidth(100) + .build(); + AcceleratedHNSWParams withOverride = + new AcceleratedHNSWParams.Builder() + .withStrategy(AcceleratedHNSWParams.Strategy.HEURISTIC) + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .withMaxConn(16) + .withBeamWidth(100) + .build(); + + CagraIndexParams derivedWithoutOverride = + CagraIndexParamsFactory.create(withoutOverride, 10_000, 128); + CagraIndexParams derivedWithOverride = + CagraIndexParamsFactory.create(withOverride, 10_000, 128); + + assertEquals( + derivedWithoutOverride.getCagraGraphBuildAlgo(), + derivedWithOverride.getCagraGraphBuildAlgo()); + assertEquals(derivedWithoutOverride.getGraphDegree(), derivedWithOverride.getGraphDegree()); + assertEquals( + derivedWithoutOverride.getIntermediateGraphDegree(), + derivedWithOverride.getIntermediateGraphDegree()); + } + /** * The heuristic type must reach cuVS rather than being pinned to the default: under * SIMILAR_SEARCH_PERFORMANCE cuVS derives {@code graph_degree = 2 + maxConn * 2 / 3} instead of @@ -254,10 +291,12 @@ public void testStrategySpecificParamsRemainAccepted() { .withStrategy(AcceleratedHNSWParams.Strategy.HEURISTIC) .withGraphDegree(96) .withIntermediateGraphDegree(192) + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) .build(); // Retained verbatim on the instance; CagraIndexParamsFactory is what declines to apply them. assertEquals(96, hnswParams.getGraphdegree()); assertEquals(192, hnswParams.getIntermediateGraphDegree()); + assertEquals(CagraGraphBuildAlgo.NN_DESCENT, hnswParams.getCagraGraphBuildAlgo()); GPUSearchParams gpuParams = new GPUSearchParams.Builder() diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestFbinVectorSource.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestFbinVectorSource.java new file mode 100644 index 0000000000..41e7d4a081 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestFbinVectorSource.java @@ -0,0 +1,120 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.TestUtils.writeFbin; + +import java.io.File; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Random; +import java.util.UUID; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +/** + * Correctness coverage for {@link FbinVectorSource}: header parsing, whole-file vs. windowed + * (sliced) reads, and the forward-only/single-consumer contract it shares with {@link + * VectorSource}. Does not require GPU support -- this is pure file I/O. + */ +public class TestFbinVectorSource extends LuceneTestCase { + + private Path fbinPath; + + @Before + public void beforeTest() { + fbinPath = Paths.get(UUID.randomUUID() + ".fbin"); + } + + @After + public void afterTest() { + if (fbinPath != null) { + new File(fbinPath.toString()).delete(); + } + } + + @Test + public void testWholeFileReadMatchesDataset() throws Exception { + Random random = new Random(1); + int numVectors = 250; + int dimension = 17; + float[][] dataset = generateDataset(random, numVectors, dimension); + writeFbin(fbinPath, dataset); + + // A small chunk size forces multiple prefetch chunks so this also exercises the chunk + // boundary/advance() path, not just a single-chunk read. + try (FbinVectorSource source = new FbinVectorSource(fbinPath, /* chunkSizeMB= */ 1)) { + assertEquals(dimension, source.dimensions()); + assertEquals(numVectors, source.size()); + + float[] scratch = new float[dimension]; + for (int i = 0; i < numVectors; i++) { + source.get(i, scratch); + assertArrayEquals("vector " + i + " mismatch", dataset[i], scratch, 0f); + } + } + } + + @Test + public void testWindowedReadServesOnlyItsSlice() throws Exception { + Random random = new Random(2); + int numVectors = 100; + int dimension = 8; + int sliceStart = 30; + int sliceSize = 25; + float[][] dataset = generateDataset(random, numVectors, dimension); + writeFbin(fbinPath, dataset); + + try (FbinVectorSource source = new FbinVectorSource(fbinPath, sliceStart, sliceSize, 1)) { + assertEquals(dimension, source.dimensions()); + assertEquals(sliceSize, source.size()); + + for (int i = 0; i < sliceSize; i++) { + float[] got = source.get(i); + assertArrayEquals( + "relative index " + i + " should map to absolute " + (sliceStart + i), + dataset[sliceStart + i], + got, + 0f); + } + } + } + + @Test + public void testOutOfOrderAccessIsRejected() throws Exception { + Random random = new Random(3); + float[][] dataset = generateDataset(random, 20, 4); + writeFbin(fbinPath, dataset); + + try (FbinVectorSource source = new FbinVectorSource(fbinPath, 1)) { + source.get(5); + try { + source.get(2); // going backwards must be rejected: forward-only contract + fail("expected UnsupportedOperationException for out-of-order access"); + } catch (UnsupportedOperationException expected) { + // expected + } + } + } + + @Test + public void testOutOfBoundsIndexIsRejected() throws Exception { + Random random = new Random(4); + float[][] dataset = generateDataset(random, 10, 4); + writeFbin(fbinPath, dataset); + + try (FbinVectorSource source = new FbinVectorSource(fbinPath, 1)) { + try { + source.get(10); // window is [0, 10) + fail("expected IndexOutOfBoundsException"); + } catch (IndexOutOfBoundsException expected) { + // expected + } + } + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMergedGraphOrdinalBounds.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMergedGraphOrdinalBounds.java new file mode 100644 index 0000000000..12b00c6e12 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMergedGraphOrdinalBounds.java @@ -0,0 +1,198 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; + +import java.util.Random; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.hnsw.HnswGraphProvider; +import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.CodecReader; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.NoMergePolicy; +import org.apache.lucene.index.Term; +import org.apache.lucene.index.TieredMergePolicy; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.tests.util.TestUtil; +import org.apache.lucene.util.hnsw.HnswGraph; +import org.junit.Test; + +/** + * Repro for the CI-observed {@code EOFException} in {@code OffHeapFloatVectorValues} during + * concurrent KNN search over an accelerated-HNSW index built via {@link + * TestAcceleratedHNSWDeletedDocuments}/{@link TestCuVSAcceleratedHNSWDeletedDocuments} (both: + * deletions + a real merge, heap-buffered path, no {@code numInputVectors}). + * + *

Rather than relying on a random concurrent search happening to traverse a bad graph node + * (which only reproduced intermittently, on one CI node), this walks the entire merged + * HNSW graph directly and asserts every neighbor ordinal is within the merged segment's actual + * flat-vector count. This targets the suspected root cause: {@code + * Lucene99AcceleratedHNSWVectorsWriter#mergeOneField} derives the merged vector set twice, + * independently -- once via the real {@code flatVectorsWriter.mergeOneField} (the authoritative + * flat {@code .vec} file) and again via {@code vectorBasedMerge}'s own call to {@code + * KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues} to build the CAGRA/HNSW graph. If + * those two independently-derived views of "the merged, post-deletion vector set" ever disagree in + * count or ordinal order, the graph ends up referencing ordinals the flat file doesn't actually + * have, which is exactly what an out-of-bounds read (EOFException) during traversal would look + * like. + * + *

This test is deterministic: it fails on any disagreement between the graph and the flat file, + * rather than depending on a search happening to reach the bad node. + */ +@SuppressSysoutChecks(bugUrl = "") +public class TestMergedGraphOrdinalBounds extends LuceneTestCase { + + private static final String ID_FIELD = "id"; + private static final String FIELD = "vector"; + + @Test + public void testMergedGraphOrdinalsStayWithinFlatVectorBounds() throws Exception { + assumeTrue("cuVS not supported", isSupported()); + + Random random = new Random(1234); + int segmentSize = 300; + int dimension = 32; + // Interspersed deletions on both segments, so the merge must drop a scattered subset of + // ordinals from each -- not just a contiguous prefix/suffix -- when it re-derives the merged + // vector set. + int deleteEveryNth = 4; + + Codec codec = TestUtil.alwaysKnnVectorsFormat(new Lucene99AcceleratedHNSWVectorsFormat()); + IndexWriterConfig config = + new IndexWriterConfig().setCodec(codec).setMergePolicy(NoMergePolicy.INSTANCE); + + int expectedLiveVectors; + try (Directory dir = newDirectory(); + IndexWriter writer = new IndexWriter(dir, config)) { + int deletedFromSegment1 = + addSegmentWithInterspersedDeletions( + writer, 0, segmentSize, dimension, deleteEveryNth, random); + writer.commit(); // segment 1, alone + int deletedFromSegment2 = + addSegmentWithInterspersedDeletions( + writer, segmentSize, segmentSize, dimension, deleteEveryNth, random); + writer.commit(); // segment 2, alone + + expectedLiveVectors = 2 * segmentSize - deletedFromSegment1 - deletedFromSegment2; + + // NoMergePolicy blocks forced merges too, so swap it out now that the two segments (each + // with their own interspersed deletions already committed) are set up. + writer.getConfig().setMergePolicy(new TieredMergePolicy()); + writer.forceMerge(1); + writer.commit(); + + try (DirectoryReader reader = DirectoryReader.open(dir)) { + assertEquals( + "expected the forced merge to produce a single segment", 1, reader.leaves().size()); + LeafReader leaf = reader.leaves().get(0).reader(); + + FloatVectorValues flatValues = leaf.getFloatVectorValues(FIELD); + assertEquals( + "merged flat vector count should equal (added - deleted)", + expectedLiveVectors, + flatValues.size()); + + HnswGraph graph = graphOf(leaf); + int level0NodeCount = graph.getNodesOnLevel(0).size(); + assertEquals( + "HNSW graph's level-0 node count disagrees with the merged flat vector file's actual" + + " count -- the graph and the flat file were derived independently by" + + " vectorBasedMerge and flatVectorsWriter.mergeOneField and disagree", + flatValues.size(), + level0NodeCount); + + assertAllNeighborOrdinalsInBounds(graph, flatValues.size()); + } + } + } + + /** + * Every neighbor referenced anywhere in the graph, at every level, must be a valid ordinal into + * the merged segment's actual flat vector data -- otherwise a reader resolving that neighbor's + * vector (e.g. mid-search, to score it) reads past the end of the flat file. + */ + private static void assertAllNeighborOrdinalsInBounds(HnswGraph graph, int liveVectorCount) + throws Exception { + for (int level = 0; level < graph.numLevels(); level++) { + HnswGraph.NodesIterator nodes = graph.getNodesOnLevel(level); + while (nodes.hasNext()) { + int node = nodes.nextInt(); + assertTrue( + "node " + + node + + " at level " + + level + + " is itself out of bounds (live vectors: " + + liveVectorCount + + ")", + node >= 0 && node < liveVectorCount); + graph.seek(level, node); + for (int neighbor = graph.nextNeighbor(); + neighbor != NO_MORE_DOCS; + neighbor = graph.nextNeighbor()) { + assertTrue( + "node " + + node + + " at level " + + level + + " has a neighbor ordinal " + + neighbor + + " out of bounds for the merged segment's " + + liveVectorCount + + " live vectors", + neighbor >= 0 && neighbor < liveVectorCount); + } + } + } + } + + private static HnswGraph graphOf(LeafReader leaf) throws Exception { + KnnVectorsReader knnReader = ((CodecReader) leaf).getVectorReader(); + if (knnReader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) { + knnReader = fieldsReader.getFieldReader(FIELD); + } + return ((HnswGraphProvider) knnReader).getGraph(FIELD); + } + + /** + * Adds {@code count} documents (global ids {@code [startId, startId + count)}), then deletes + * every {@code deleteEveryNth}-th one by id, scattering the deletions across the segment rather + * than leaving a contiguous surviving range. + * + * @return the number of documents deleted from this segment + */ + private static int addSegmentWithInterspersedDeletions( + IndexWriter writer, int startId, int count, int dimension, int deleteEveryNth, Random random) + throws Exception { + float[][] dataset = generateDataset(random, count, dimension); + for (int i = 0; i < count; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(startId + i), Field.Store.YES)); + document.add(new KnnFloatVectorField(FIELD, dataset[i], EUCLIDEAN)); + writer.addDocument(document); + } + int deleted = 0; + for (int i = 0; i < count; i += deleteEveryNth) { + writer.deleteDocuments(new Term(ID_FIELD, Integer.toString(startId + i))); + deleted++; + } + return deleted; + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingGuardRails.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingGuardRails.java new file mode 100644 index 0000000000..2645128419 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingGuardRails.java @@ -0,0 +1,221 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import java.io.File; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Random; +import java.util.UUID; +import org.apache.commons.io.FileUtils; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.NumericDocValuesField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.NoMergePolicy; +import org.apache.lucene.index.SerialMergeScheduler; +import org.apache.lucene.index.TieredMergePolicy; +import org.apache.lucene.search.Sort; +import org.apache.lucene.search.SortField; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +/** + * Negative-path coverage for the three guard rails in {@code Lucene99AcceleratedHNSWVectorsWriter} + * around {@code AcceleratedHNSWParams.numInputVectors} (native flat buffering): a count mismatch, + * an index-sorted segment, and a merge attempt. None of these had test coverage before. + * + *

Positive-path coverage (does a natively-buffered index actually search correctly, tolerate + * deletions, etc.) lives separately in {@link TestNativeFlatBufferingIndexAndSearch}. + */ +@SuppressSysoutChecks(bugUrl = "") +public class TestNativeFlatBufferingGuardRails extends LuceneTestCase { + + private static final String ID_FIELD = "id"; + private static final String VECTOR_FIELD = "vector_field"; + + private Random random; + private Path indexDirPath; + + @Before + public void beforeTest() throws Exception { + assumeTrue("cuVS not supported", isSupported()); + random = new Random(222); + indexDirPath = Paths.get(UUID.randomUUID().toString()); + } + + @After + public void afterTest() throws Exception { + if (indexDirPath == null) { + return; + } + File indexDirPathFile = indexDirPath.toFile(); + if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { + FileUtils.deleteDirectory(indexDirPathFile); + } + } + + /** + * The count-mismatch guard exists for a caller-side bookkeeping error: declaring {@code + * numInputVectors} against the pre-filter document count instead of the number of vectors that + * actually reach {@code addValue} (e.g. an ingest-time filter skips some documents' vector + * field). It is unrelated to, and not triggered by, Lucene-level deletion -- see the javadoc on + * {@link TestNativeFlatBufferingIndexAndSearch#testDeletedDocsAfterNativeFlatBufferedFlush}. + */ + @Test + public void testCountMismatchFromIngestTimeFilterIsRejected() throws Exception { + int declaredNumInputVectors = 100; + int actuallyIndexed = declaredNumInputVectors - 1; // one doc "filtered out" before addValue + int dimension = 32; + float[][] dataset = generateDataset(random, actuallyIndexed, dimension); + + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder().withNumInputVectors(declaredNumInputVectors).build(); + Codec codec = new Lucene101AcceleratedHNSWCodec(params); + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(codec) + .setUseCompoundFile(false) + .setMaxBufferedDocs(declaredNumInputVectors + 1) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) + .setMergePolicy(NoMergePolicy.INSTANCE); + + try (Directory dir = FSDirectory.open(indexDirPath); + IndexWriter writer = new IndexWriter(dir, config)) { + for (int i = 0; i < actuallyIndexed; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + writer.addDocument(document); + } + IllegalStateException thrown = expectThrows(IllegalStateException.class, writer::commit); + assertTrue( + "unexpected message: " + thrown.getMessage(), + thrown.getMessage().contains("numInputVectors")); + } + } + + /** Native flat buffering pre-sizes a single flush's buffer and cannot support sorted flushes. */ + @Test + public void testIndexSortedSegmentIsRejected() throws Exception { + int numDocs = 50; + int dimension = 32; + float[][] dataset = generateDataset(random, numDocs, dimension); + + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder().withNumInputVectors(numDocs).build(); + Codec codec = new Lucene101AcceleratedHNSWCodec(params); + Sort indexSort = new Sort(new SortField("sort_key", SortField.Type.LONG)); + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(codec) + .setUseCompoundFile(false) + .setIndexSort(indexSort) + .setMaxBufferedDocs(numDocs + 1) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) + .setMergePolicy(NoMergePolicy.INSTANCE); + + try (Directory dir = FSDirectory.open(indexDirPath); + IndexWriter writer = new IndexWriter(dir, config)) { + // KnnVectorsFormat#fieldsWriter (and so the writer's index-sort check in its constructor) + // is invoked on the first addDocument() for the segment, not at commit() -- so the guard + // must be expected around the whole indexing loop, not just the flush. + IllegalArgumentException thrown = + expectThrows( + IllegalArgumentException.class, + () -> { + for (int i = 0; i < numDocs; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); + document.add(new NumericDocValuesField("sort_key", numDocs - i)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + writer.addDocument(document); + } + }); + assertTrue( + "unexpected message: " + thrown.getMessage(), + thrown.getMessage().contains("index-sorted")); + } + } + + /** + * Native flat buffering supports only the unsorted single-segment flush path; merging two + * natively-buffered segments must be rejected rather than silently mis-sizing the native buffer. + */ + @Test + public void testMergeIsRejected() throws Exception { + int segmentSize = 40; + int dimension = 32; + + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder().withNumInputVectors(segmentSize).build(); + Codec codec = new Lucene101AcceleratedHNSWCodec(params); + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(codec) + .setUseCompoundFile(false) + .setMaxBufferedDocs(segmentSize + 1) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) + // Keep the two flushes below as separate segments; NoMergePolicy also blocks forced + // merges, so it is swapped out before forceMerge() is called. + .setMergePolicy(NoMergePolicy.INSTANCE) + // Force merges to run synchronously on the calling thread, so the guard's exception + // (or whatever IndexWriter/SegmentMerger wraps it as) surfaces directly from + // forceMerge() instead of on a background merge thread. + .setMergeScheduler(new SerialMergeScheduler()); + + try (Directory dir = FSDirectory.open(indexDirPath); + IndexWriter writer = new IndexWriter(dir, config)) { + addSegment(writer, 0, segmentSize, dimension); + writer.commit(); // segment 1: exactly segmentSize vectors, matching numInputVectors + addSegment(writer, segmentSize, segmentSize, dimension); + writer.commit(); // segment 2: exactly segmentSize vectors, matching numInputVectors + + writer.getConfig().setMergePolicy(new TieredMergePolicy()); + Throwable thrown = expectThrows(Throwable.class, () -> writer.forceMerge(1)); + assertTrue( + "expected UnsupportedOperationException somewhere in the cause chain of: " + thrown, + causedBy(thrown, UnsupportedOperationException.class)); + } + } + + private void addSegment(IndexWriter writer, int startId, int count, int dimension) + throws Exception { + float[][] dataset = generateDataset(random, count, dimension); + for (int i = 0; i < count; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(startId + i), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + writer.addDocument(document); + } + } + + private static boolean causedBy(Throwable t, Class type) { + for (Throwable cur = t; cur != null; cur = cur.getCause()) { + if (type.isInstance(cur)) { + return true; + } + for (Throwable suppressed : cur.getSuppressed()) { + if (causedBy(suppressed, type)) { + return true; + } + } + } + return false; + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingIndexAndSearch.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingIndexAndSearch.java new file mode 100644 index 0000000000..147b7aca24 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingIndexAndSearch.java @@ -0,0 +1,225 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; +import java.io.File; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Random; +import java.util.UUID; +import org.apache.commons.io.FileUtils; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.NoMergePolicy; +import org.apache.lucene.index.Term; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +/** + * Positive-path functional coverage for native flat buffering ({@code + * AcceleratedHNSWParams.numInputVectors}) beyond {@code TestNativeFlatVectorsWriterRoundTrip}, + * which only checks that the flat {@code .vec} file round-trips -- not that the resulting index is + * actually searchable, tolerates deletions, or composes correctly with the odd-graph-degree fix. + * + *

The negative/guard-rail paths (count mismatch, index-sorted segments, merges) are covered + * separately in {@link TestNativeFlatBufferingGuardRails}. + */ +@SuppressSysoutChecks(bugUrl = "") +public class TestNativeFlatBufferingIndexAndSearch extends LuceneTestCase { + + private static final String ID_FIELD = "id"; + private static final String VECTOR_FIELD = "vector_field"; + + private Random random; + private Path indexDirPath; + + @Before + public void beforeTest() throws Exception { + assumeTrue("cuVS not supported", isSupported()); + random = new Random(222); + indexDirPath = Paths.get(UUID.randomUUID().toString()); + } + + @After + public void afterTest() throws Exception { + if (indexDirPath == null) { + return; + } + File indexDirPathFile = indexDirPath.toFile(); + if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { + FileUtils.deleteDirectory(indexDirPathFile); + } + } + + /** A natively-buffered index must still be searchable through the normal Lucene KNN query API. */ + @Test + public void testIndexAndSearch() throws Exception { + int numDocs = 500; + int dimension = 32; + int topK = 10; + float[][] dataset = generateDataset(random, numDocs, dimension); + + buildNativeFlatBufferedIndex(numDocs, dimension, dataset); + + try (Directory dir = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(dir)) { + assertEquals("expected a single native-flat-buffered segment", 1, reader.leaves().size()); + + IndexSearcher searcher = new IndexSearcher(reader); + float[] queryVector = generateDataset(random, 1, dimension)[0]; + TopDocs results = + searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK), topK); + + assertEquals("expected topK results", topK, results.scoreDocs.length); + for (var scoreDoc : results.scoreDocs) { + String id = searcher.storedFields().document(scoreDoc.doc).get(ID_FIELD); + int idValue = Integer.parseInt(id); + assertTrue("returned id out of range: " + id, idValue >= 0 && idValue < numDocs); + } + } + } + + /** + * Deletion is orthogonal to native flat buffering: {@code IndexWriter.deleteDocuments} only + * updates Lucene's liveDocs bitset at search time -- it never touches {@code FieldWriter} or the + * native host matrix, so it cannot trip (and isn't meant to be caught by) the count-mismatch + * guard rail tested in {@link TestNativeFlatBufferingGuardRails}. This test instead confirms that + * deletions applied after a natively-buffered flush are still honored correctly at search time. + */ + @Test + public void testDeletedDocsAfterNativeFlatBufferedFlush() throws Exception { + int numDocs = 300; + int dimension = 32; + float[][] dataset = generateDataset(random, numDocs, dimension); + + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder().withNumInputVectors(numDocs).build(); + Codec codec = new Lucene101AcceleratedHNSWCodec(params); + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(codec) + .setUseCompoundFile(false) + .setMaxBufferedDocs(numDocs + 1) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) + .setMergePolicy(NoMergePolicy.INSTANCE); + + try (Directory dir = FSDirectory.open(indexDirPath); + IndexWriter writer = new IndexWriter(dir, config)) { + for (int i = 0; i < numDocs; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + writer.addDocument(document); + } + writer.commit(); // single natively-buffered flush: FieldWriter's count matches numDocs + + // Delete every 3rd doc. No new vectors are added, so this does not trigger another flush of + // the vector field and cannot interact with the numInputVectors hint. + for (int i = 0; i < numDocs; i += 3) { + writer.deleteDocuments(new Term(ID_FIELD, Integer.toString(i))); + } + writer.commit(); + } + + try (Directory dir = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(dir)) { + assertEquals( + "expected the deletions to land in the same single segment", 1, reader.leaves().size()); + assertTrue("expected some deleted docs", reader.numDeletedDocs() > 0); + + IndexSearcher searcher = new IndexSearcher(reader); + TopDocs results = + searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, dataset[0], numDocs), numDocs); + for (var scoreDoc : results.scoreDocs) { + String id = searcher.storedFields().document(scoreDoc.doc).get(ID_FIELD); + assertNotEquals( + "deleted doc id=" + id + " was still returned by search", 0, Integer.parseInt(id) % 3); + } + } + } + + /** + * The M = ceil(cagraGraphDegree / 2) fix ({@link TestAcceleratedHNSWOddGraphDegree}) must also + * hold on the native-flat-buffered write path ({@code writeFieldNative}), which is a distinct + * call path from the heap-buffered one that test exercises. + */ + @Test + public void testOddGraphDegreeWithNativeFlatBuffering() throws Exception { + int numDocs = 200; + int dimension = 32; + int oddGraphDegree = 63; + float[][] dataset = generateDataset(random, numDocs, dimension); + + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder() + .withStrategy(AcceleratedHNSWParams.Strategy.CUSTOM) + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .withIntermediateGraphDegree(128) + .withGraphDegree(oddGraphDegree) + .withNumInputVectors(numDocs) + .build(); + + buildNativeFlatBufferedIndex(numDocs, dimension, dataset, params); + + try (Directory dir = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(dir)) { + IndexSearcher searcher = new IndexSearcher(reader); + TopDocs results = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, dataset[0], 5), 5); + assertEquals(5, results.scoreDocs.length); + } + } + + private void buildNativeFlatBufferedIndex(int numDocs, int dimension, float[][] dataset) + throws Exception { + buildNativeFlatBufferedIndex( + numDocs, + dimension, + dataset, + new AcceleratedHNSWParams.Builder().withNumInputVectors(numDocs).build()); + } + + private void buildNativeFlatBufferedIndex( + int numDocs, int dimension, float[][] dataset, AcceleratedHNSWParams params) + throws Exception { + Codec codec = new Lucene101AcceleratedHNSWCodec(params); + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(codec) + .setUseCompoundFile(false) + .setMaxBufferedDocs(numDocs + 1) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) + .setMergePolicy(NoMergePolicy.INSTANCE); + + try (Directory dir = FSDirectory.open(indexDirPath); + IndexWriter writer = new IndexWriter(dir, config)) { + for (int i = 0; i < numDocs; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + writer.addDocument(document); + } + writer.commit(); + } + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatVectorsWriterRoundTrip.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatVectorsWriterRoundTrip.java new file mode 100644 index 0000000000..c503f74b20 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatVectorsWriterRoundTrip.java @@ -0,0 +1,282 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.Set; +import java.util.UUID; +import java.util.function.IntFunction; +import org.apache.commons.io.FileUtils; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.NoMergePolicy; +import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +/** + * Positive round-trip check for {@link NativeFlatVectorsWriter}: builds a single-segment index + * with {@code AcceleratedHNSWParams.numInputVectors} set (native flat buffering) and reads the + * {@code .vec}/{@code .vemf} files back through the stock {@code Lucene99FlatVectorsReader}, + * asserting every vector round-trips byte-exact. + * + *

This is the check called for by {@link NativeFlatVectorsWriter}'s "on a Lucene upgrade" class + * javadoc: it confirms the hand-transcribed format is still readable by Lucene's real reader, and + * should pass on every {@code lucene-core} version bump. + */ +@SuppressSysoutChecks(bugUrl = "") +public class TestNativeFlatVectorsWriterRoundTrip extends LuceneTestCase { + + private static final String ID_FIELD = "id"; + private static final String VECTOR_FIELD = "vector_field"; + + private Random random; + private Path indexDirPath; + + @Before + public void beforeTest() throws Exception { + assumeTrue("cuVS not supported", isSupported()); + random = new Random(222); + indexDirPath = Paths.get(UUID.randomUUID().toString()); + } + + @After + public void afterTest() throws Exception { + if (indexDirPath == null) { + return; + } + File indexDirPathFile = indexDirPath.toFile(); + if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { + FileUtils.deleteDirectory(indexDirPathFile); + } + } + + @Test + public void vectorsRoundTripThroughStockLucene99FlatVectorsReader() throws Exception { + int numDocs = 500; + int dimension = 32; + float[][] dataset = generateDataset(random, numDocs, dimension); + + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder().withNumInputVectors(numDocs).build(); + Codec codec = new Lucene101AcceleratedHNSWCodec(params); + + // Force everything into a single unsorted, unmerged flush: native flat buffering requires + // numInputVectors to equal the exact number of vectors landing in that one flush. + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(codec) + .setUseCompoundFile(false) + .setMaxBufferedDocs(numDocs + 1) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) + .setMergePolicy(NoMergePolicy.INSTANCE); + + try (Directory dir = FSDirectory.open(indexDirPath); + IndexWriter writer = new IndexWriter(dir, config)) { + for (int i = 0; i < numDocs; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + writer.addDocument(document); + } + writer.commit(); + } + + try (Directory dir = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(dir)) { + assertEquals("expected a single native-flat-buffered segment", 1, reader.leaves().size()); + LeafReader leafReader = reader.leaves().get(0).reader(); + assertFieldRoundTrips(leafReader, VECTOR_FIELD, dimension, numDocs, docId -> dataset[docId]); + } + } + + /** + * Cardinality (8000) exceeds {@code IndexedDISI.MAX_ARRAY_LENGTH} (4095), so {@code + * OrdToDocDISIReaderConfiguration} picks the DENSE (bitset) {@code docsWithField} encoding for + * this field. + */ + @Test + public void vectorsRoundTripWithDenseDocsWithFieldEncoding() throws Exception { + roundTripPartialField(10000, 8000, 32); + } + + /** + * Cardinality (800) stays under {@code IndexedDISI.MAX_ARRAY_LENGTH} (4095), so {@code + * OrdToDocDISIReaderConfiguration} picks the SPARSE (array) {@code docsWithField} encoding + * instead. + */ + @Test + public void vectorsRoundTripWithSparseDocsWithFieldEncoding() throws Exception { + roundTripPartialField(1000, 800, 32); + } + + /** + * Two vector fields written to the same segment, exercising the boundary between consecutive + * per-field records in {@code .vemf}: the second field's header must be found where the first + * field's record actually ends. + */ + @Test + public void vectorsRoundTripAcrossMultipleFieldsInSameSegment() throws Exception { + int numDocs = 500; + int dimensionA = 32; + int dimensionB = 16; + String fieldB = "vector_field_two"; + float[][] datasetA = generateDataset(random, numDocs, dimensionA); + float[][] datasetB = generateDataset(random, numDocs, dimensionB); + + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder().withNumInputVectors(numDocs).build(); + Codec codec = new Lucene101AcceleratedHNSWCodec(params); + + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(codec) + .setUseCompoundFile(false) + .setMaxBufferedDocs(numDocs + 1) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) + .setMergePolicy(NoMergePolicy.INSTANCE); + + try (Directory dir = FSDirectory.open(indexDirPath); + IndexWriter writer = new IndexWriter(dir, config)) { + for (int i = 0; i < numDocs; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, datasetA[i], EUCLIDEAN)); + document.add(new KnnFloatVectorField(fieldB, datasetB[i], EUCLIDEAN)); + writer.addDocument(document); + } + writer.commit(); + } + + try (Directory dir = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(dir)) { + assertEquals("expected a single native-flat-buffered segment", 1, reader.leaves().size()); + LeafReader leafReader = reader.leaves().get(0).reader(); + assertFieldRoundTrips( + leafReader, VECTOR_FIELD, dimensionA, numDocs, docId -> datasetA[docId]); + assertFieldRoundTrips(leafReader, fieldB, dimensionB, numDocs, docId -> datasetB[docId]); + } + } + + /** + * Builds a segment of {@code numDocs} documents where only a random {@code numDocsWithVector} + * of them carry {@code VECTOR_FIELD}, then verifies the round trip. + */ + private void roundTripPartialField(int numDocs, int numDocsWithVector, int dimension) + throws Exception { + Set docsWithVector = randomDocSubset(numDocs, numDocsWithVector); + float[][] dataset = generateDataset(random, numDocsWithVector, dimension); + + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder().withNumInputVectors(numDocsWithVector).build(); + Codec codec = new Lucene101AcceleratedHNSWCodec(params); + + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(codec) + .setUseCompoundFile(false) + .setMaxBufferedDocs(numDocs + 1) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) + .setMergePolicy(NoMergePolicy.INSTANCE); + + Map docIdToVectorIndex = new HashMap<>(); + try (Directory dir = FSDirectory.open(indexDirPath); + IndexWriter writer = new IndexWriter(dir, config)) { + int vectorIndex = 0; + for (int i = 0; i < numDocs; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); + if (docsWithVector.contains(i)) { + document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[vectorIndex], EUCLIDEAN)); + docIdToVectorIndex.put(i, vectorIndex); + vectorIndex++; + } + writer.addDocument(document); + } + writer.commit(); + } + + try (Directory dir = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(dir)) { + assertEquals("expected a single native-flat-buffered segment", 1, reader.leaves().size()); + LeafReader leafReader = reader.leaves().get(0).reader(); + assertFieldRoundTrips( + leafReader, + VECTOR_FIELD, + dimension, + numDocsWithVector, + docId -> dataset[docIdToVectorIndex.get(docId)]); + } + } + + /** Picks {@code count} distinct doc ids out of {@code [0, numDocs)}. */ + private Set randomDocSubset(int numDocs, int count) { + List allDocs = new ArrayList<>(numDocs); + for (int i = 0; i < numDocs; i++) { + allDocs.add(i); + } + Collections.shuffle(allDocs, random); + return new HashSet<>(allDocs.subList(0, count)); + } + + private void assertFieldRoundTrips( + LeafReader leafReader, + String fieldName, + int dimension, + int expectedCount, + IntFunction expectedVectorForDocId) + throws IOException { + FloatVectorValues values = leafReader.getFloatVectorValues(fieldName); + assertNotNull(values); + assertEquals(expectedCount, values.size()); + assertEquals(dimension, values.dimension()); + + int seen = 0; + KnnVectorValues.DocIndexIterator it = values.iterator(); + for (int doc = it.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = it.nextDoc()) { + String id = leafReader.storedFields().document(doc).get(ID_FIELD); + float[] roundTripped = values.vectorValue(it.index()); + assertArrayEquals( + "vector for field=" + + fieldName + + " id=" + + id + + " did not round-trip byte-exact through the stock Lucene99FlatVectorsReader", + expectedVectorForDocId.apply(Integer.parseInt(id)), + roundTripped, + 0f); + seen++; + } + assertEquals("did not visit every vector for field=" + fieldName, expectedCount, seen); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java index 88ca3150fd..e19e5ce64f 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java @@ -4,10 +4,43 @@ */ package com.nvidia.cuvs.lucene; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.channels.FileChannel; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; import java.util.Random; public class TestUtils { + /** Writes {@code dataset} as an uncompressed {@code .fbin} file (little-endian float32 rows). */ + public static void writeFbin(Path path, float[][] dataset) throws IOException { + int numVectors = dataset.length; + int dimension = numVectors == 0 ? 0 : dataset[0].length; + ByteBuffer buf = + ByteBuffer.allocate(8 + numVectors * dimension * Float.BYTES) + .order(ByteOrder.LITTLE_ENDIAN); + buf.putInt(numVectors); + buf.putInt(dimension); + for (float[] vector : dataset) { + for (float v : vector) { + buf.putFloat(v); + } + } + buf.flip(); + try (FileChannel ch = + FileChannel.open( + path, + StandardOpenOption.CREATE, + StandardOpenOption.WRITE, + StandardOpenOption.TRUNCATE_EXISTING)) { + while (buf.hasRemaining()) { + ch.write(buf); + } + } + } + public static float[][] generateDataset(Random random, int size, int dimensions) { float[][] dataset = new float[size][dimensions]; for (int i = 0; i < size; i++) { diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphEquivalence.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphEquivalence.java new file mode 100644 index 0000000000..fd4f9079f4 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphEquivalence.java @@ -0,0 +1,154 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; + +import com.nvidia.cuvs.CuVSMatrix; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Random; +import org.apache.lucene.store.ByteBuffersDirectory; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.util.hnsw.HnswGraph; +import org.apache.lucene.util.hnsw.HnswGraph.NodesIterator; +import org.junit.Test; + +/** + * Verifies that {@code writerThreads > 1} produces the same result as the serial path for the + * two parallelizations gated on {@code AcceleratedHNSWUtils}/{@code GPUBuiltHnswGraph}'s {@code + * PARALLEL_MIN_NODES} threshold: materializing the CAGRA adjacency into {@code NeighborArray}s + * (the {@code GPUBuiltHnswGraph} constructor), and encoding level 0 to disk ({@code + * AcceleratedHNSWUtils#writeGraph}). + * + *

Both tests use a synthetic adjacency ({@link CuVSMatrix#ofArray(int[][])}, the same host-matrix + * construction the higher-layer subset builder already uses) rather than a real CAGRA build, so + * that the comparison isolates these two parallelizations from CAGRA's own build-to-build + * variance -- a real GPU build is not guaranteed to produce the identical graph twice even with the + * same input and thread count, which would make an end-to-end build comparison unreliable for this + * purpose. This needs cuVS/GPU only to allocate the host matrix, not to run a build. + */ +@SuppressSysoutChecks(bugUrl = "") +public class TestWriterThreadsGraphEquivalence extends LuceneTestCase { + + // Must be >= PARALLEL_MIN_NODES (1 << 16) in both AcceleratedHNSWUtils and + // GPUBuiltHnswGraph, or the "parallel" runs below silently fall through to the serial branch and + // the test would pass without exercising anything. + private static final int NUM_NODES = (1 << 16) + 1000; + private static final int DEGREE = 12; + private static final int NUM_THREADS = 4; + + @Test + public void fillNeighborArrayParallelMatchesSerial() throws Exception { + assumeTrue("cuVS not supported", isSupported()); + int[][] adjacency = randomAdjacency(NUM_NODES, DEGREE, new Random(1)); + + try (CuVSMatrix matrix = CuVSMatrix.ofArray(adjacency)) { + GPUBuiltHnswGraph serial = newSingleLayerGraph(matrix, 1); + GPUBuiltHnswGraph parallel = newSingleLayerGraph(matrix, NUM_THREADS); + assertGraphsEqual(serial, parallel); + } + } + + @Test + public void writeGraphParallelMatchesSerial() throws Exception { + assumeTrue("cuVS not supported", isSupported()); + int[][] adjacency = randomAdjacency(NUM_NODES, DEGREE, new Random(2)); + + try (CuVSMatrix matrix = CuVSMatrix.ofArray(adjacency); + Directory dir = new ByteBuffersDirectory()) { + // Materialize once, serially, so any difference found below is attributable only to + // writeGraph's own parallelization, not to fillNeighborArray's. + GPUBuiltHnswGraph graph = newSingleLayerGraph(matrix, 1); + + int[][] serialOffsets; + try (IndexOutput out = dir.createOutput("serial", IOContext.DEFAULT)) { + serialOffsets = AcceleratedHNSWUtils.writeGraph(graph, out, 1); + } + int[][] parallelOffsets; + try (IndexOutput out = dir.createOutput("parallel", IOContext.DEFAULT)) { + parallelOffsets = AcceleratedHNSWUtils.writeGraph(graph, out, NUM_THREADS); + } + + assertEquals(serialOffsets.length, parallelOffsets.length); + for (int level = 0; level < serialOffsets.length; level++) { + assertArrayEquals( + "per-node byte-length offsets differ for level " + level, + serialOffsets[level], + parallelOffsets[level]); + } + + assertArrayEquals( + "writeGraph's parallel level-0 encoding produced different bytes than the serial path", + readAllBytes(dir, "serial"), + readAllBytes(dir, "parallel")); + } + } + + private static GPUBuiltHnswGraph newSingleLayerGraph(CuVSMatrix layer0Adjacency, int numThreads) { + // A single layer (layer 0 only): the constructor never consults layerNodes in that case, so + // the placeholder null entry mirrors the convention used elsewhere for "layer 0 needs no node + // list" without actually being read. + return new GPUBuiltHnswGraph( + NUM_NODES, + /* dimensions= */ 4, + Arrays.asList((int[]) null), + List.of(layer0Adjacency), + numThreads); + } + + /** Every node/level's in-order arc list must match exactly between the two graphs. */ + private static void assertGraphsEqual(HnswGraph a, HnswGraph b) throws Exception { + assertEquals(a.numLevels(), b.numLevels()); + for (int level = 0; level < a.numLevels(); level++) { + int[] nodes = NodesIterator.getSortedNodes(a.getNodesOnLevel(level)); + for (int node : nodes) { + assertArrayEquals( + "node " + node + " at level " + level + " has different neighbors", + arcsOf(a, level, node), + arcsOf(b, level, node)); + } + } + } + + private static int[] arcsOf(HnswGraph graph, int level, int node) throws Exception { + graph.seek(level, node); + List arcs = new ArrayList<>(); + for (int n = graph.nextNeighbor(); n != NO_MORE_DOCS; n = graph.nextNeighbor()) { + arcs.add(n); + } + return arcs.stream().mapToInt(Integer::intValue).toArray(); + } + + private static byte[] readAllBytes(Directory dir, String name) throws Exception { + try (IndexInput in = dir.openInput(name, IOContext.DEFAULT)) { + byte[] bytes = new byte[(int) in.length()]; + in.readBytes(bytes, 0, bytes.length); + return bytes; + } + } + + /** + * A deterministic, seeded pseudo-adjacency. It doesn't need to be a real CAGRA graph -- only a + * realistic shape (fixed degree, valid node ids) -- since {@code GPUBuiltHnswGraph} and {@code + * AcceleratedHNSWUtils#writeGraph} don't interpret the neighbor ids semantically. + */ + private static int[][] randomAdjacency(int numNodes, int degree, Random random) { + int[][] adjacency = new int[numNodes][degree]; + for (int[] row : adjacency) { + for (int j = 0; j < degree; j++) { + row[j] = random.nextInt(numNodes); + } + } + return adjacency; + } +}