From 5e6ce7e83682d01f08a58c3e867cb3192e8acf12 Mon Sep 17 00:00:00 2001 From: EC2 Default User Date: Fri, 11 Sep 2026 02:09:14 +0000 Subject: [PATCH 1/3] Optimize CAGRA HNSW FBIN bulk indexing (cuvs-issue-2592) --- java/cuvs-lucene/README.md | 130 ++++- .../cuvs/lucene/AcceleratedHNSWParams.java | 20 + .../cuvs/lucene/AcceleratedHNSWUtils.java | 220 ++++---- .../lucene/AcceleratedHnswGraphOutput.java | 59 +- .../lucene/BorrowedDatasetFieldWriter.java | 79 +++ .../BorrowedDatasetHnswVectorsWriter.java | 144 +++++ .../cuvs/lucene/BorrowedDatasetOutput.java | 24 + .../cuvs/lucene/BulkIndexingContext.java | 95 ++++ .../cuvs/lucene/CagraHnswBuildMetrics.java | 92 ++++ .../cuvs/lucene/CagraHnswBulkIndexWriter.java | 266 ++++++++- .../cuvs/lucene/CuVS2510GPUVectorsReader.java | 16 +- .../lucene/ExternalFbinBuildValidation.java | 26 + .../cuvs/lucene/ExternalFbinFileRegistry.java | 154 ++++++ .../lucene/ExternalFbinFlatVectorsReader.java | 312 +++++++++++ .../nvidia/cuvs/lucene/ExternalFbinIO.java | 148 +++++ .../cuvs/lucene/ExternalFbinOptions.java | 27 + .../cuvs/lucene/ExternalFbinReference.java | 265 +++++++++ .../lucene/ExternalFbinReferenceWriter.java | 116 ++++ .../lucene/ExternalFbinScanCoordinator.java | 283 ++++++++++ .../cuvs/lucene/ExternalFloat32Dataset.java | 138 +++++ .../nvidia/cuvs/lucene/FbinFileMetadata.java | 61 +++ .../nvidia/cuvs/lucene/GPUBuiltHnswGraph.java | 15 +- .../lucene/ImmutableExternalFbinDataset.java | 104 ++++ .../lucene/ImmutableExternalFbinOutput.java | 69 +++ .../lucene/Lucene101AcceleratedHNSWCodec.java | 23 +- .../Lucene99AcceleratedHNSWVectorsFormat.java | 39 +- ...ratedHNSWBinaryQuantizedVectorsWriter.java | 3 +- ...ratedHNSWScalarQuantizedVectorsWriter.java | 3 +- .../nvidia/cuvs/lucene/MappedFbinDataset.java | 117 ++++ .../lucene/MappedSelfContainedOutput.java | 118 ++++ .../NativeFlatBufferedHNSWVectorsWriter.java | 23 +- .../cuvs/lucene/NativeFlatVectorsWriter.java | 66 ++- .../lucene/TestAcceleratedHNSWParams.java | 10 + .../cuvs/lucene/TestAcceleratedHNSWUtils.java | 61 +++ .../TestBulkExternalFbinApiSurface.java | 62 +++ .../cuvs/lucene/TestBulkFbinIndexStorage.java | 232 ++++++++ .../TestCagraHnswBuildMetricsConcurrency.java | 97 ++++ .../lucene/TestCagraHnswBulkIndexWriter.java | 26 +- .../TestCagraHnswBulkManualFieldContract.java | 91 +++ ...510GPUVectorsReaderScoreNormalization.java | 41 ++ .../TestExternalFbinFlatVectorsReader.java | 516 ++++++++++++++++++ .../lucene/TestExternalFbinReference.java | 213 ++++++++ .../TestExternalFbinScanCoordinator.java | 200 +++++++ .../lucene/TestExternalFloat32Dataset.java | 36 ++ .../cuvs/lucene/TestFbinFileMetadata.java | 33 ++ .../cuvs/lucene/TestGPUBuiltHnswGraph.java | 49 ++ .../cuvs/lucene/TestMappedFbinDataset.java | 64 +++ .../lucene/TestNativeFlatBufferedCleanup.java | 49 ++ 48 files changed, 4899 insertions(+), 136 deletions(-) create mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/BorrowedDatasetFieldWriter.java create mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/BorrowedDatasetHnswVectorsWriter.java create mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/BorrowedDatasetOutput.java create mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/BulkIndexingContext.java create mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraHnswBuildMetrics.java create mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinBuildValidation.java create mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinFileRegistry.java create mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinFlatVectorsReader.java create mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinIO.java create mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinOptions.java create mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinReference.java create mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinReferenceWriter.java create mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinScanCoordinator.java create mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFloat32Dataset.java create mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FbinFileMetadata.java create mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ImmutableExternalFbinDataset.java create mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ImmutableExternalFbinOutput.java create mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/MappedFbinDataset.java create mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/MappedSelfContainedOutput.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWUtils.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBulkExternalFbinApiSurface.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBulkFbinIndexStorage.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBuildMetricsConcurrency.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBulkManualFieldContract.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVS2510GPUVectorsReaderScoreNormalization.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestExternalFbinFlatVectorsReader.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestExternalFbinReference.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestExternalFbinScanCoordinator.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestExternalFloat32Dataset.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestFbinFileMetadata.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGPUBuiltHnswGraph.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMappedFbinDataset.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferedCleanup.java diff --git a/java/cuvs-lucene/README.md b/java/cuvs-lucene/README.md index ae7885d142..2a192abebd 100644 --- a/java/cuvs-lucene/README.md +++ b/java/cuvs-lucene/README.md @@ -7,8 +7,9 @@ This is a project for using [cuVS](https://github.com/rapidsai/cuvs), NVIDIA's G 1. [What is cuvs-lucene?](#what-is-cuvs-lucene) 2. [Installing cuvs-lucene](#installing-cuvs-lucene) 3. [Getting Started](#getting-started) -4. [Contributing](#contributing) -5. [References](#references) +4. [CAGRA/HNSW bulk indexing](#cagrahnsw-bulk-indexing) +5. [Contributing](#contributing) +6. [References](#references) ## What is cuvs-lucene? @@ -130,6 +131,131 @@ mvn -q compile org.codehaus.mojo:exec-maven-plugin:3.5.1:java \ For more examples, including one that indexes and searches entirely on the GPU using `CuVS2510GPUSearchCodec`, please refer to the [`examples/`](../../examples/java/cuvs-lucene) directory. +## CAGRA/HNSW bulk indexing + +`CagraHnswBulkIndexWriter` is an opt-in API for controlled, offline creation of +`Lucene101AcceleratedHNSWCodec` indexes. It owns the `IndexWriter` configuration and flush +lifecycle needed by native buffering. It supports GPU CAGRA graph construction followed by CPU +HNSW search; it does not build `CuVS2510GPUSearchCodec` indexes and, unlike the generic codec, +requires cuVS GPU support instead of falling back to CPU indexing. + +### Ingestion modes + +| Mode | Public entry point | Vector handling | Finished index | +| --- | --- | --- | --- | +| Generic Lucene codec | Configure `Lucene101AcceleratedHNSWCodec` on an application-owned `IndexWriter` | Normal Lucene document ingestion and lifecycle | Self-contained; normal Lucene flush, sort, and merge behavior | +| Streaming/native-buffered bulk | Construct `CagraHnswBulkIndexWriter` directly, call `build(VectorSource, Config)`, or call `indexFbin(Path, Config)` | Copies each vector into an exactly sized native host matrix before that segment's GPU build | Self-contained; the source can be removed after a successful build | +| Mapped, self-contained bulk | `indexMappedFbin(Path, Config)` | CAGRA and the flat-vector writer share a read-only mapping, avoiding per-row decode and ingest copies; the flat vectors are still written into the index | Self-contained; the FBIN can be removed after a successful build | +| Immutable external-FBIN bulk | Register the FBIN, then call `indexImmutableFbin(Registration, Config, ExternalFbinOptions)` | CAGRA reads the mapping and the index stores a content-addressed descriptor instead of duplicating the flat-vector payload | Not self-contained; exact scoring reads the registered FBIN | + +The streaming/native-buffered mode has three entry shapes: + +- The direct constructor is the manual, single-segment form. The caller creates every document, + promises the exact document count, and calls `addDocument` exactly that many times. +- `build(VectorSource, Config)` owns document creation for a caller-owned, forward-only source. It + can build multiple segments sequentially, but does not close the source and does not support the + overlapped pipeline. +- `indexFbin(Path, Config)` owns the FBIN reader. It supports sequential segment builds and an + overlapped multi-segment pipeline in which host ingest can overlap a serialized GPU build. + `FieldCallback` can add non-vector fields in either one-shot form. + +Both mapped modes currently require `segments(1, false)`. They also require exactly one dense +`FLOAT32` vector field, one vector per document, and row `i` to correspond to Lucene document and +vector ordinal `i`. + +### Ownership and sharing + +All bulk forms own their internal `IndexWriter`, disable index sorting, automatic flushes, compound +files, and merges, and commit only after the promised vector count has been written. Closing a +manual writer with too few vectors rolls it back; adding too many vectors is rejected. A failed +segment is rolled back. A sequential multi-segment build can already have committed earlier +segments when a later segment fails, so callers must treat and replace that target as a failed +build. + +The generic, streaming, and mapped self-contained results can be moved by copying the Lucene index +directory. The immutable external-FBIN result must be shipped as two artifacts: the Lucene index +directory and the exact FBIN identified by the descriptor's complete-file SHA-256. The descriptor +does not persist a host path, so another process or node may place the FBIN at a different local +path. Before opening the index, that process must register the path and digest: + +```java +try (ExternalFbinFileRegistry.Registration lease = + ExternalFbinFileRegistry.register(localFbin, sha256Hex); + Directory directory = FSDirectory.open(indexPath); + DirectoryReader reader = DirectoryReader.open(directory)) { + // Search while both the reader and registration lease are open. +} +``` + +Registrations are process-local and reference counted. Keep at least one lease open for the entire +lifetime of every reader using that content ID. Do not modify, replace, relocate, or delete the +registered FBIN while a build or reader is alive. An index directory copied without its FBIN, or a +process that has not registered the local FBIN, cannot be opened. Registering a different path for +the same content ID while an existing registration is live is rejected. External-FBIN indexes also +require a cuvs-lucene runtime that understands their external-vector descriptor; they are not +readable by a stock Lucene runtime alone. + +### Validation + +Every FBIN path is structurally checked for a positive shape and an exact +`8 + rows * dimensions * 4` byte length. `Config` also rejects a mismatch between the Lucene +similarity and the cuVS graph metric: + +- `EUCLIDEAN` requires `L2Expanded`. +- `DOT_PRODUCT` and `MAXIMUM_INNER_PRODUCT` require `InnerProduct`. +- `COSINE` requires `CosineExpanded`. + +The manual bulk form additionally validates that every document has exactly one vector field with +the configured name, dimension, `FLOAT32` encoding, and similarity. Borrowed mapped data is not +scanned component by component, so the caller must ensure that every source float is finite. + +Immutable external-FBIN builds require the caller to supply a previously established SHA-256 for +the complete file. `ExternalFbinFileRegistry.register` validates and allowlists the real path, but +does not compute the digest. `ExternalFbinBuildValidation` controls the build-time scan: + +| Validation | Build-time behavior | +| --- | --- | +| `TRUSTED_IMMUTABLE` | Checks the registration, reference metadata, header, range, and file length; does not scan the payload or verify the digest. Use only with storage that already enforces the content identity. | +| `PREFETCH` | Sequentially scans the referenced payload as read-ahead while graph construction runs; trusts the supplied digest. | +| `VERIFY_SHA256` | Hashes the complete FBIN and compares it with the supplied digest while graph construction runs. A mismatch prevents the commit. | + +`ExternalFbinOptions.scanHeadStartBytes()` can delay graph construction until a requested amount of +the selected payload has been scanned. It is a scheduling control, not a reduction in validation: +the selected scan still completes before commit. It must be zero for `TRUSTED_IMMUTABLE` and no +larger than the referenced payload. + +Opening an external-FBIN index verifies the descriptor checksum, field metadata, registered path, +header, range, and file length, but does not hash the full file. Lucene's `checkIntegrity()` hashes +the complete external FBIN and compares it with the persisted SHA-256. + +### Metrics + +Supply one `CagraHnswBuildMetrics` per build through `Config.Builder.metrics`. After the build, +`snapshot()` returns a stable `Map`, and `appendTo(target, prefix)` adds that +snapshot to a benchmark result map. + +Metric keys use three namespaces: + +- `stage//seconds`, `stage//count`, and optional `stage//bytes` report aggregated + graph build, graph conversion and output, mapped flat output, external scan/overlap, and bulk + writer commit/close stages. The commit wall includes the flush and its nested GPU/output work; + nested stage durations must not be added to it. +- `counter/` reports values such as logical adjacency bytes, mapped flat chunks, and external + scan progress at CAGRA start and end. +- `gauge/` records effective graph-build parameters, including graph degrees, writer threads, + NN-descent iterations, and IVF-PQ dimensions, lists, probes, and k-means iterations when used. + +### Why these controls are bulk-only + +Native buffering and borrowed FBIN storage depend on guarantees that a general Lucene codec cannot +make: one controlled flush, an exact dense vector count and order, no index sort, no merge during +the build, and an external file whose immutability and reader lifetime are managed by the +application. `CagraHnswBulkIndexWriter` owns those conditions. Storage and external-file controls +therefore remain on the bulk API (`Config`, `ExternalFbinOptions`, and +`ExternalFbinFileRegistry`) rather than `AcceleratedHNSWParams` or public codec constructors. The +generic codec remains suitable for application-owned Lucene, Solr, Elasticsearch, and OpenSearch +indexing lifecycles without adding an external-file contract they cannot enforce. + ## Contributing If you are interested in contributing to cuvs-lucene, please read the cuVS [Contributing guide](https://docs.nvidia.com/cuvs/developer-guide/contributing). 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 b57614539a..a7463f0a03 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 @@ -56,6 +56,7 @@ public static enum Strategy { public static final int DEFAULT_INT_GRAPH_DEGREE = 128; public static final int DEFAULT_GRAPH_DEGREE = 64; public static final int DEFAULT_HNSW_LAYERS = 1; + public static final long DEFAULT_HNSW_LAYER_SEED = 44L; 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 = @@ -81,6 +82,7 @@ public static enum Strategy { private final int intermediateGraphDegree; private final int graphdegree; private final int hnswLayers; + private final long hnswLayerSeed; private final int maxConn; private final int beamWidth; private final CagraGraphBuildAlgo cagraGraphBuildAlgo; @@ -101,6 +103,7 @@ public static enum 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 hnswLayerSeed The deterministic seed used to sample HNSW upper-layer nodes. * @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]. Only @@ -120,6 +123,7 @@ private AcceleratedHNSWParams( int intermediateGraphDegree, int graphdegree, int hnswLayers, + long hnswLayerSeed, int maxConn, int beamWidth, CagraGraphBuildAlgo cagraGraphBuildAlgo, @@ -135,6 +139,7 @@ private AcceleratedHNSWParams( this.intermediateGraphDegree = intermediateGraphDegree; this.graphdegree = graphdegree; this.hnswLayers = hnswLayers; + this.hnswLayerSeed = hnswLayerSeed; this.maxConn = maxConn; this.beamWidth = beamWidth; this.cagraGraphBuildAlgo = cagraGraphBuildAlgo; @@ -183,6 +188,11 @@ public int getHnswLayers() { return hnswLayers; } + /** Returns the deterministic seed used to sample nodes for HNSW upper layers. */ + public long getHnswLayerSeed() { + return hnswLayerSeed; + } + /** * Get the max connection parameter * @@ -290,6 +300,8 @@ public String toString() { + graphdegree + ", hnswLayers=" + hnswLayers + + ", hnswLayerSeed=" + + hnswLayerSeed + ", maxConn=" + maxConn + ", beamWidth=" @@ -322,6 +334,7 @@ public static class Builder { private int intermediateGraphDegree = DEFAULT_INT_GRAPH_DEGREE; private int graphdegree = DEFAULT_GRAPH_DEGREE; private int hnswLayers = DEFAULT_HNSW_LAYERS; + private long hnswLayerSeed = DEFAULT_HNSW_LAYER_SEED; private int maxConn = DEFAULT_MAX_CONN; private int beamWidth = DEFAULT_BEAM_WIDTH; private CagraGraphBuildAlgo cagraGraphBuildAlgo = DEFAULT_CAGRA_GRAPH_BUILD_ALGO; @@ -387,6 +400,12 @@ public Builder withHNSWLayer(int hnswLayers) { return this; } + /** Sets the deterministic seed used to sample nested HNSW upper-layer nodes. */ + public Builder withHnswLayerSeed(long hnswLayerSeed) { + this.hnswLayerSeed = hnswLayerSeed; + return this; + } + /** * Set the max connections parameter while building HNSW index with fallback mechanism * Valid range - Minimum: {@value MIN_MAX_CONN}, Maximum: {@value MAX_MAX_CONN} @@ -624,6 +643,7 @@ public AcceleratedHNSWParams build() { intermediateGraphDegree, graphdegree, hnswLayers, + hnswLayerSeed, maxConn, beamWidth, cagraGraphBuildAlgo, 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 fdd7dd441f..9cf30f02c2 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 @@ -5,6 +5,7 @@ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_HNSW_LAYER_SEED; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.getCuVSResourcesInstance; import static com.nvidia.cuvs.lucene.Utils.createByteMatrixFromArray; @@ -101,9 +102,32 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( QuantizationType quantization, int numThreads) throws Throwable { + return createMultiLayerHnswGraph( + fieldInfo, + dimensions, + adjacencyListMatrix, + vectorDataset, + hnswLayers, + params, + quantization, + numThreads, + DEFAULT_HNSW_LAYER_SEED); + } + + static GPUBuiltHnswGraph createMultiLayerHnswGraph( + FieldInfo fieldInfo, + int dimensions, + CuVSMatrix adjacencyListMatrix, + CuVSMatrix vectorDataset, + int hnswLayers, + CagraIndexParams params, + QuantizationType quantization, + int numThreads, + long hnswLayerSeed) + throws Throwable { int size = (int) vectorDataset.size(); - int M = Math.ceilDiv((int) adjacencyListMatrix.columns(), 2); + int m = hnswM(adjacencyListMatrix.columns()); // Store all layers data List layerNodes = new ArrayList<>(); @@ -113,70 +137,94 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( layerNodes.add(null); // Layer 0 contains all nodes, so we don't need to store node list layerAdjacencies.add(adjacencyListMatrix); - int currentLayerSize = size; - int layerIndex = 1; - Random random = new Random(); - - while (layerIndex < hnswLayers && currentLayerSize > 1) { - // Calculate size for next layer (1/M of current layer) - int nextLayerSize = Math.max(2, currentLayerSize / M); - // Select nodes for this layer - SortedSet selectedNodesSet = new TreeSet<>(); - - if (layerIndex == 1) { - // Select from all nodes (Layer 0) - while (selectedNodesSet.size() < nextLayerSize) { - selectedNodesSet.add(random.nextInt(size)); - } - } else { - // Select from previous layer nodes - int[] prevLayerNodes = layerNodes.get(layerNodes.size() - 1); - while (selectedNodesSet.size() < nextLayerSize) { - selectedNodesSet.add(prevLayerNodes[random.nextInt(prevLayerNodes.length)]); + List selectedLayers = selectUpperLayerNodes(size, hnswLayers, m, hnswLayerSeed); + try (OwnedMatrices ownedUpperAdjacencies = new OwnedMatrices()) { + for (int[] selectedNodes : selectedLayers) { + layerNodes.add(selectedNodes); + + CuVSMatrix upperAdjacency; + if (quantization == QuantizationType.NONE) { + // Read only the sampled rows from the native matrix, without a full-dataset heap copy. + float[][] selectedVectors = new float[selectedNodes.length][dimensions]; + for (int i = 0; i < selectedNodes.length; i++) { + vectorDataset.getRow(selectedNodes[i]).toArray(selectedVectors[i]); + } + upperAdjacency = + buildCagraGraphForSubset( + selectedVectors, selectedNodes, 0, params, dimensions, quantization); + } else { + // Binary packs 8 dims/byte; scalar uses one byte per dimension. + int bytesPerVector = (int) vectorDataset.columns(); + byte[][] selectedVectors = new byte[selectedNodes.length][bytesPerVector]; + for (int i = 0; i < selectedNodes.length; i++) { + vectorDataset.getRow(selectedNodes[i]).toArray(selectedVectors[i]); + } + upperAdjacency = + buildCagraGraphForSubset( + selectedVectors, selectedNodes, bytesPerVector, params, dimensions, quantization); } + layerAdjacencies.add(ownedUpperAdjacencies.add(upperAdjacency)); } - // Convert to sorted array - int[] selectedNodes = - selectedNodesSet.stream().mapToInt(Integer::intValue).sorted().toArray(); - - layerNodes.add(selectedNodes); + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies, numThreads); + } + } - if (quantization == QuantizationType.NONE) { - // 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++) { - vectorDataset.getRow(selectedNodes[i]).toArray(selectedVectors[i]); - } + static int hnswM(long graphDegree) { + return Math.toIntExact(Math.max(1L, Math.ceilDiv(graphDegree, 2L))); + } - // Build CAGRA graph for this layer - layerAdjacencies.add( - buildCagraGraphForSubset( - selectedVectors, selectedNodes, 0, params, dimensions, quantization)); - } else { - // 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++) { - vectorDataset.getRow(selectedNodes[i]).toArray(selectedVectors[i]); - } + static List selectUpperLayerNodes(int size, int hnswLayers, int m, long hnswLayerSeed) { + List selectedLayers = new ArrayList<>(); + if (hnswLayers <= 1 || m <= 1 || size <= m) { + return selectedLayers; + } - // Build CAGRA graph for this layer - layerAdjacencies.add( - buildCagraGraphForSubset( - selectedVectors, selectedNodes, bytesPerVector, params, dimensions, quantization)); + Random random = new Random(hnswLayerSeed); + int currentLayerSize = size; + int[] previousLayerNodes = null; + while (selectedLayers.size() + 1 < hnswLayers && currentLayerSize > m) { + int nextLayerSize = Math.ceilDiv(currentLayerSize, m); + SortedSet selectedNodesSet = new TreeSet<>(); + while (selectedNodesSet.size() < nextLayerSize) { + int sampledIndex = random.nextInt(currentLayerSize); + selectedNodesSet.add( + previousLayerNodes == null ? sampledIndex : previousLayerNodes[sampledIndex]); } - - // Update for next iteration + int[] selectedNodes = selectedNodesSet.stream().mapToInt(Integer::intValue).toArray(); + selectedLayers.add(selectedNodes); + previousLayerNodes = selectedNodes; currentLayerSize = nextLayerSize; - layerIndex++; + } + return selectedLayers; + } - // Use different seed for each layer - random = new Random(new Random().nextLong()); + private static final class OwnedMatrices implements AutoCloseable { + private final List matrices = new ArrayList<>(); + + CuVSMatrix add(CuVSMatrix matrix) { + matrices.add(matrix); + return matrix; } - // Create the multi-layer graph with all layers - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies, numThreads); + @Override + public void close() { + RuntimeException failure = null; + for (int i = matrices.size() - 1; i >= 0; i--) { + try { + matrices.get(i).close(); + } catch (RuntimeException e) { + if (failure == null) { + failure = e; + } else { + failure.addSuppressed(e); + } + } + } + if (failure != null) { + throw failure; + } + } } /** @@ -191,48 +239,34 @@ private static CuVSMatrix buildCagraGraphForSubset( QuantizationType quantization) throws Throwable { - CuVSMatrix subsetDataset; - - if (quantization == QuantizationType.BINARY) { - subsetDataset = createByteMatrixFromArray((byte[][]) vectors, bytesPerVector); - } else if (quantization == QuantizationType.SCALAR) { - subsetDataset = createByteMatrixFromArray((byte[][]) vectors, dimensions); - } else { - subsetDataset = CuVSMatrix.ofArray((float[][]) vectors); - } - - // Build CAGRA index for the subset - CagraIndex subsetIndex = - CagraIndex.newBuilder(getCuVSResourcesInstance()) - .withDataset(subsetDataset) - .withIndexParams(params) - .build(); - - // Get adjacency list from subset CAGRA index - CuVSMatrix cagraGraph = subsetIndex.getGraph(); - - long numNodes = cagraGraph.size(); - long degree = cagraGraph.columns(); - - // Create a re-mapped adjacency list - int[][] remappedAdjacency = new int[(int) numNodes][(int) degree]; - - for (int i = 0; i < numNodes; i++) { - RowView rv = cagraGraph.getRow(i); - for (int j = 0; j < degree && j < rv.size(); j++) { - int subsetIndex1 = rv.getAsInt(j); - // Map subset index to original node ID - if (subsetIndex1 >= 0 && subsetIndex1 < selectedNodes.length) { - remappedAdjacency[i][j] = selectedNodes[subsetIndex1]; - } else { - // Invalid index, use self-reference - remappedAdjacency[i][j] = selectedNodes[i]; + try (CuVSMatrix subsetDataset = + quantization == QuantizationType.BINARY + ? createByteMatrixFromArray((byte[][]) vectors, bytesPerVector) + : quantization == QuantizationType.SCALAR + ? createByteMatrixFromArray((byte[][]) vectors, dimensions) + : CuVSMatrix.ofArray((float[][]) vectors); + CagraIndex subsetIndex = + CagraIndex.newBuilder(getCuVSResourcesInstance()) + .withDataset(subsetDataset) + .withIndexParams(params) + .build()) { + CuVSMatrix cagraGraph = subsetIndex.getGraph(); + int numNodes = Math.toIntExact(cagraGraph.size()); + int degree = Math.toIntExact(cagraGraph.columns()); + int[][] remappedAdjacency = new int[numNodes][degree]; + + for (int i = 0; i < numNodes; i++) { + RowView rv = cagraGraph.getRow(i); + for (int j = 0; j < degree && j < rv.size(); j++) { + int subsetNeighbor = rv.getAsInt(j); + remappedAdjacency[i][j] = + subsetNeighbor >= 0 && subsetNeighbor < selectedNodes.length + ? selectedNodes[subsetNeighbor] + : selectedNodes[i]; } } + return CuVSMatrix.ofArray(remappedAdjacency); } - - subsetIndex.close(); - return CuVSMatrix.ofArray(remappedAdjacency); } /** diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHnswGraphOutput.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHnswGraphOutput.java index e5317f4a01..db9dbb9366 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHnswGraphOutput.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHnswGraphOutput.java @@ -51,13 +51,23 @@ final class AcceleratedHnswGraphOutput implements Closeable { } private final AcceleratedHNSWParams acceleratedHNSWParams; + private final CagraHnswBuildMetrics metrics; private final IndexOutput hnswMeta; private final IndexOutput hnswVectorIndex; private boolean finished; AcceleratedHnswGraphOutput(SegmentWriteState state, AcceleratedHNSWParams acceleratedHNSWParams) throws IOException { + this(state, acceleratedHNSWParams, new CagraHnswBuildMetrics()); + } + + AcceleratedHnswGraphOutput( + SegmentWriteState state, + AcceleratedHNSWParams acceleratedHNSWParams, + CagraHnswBuildMetrics metrics) + throws IOException { this.acceleratedHNSWParams = acceleratedHNSWParams; + this.metrics = metrics; String vemFileName = IndexFileNames.segmentFileName( state.segmentInfo.name, state.segmentSuffix, HNSW_META_CODEC_EXT); @@ -115,7 +125,18 @@ void writeField(FieldInfo fieldInfo, List vectors) throws IOException { * never double-materialised on the Java heap. */ void writeField(FieldInfo fieldInfo, CuVSMatrix dataset) throws IOException { - try (dataset) { + writeField(fieldInfo, dataset, true); + } + + /** Builds a graph from a caller-owned matrix without closing that matrix. */ + void writeBorrowedField(FieldInfo fieldInfo, CuVSMatrix dataset) throws IOException { + writeField(fieldInfo, dataset, false); + } + + private void writeField(FieldInfo fieldInfo, CuVSMatrix dataset, boolean closeDataset) + throws IOException { + long graphOutputStart = CagraHnswBuildMetrics.start(); + try { int size = (int) dataset.size(); if (size == 0) { writeEmpty(fieldInfo, hnswMeta); @@ -131,13 +152,23 @@ void writeField(FieldInfo fieldInfo, CuVSMatrix dataset) throws IOException { CagraIndexParams params = CagraIndexParamsFactory.create( acceleratedHNSWParams, dataset.size(), dataset.columns()); + recordEffectiveParams(params); + long adjacencyBytes = + Math.multiplyExact( + Math.multiplyExact((long) size, params.getGraphDegree()), Integer.BYTES); + metrics.addCounter("logical cagra adjacency bytes", adjacencyBytes); + long stageStart = CagraHnswBuildMetrics.start(); try (CagraIndex cagraIndex = CagraIndex.newBuilder(getCuVSResourcesInstance()) .withDataset(dataset) .withIndexParams(params) .build()) { + metrics.stop("cagra-build [GPU]", stageStart); + stageStart = CagraHnswBuildMetrics.start(); CuVSMatrix adjacencyListMatrix = cagraIndex.getGraph(); + metrics.stop("graph-view [GPU]", stageStart); int dimensions = fieldInfo.getVectorDimension(); + stageStart = CagraHnswBuildMetrics.start(); GPUBuiltHnswGraph hnswGraph = createMultiLayerHnswGraph( fieldInfo, @@ -147,11 +178,15 @@ void writeField(FieldInfo fieldInfo, CuVSMatrix dataset) throws IOException { acceleratedHNSWParams.getHnswLayers(), params, QuantizationType.NONE, - acceleratedHNSWParams.getWriterThreads()); + acceleratedHNSWParams.getWriterThreads(), + acceleratedHNSWParams.getHnswLayerSeed()); + metrics.stop("hnsw-convert [PCIe+CPU]", stageStart, adjacencyBytes); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); + stageStart = CagraHnswBuildMetrics.start(); int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; + metrics.stop("write-graph [CPU+DISK]", stageStart, vectorIndexLength); writeMeta( hnswVectorIndex, hnswMeta, @@ -165,6 +200,26 @@ void writeField(FieldInfo fieldInfo, CuVSMatrix dataset) throws IOException { } catch (Throwable t) { Utils.handleThrowable(t); } + } finally { + metrics.stop("graph-output wall [CPU+GPU+DISK]", graphOutputStart); + if (closeDataset) { + dataset.close(); + } + } + } + + private void recordEffectiveParams(CagraIndexParams params) { + metrics.setGauge("effective graph degree", params.getGraphDegree()); + metrics.setGauge("effective intermediate graph degree", params.getIntermediateGraphDegree()); + metrics.setGauge("effective writer threads", params.getNumWriterThreads()); + metrics.setGauge("effective nn-descent iterations", params.getNNDescentNumIterations()); + if (params.getCagraGraphBuildAlgo() == CagraIndexParams.CagraGraphBuildAlgo.IVF_PQ) { + var ivf = params.getCuVSIvfPqParams(); + metrics.setGauge("effective ivf-pq dimensions", ivf.getIndexParams().getPqDim()); + metrics.setGauge("effective ivf-pq lists", ivf.getIndexParams().getnLists()); + metrics.setGauge("effective ivf-pq probes", ivf.getSearchParams().getnProbes()); + metrics.setGauge( + "effective ivf-pq kmeans iterations", ivf.getIndexParams().getKmeansNIters()); } } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/BorrowedDatasetFieldWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/BorrowedDatasetFieldWriter.java new file mode 100644 index 0000000000..db1558812d --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/BorrowedDatasetFieldWriter.java @@ -0,0 +1,79 @@ +/* + * 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 org.apache.lucene.codecs.KnnFieldVectorsWriter; +import org.apache.lucene.index.DocsWithFieldSet; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.util.RamUsageEstimator; + +/** Validates dense row-ordered placeholders while vectors remain in a borrowed native dataset. */ +final class BorrowedDatasetFieldWriter extends KnnFieldVectorsWriter { + + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(BorrowedDatasetFieldWriter.class); + + private final FieldInfo fieldInfo; + private final int expectedCount; + private final DocsWithFieldSet docsWithField = new DocsWithFieldSet(); + private int count; + + BorrowedDatasetFieldWriter(FieldInfo fieldInfo, int expectedCount) { + if (fieldInfo.getVectorEncoding() != VectorEncoding.FLOAT32) { + throw new IllegalArgumentException("Borrowed FBIN datasets require FLOAT32 vectors"); + } + this.fieldInfo = fieldInfo; + this.expectedCount = expectedCount; + } + + @Override + public void addValue(int docID, Object vectorValue) throws IOException { + if (docID != count) { + throw new IllegalArgumentException( + "Borrowed FBIN datasets require dense document IDs in row order; expected " + + count + + " but got " + + docID); + } + if (!(vectorValue instanceof float[] vector) + || vector.length != fieldInfo.getVectorDimension()) { + throw new IllegalArgumentException( + "Expected a float vector with dimension " + + fieldInfo.getVectorDimension() + + " for field " + + fieldInfo.name); + } + if (count >= expectedCount) { + throw new IllegalStateException( + "Borrowed FBIN vector count exceeds expected count " + expectedCount); + } + docsWithField.add(docID); + count++; + } + + FieldInfo fieldInfo() { + return fieldInfo; + } + + DocsWithFieldSet docsWithField() { + return docsWithField; + } + + int count() { + return count; + } + + @Override + public Object copyValue(Object vectorValue) { + throw new UnsupportedOperationException(); + } + + @Override + public long ramBytesUsed() { + return SHALLOW_SIZE; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/BorrowedDatasetHnswVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/BorrowedDatasetHnswVectorsWriter.java new file mode 100644 index 0000000000..324e31ccee --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/BorrowedDatasetHnswVectorsWriter.java @@ -0,0 +1,144 @@ +/* + * 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.AcceleratedHNSWUtils.printInfoStream; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.closeCuVSResourcesInstance; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import org.apache.lucene.codecs.KnnFieldVectorsWriter; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.MergeState; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.Sorter.DocMap; +import org.apache.lucene.util.IOUtils; +import org.apache.lucene.util.InfoStream; +import org.apache.lucene.util.RamUsageEstimator; + +/** Bulk-only writer for vectors already resident in a caller-owned native FBIN mapping. */ +final class BorrowedDatasetHnswVectorsWriter extends KnnVectorsWriter { + + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(BorrowedDatasetHnswVectorsWriter.class); + + private final BulkIndexingContext context; + private final InfoStream infoStream; + private final List fields = new ArrayList<>(1); + private final AcceleratedHnswGraphOutput graphOutput; + private final BorrowedDatasetOutput vectorOutput; + private boolean finished; + + BorrowedDatasetHnswVectorsWriter( + SegmentWriteState state, AcceleratedHNSWParams graphBuildParams, BulkIndexingContext context) + throws IOException { + this.context = context; + this.infoStream = state.infoStream; + AcceleratedHnswGraphOutput newGraphOutput = null; + BorrowedDatasetOutput newVectorOutput = null; + boolean success = false; + try { + newGraphOutput = new AcceleratedHnswGraphOutput(state, graphBuildParams, context.metrics()); + newVectorOutput = + switch (context.storage()) { + case MAPPED_SELF_CONTAINED -> new MappedSelfContainedOutput(state, context.metrics()); + case IMMUTABLE_EXTERNAL -> new ImmutableExternalFbinOutput(state, context); + case NATIVE_BUFFERED -> + throw new IllegalArgumentException("Borrowed writer requires a mapped dataset"); + }; + success = true; + } finally { + graphOutput = newGraphOutput; + vectorOutput = newVectorOutput; + if (!success) { + IOUtils.closeWhileHandlingException(newVectorOutput, newGraphOutput); + } + } + printInfoStream(infoStream, getClass().getSimpleName(), "borrowed FBIN writer initialized"); + } + + @Override + public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException { + if (!fields.isEmpty()) { + throw new UnsupportedOperationException( + "Mapped FBIN bulk builds support exactly one vector field"); + } + if (fieldInfo.getVectorDimension() != context.dataset().dimensions()) { + throw new IllegalArgumentException( + "Vector field dimension " + + fieldInfo.getVectorDimension() + + " does not match mapped FBIN dimension " + + context.dataset().dimensions()); + } + BorrowedDatasetFieldWriter field = + new BorrowedDatasetFieldWriter(fieldInfo, context.exactVectorCount()); + fields.add(field); + return field; + } + + @Override + public void flush(int maxDoc, DocMap sortMap) throws IOException { + if (sortMap != null) { + throw new UnsupportedOperationException("Mapped FBIN bulk builds do not support index sort"); + } + if (fields.size() != 1) { + throw new IllegalStateException( + "Mapped FBIN bulk builds require exactly one vector field, got " + fields.size()); + } + BorrowedDatasetFieldWriter field = fields.getFirst(); + if (field.count() != context.exactVectorCount()) { + throw new IllegalStateException( + "Expected " + + context.exactVectorCount() + + " mapped FBIN vectors but received " + + field.count()); + } + if (maxDoc != context.exactVectorCount()) { + throw new IllegalStateException( + "Mapped FBIN bulk builds require one vector per document; maxDoc=" + + maxDoc + + ", expected=" + + context.exactVectorCount()); + } + vectorOutput.writeField( + field.fieldInfo(), context.dataset(), maxDoc, field.docsWithField(), graphOutput); + } + + @Override + public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) { + throw new UnsupportedOperationException("Mapped FBIN bulk builds do not support merges"); + } + + @Override + public void finish() throws IOException { + if (finished) { + throw new IllegalStateException("already finished"); + } + finished = true; + vectorOutput.finish(); + graphOutput.finish(); + } + + @Override + public void close() throws IOException { + printInfoStream(infoStream, getClass().getSimpleName(), "closing resources"); + try { + IOUtils.close(vectorOutput, graphOutput); + } finally { + closeCuVSResourcesInstance(); + } + } + + @Override + public long ramBytesUsed() { + long total = SHALLOW_SIZE; + for (BorrowedDatasetFieldWriter field : fields) { + total += field.ramBytesUsed(); + } + return total; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/BorrowedDatasetOutput.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/BorrowedDatasetOutput.java new file mode 100644 index 0000000000..130db2c7de --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/BorrowedDatasetOutput.java @@ -0,0 +1,24 @@ +/* + * 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; +import org.apache.lucene.index.DocsWithFieldSet; +import org.apache.lucene.index.FieldInfo; + +/** Persistence strategy for a bulk build whose vectors live in a borrowed native dataset. */ +interface BorrowedDatasetOutput extends Closeable { + + void writeField( + FieldInfo field, + ExternalFloat32Dataset dataset, + int maxDoc, + DocsWithFieldSet docsWithField, + AcceleratedHnswGraphOutput graphOutput) + throws IOException; + + void finish() throws IOException; +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/BulkIndexingContext.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/BulkIndexingContext.java new file mode 100644 index 0000000000..5151c23587 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/BulkIndexingContext.java @@ -0,0 +1,95 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import java.util.Objects; + +/** Package-private build state reachable only through {@link CagraHnswBulkIndexWriter}. */ +final class BulkIndexingContext { + + enum Storage { + NATIVE_BUFFERED, + MAPPED_SELF_CONTAINED, + IMMUTABLE_EXTERNAL + } + + private final int exactVectorCount; + private final Storage storage; + private final ExternalFloat32Dataset dataset; + private final ExternalFbinReference reference; + private final ExternalFbinOptions externalOptions; + private final CagraHnswBuildMetrics metrics; + + static BulkIndexingContext nativeBuffered(int exactVectorCount, CagraHnswBuildMetrics metrics) { + return new BulkIndexingContext( + exactVectorCount, Storage.NATIVE_BUFFERED, null, null, null, metrics); + } + + static BulkIndexingContext mapped(ExternalFloat32Dataset dataset, CagraHnswBuildMetrics metrics) { + Objects.requireNonNull(dataset, "dataset"); + return new BulkIndexingContext( + dataset.rows(), Storage.MAPPED_SELF_CONTAINED, dataset, null, null, metrics); + } + + static BulkIndexingContext external( + ImmutableExternalFbinDataset dataset, + ExternalFbinOptions options, + CagraHnswBuildMetrics metrics) { + Objects.requireNonNull(dataset, "dataset"); + Objects.requireNonNull(options, "options"); + ExternalFbinReference reference = dataset.reference(); + if (options.scanHeadStartBytes() > reference.payloadLength()) { + throw new IllegalArgumentException( + "scanHeadStartBytes exceeds the referenced payload length: " + + options.scanHeadStartBytes() + + " > " + + reference.payloadLength()); + } + return new BulkIndexingContext( + dataset.rows(), Storage.IMMUTABLE_EXTERNAL, dataset.dataset(), reference, options, metrics); + } + + private BulkIndexingContext( + int exactVectorCount, + Storage storage, + ExternalFloat32Dataset dataset, + ExternalFbinReference reference, + ExternalFbinOptions externalOptions, + CagraHnswBuildMetrics metrics) { + if (exactVectorCount <= 0) { + throw new IllegalArgumentException("exactVectorCount must be positive"); + } + this.exactVectorCount = exactVectorCount; + this.storage = Objects.requireNonNull(storage, "storage"); + this.dataset = dataset; + this.reference = reference; + this.externalOptions = externalOptions; + this.metrics = Objects.requireNonNull(metrics, "metrics"); + } + + int exactVectorCount() { + return exactVectorCount; + } + + Storage storage() { + return storage; + } + + ExternalFloat32Dataset dataset() { + return dataset; + } + + ExternalFbinReference reference() { + return reference; + } + + ExternalFbinOptions externalOptions() { + return externalOptions; + } + + CagraHnswBuildMetrics metrics() { + return metrics; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraHnswBuildMetrics.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraHnswBuildMetrics.java new file mode 100644 index 0000000000..297e54a63b --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraHnswBuildMetrics.java @@ -0,0 +1,92 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Objects; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.LongAdder; + +/** Per-build, thread-safe measurements for the bulk CAGRA-to-HNSW pipeline. */ +public final class CagraHnswBuildMetrics { + + private final Map stages = new ConcurrentHashMap<>(); + private final Map counters = new ConcurrentHashMap<>(); + private final Map gauges = new ConcurrentHashMap<>(); + + static long start() { + return System.nanoTime(); + } + + void stop(String stage, long startedAtNanos) { + record(stage, System.nanoTime() - startedAtNanos, 0L); + } + + void stop(String stage, long startedAtNanos, long byteCount) { + record(stage, System.nanoTime() - startedAtNanos, byteCount); + } + + void record(String stage, long elapsedNanos, long byteCount) { + if (elapsedNanos < 0L || byteCount < 0L) { + throw new IllegalArgumentException("elapsedNanos and byteCount must be non-negative"); + } + stages.compute( + stage, + (ignored, current) -> + current == null + ? new StageMeasurement(elapsedNanos, 1L, byteCount) + : current.plus(elapsedNanos, byteCount)); + } + + void addCounter(String name, long value) { + counters.computeIfAbsent(name, ignored -> new LongAdder()).add(value); + } + + /** Records a configuration value that must remain identical across every segment in a build. */ + void setGauge(String name, long value) { + Objects.requireNonNull(name, "name"); + Long existing = gauges.putIfAbsent(name, value); + if (existing != null && existing.longValue() != value) { + throw new IllegalStateException( + "Gauge \"" + name + "\" changed from " + existing + " to " + value); + } + } + + /** Returns a stable machine-readable copy; later measurements do not mutate it. */ + public Map snapshot() { + Map result = new LinkedHashMap<>(); + stages.keySet().stream() + .sorted() + .forEach( + stage -> { + StageMeasurement measurement = stages.get(stage); + String prefix = "stage/" + stage; + result.put(prefix + "/seconds", measurement.nanos() / 1e9); + result.put(prefix + "/count", measurement.count()); + if (measurement.bytes() != 0L) { + result.put(prefix + "/bytes", measurement.bytes()); + } + }); + counters.keySet().stream() + .sorted() + .forEach(name -> result.put("counter/" + name, counters.get(name).sum())); + gauges.keySet().stream() + .sorted() + .forEach(name -> result.put("gauge/" + name, gauges.get(name))); + return Map.copyOf(result); + } + + /** Adds this build's snapshot under {@code prefix} to a benchmark result map. */ + public void appendTo(Map target, String prefix) { + snapshot().forEach((key, value) -> target.put(prefix + "/" + key, value)); + } + + private record StageMeasurement(long nanos, long count, long bytes) { + StageMeasurement plus(long additionalNanos, long additionalBytes) { + return new StageMeasurement(nanos + additionalNanos, count + 1L, bytes + additionalBytes); + } + } +} 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 index c1cd5e71b4..161bea8cd2 100644 --- 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 @@ -4,6 +4,7 @@ */ package com.nvidia.cuvs.lucene; +import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; import java.io.Closeable; import java.io.IOException; import java.nio.file.Files; @@ -26,6 +27,7 @@ import org.apache.lucene.index.IndexWriterConfig; import org.apache.lucene.index.IndexableField; import org.apache.lucene.index.NoMergePolicy; +import org.apache.lucene.index.VectorEncoding; import org.apache.lucene.index.VectorSimilarityFunction; import org.apache.lucene.misc.store.HardlinkCopyDirectoryWrapper; import org.apache.lucene.store.Directory; @@ -68,6 +70,7 @@ public final class CagraHnswBulkIndexWriter implements Closeable { private static final int DEFAULT_CHUNK_SIZE_MB = 32; private final IndexWriter writer; + private final Config config; private final int exactVectorCount; private int documentsAdded; private boolean closed; @@ -95,9 +98,21 @@ public final class CagraHnswBulkIndexWriter implements Closeable { public CagraHnswBulkIndexWriter( Directory directory, IndexWriterConfig conf, Config config, int exactVectorCount) throws Exception { + this( + directory, + conf, + config, + BulkIndexingContext.nativeBuffered(exactVectorCount, config.metrics())); + } + + private CagraHnswBulkIndexWriter( + Directory directory, IndexWriterConfig conf, Config config, BulkIndexingContext bulkContext) + throws Exception { Objects.requireNonNull(directory, "directory"); Objects.requireNonNull(conf, "conf"); Objects.requireNonNull(config, "config"); + Objects.requireNonNull(bulkContext, "bulkContext"); + int exactVectorCount = bulkContext.exactVectorCount(); if (exactVectorCount <= 0) { throw new IllegalArgumentException("exactVectorCount must be > 0, got " + exactVectorCount); } @@ -107,9 +122,10 @@ public CagraHnswBulkIndexWriter( + " buffering requires an unsorted single-segment build); leave" + " IndexWriterConfig.indexSort unset"); } + this.config = config; this.exactVectorCount = exactVectorCount; - Codec codec = new Lucene101AcceleratedHNSWCodec(config.graphBuildParams(), exactVectorCount); + Codec codec = new Lucene101AcceleratedHNSWCodec(config.graphBuildParams(), bulkContext); IndexWriterConfig ownedConf = new IndexWriterConfig(conf.getAnalyzer()) .setSimilarity(conf.getSimilarity()) @@ -124,10 +140,10 @@ public CagraHnswBulkIndexWriter( } /** - * 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. + * Adds one document, exactly like {@link IndexWriter#addDocument}. {@code doc} must contain + * exactly one vector field whose name, dimension, and similarity match {@code config}; that field + * is routed into the native flat buffer automatically by the underlying codec. Any number of + * non-vector fields may also be present and are indexed normally. */ public long addDocument(Iterable doc) throws IOException { if (closed) { @@ -137,11 +153,58 @@ public long addDocument(Iterable doc) throws IOExcepti throw new IllegalStateException( "addDocument called more than exactVectorCount (" + exactVectorCount + ") times"); } - long seqNo = writer.addDocument(doc); + List fields = validateAndSnapshotDocument(doc); + long seqNo = writer.addDocument(fields); documentsAdded++; return seqNo; } + private List validateAndSnapshotDocument( + Iterable document) { + Objects.requireNonNull(document, "doc"); + List fields = new ArrayList<>(); + int vectorFields = 0; + for (IndexableField field : document) { + fields.add(field); + var fieldType = field.fieldType(); + if (fieldType.vectorDimension() == 0) { + continue; + } + vectorFields++; + if (!field.name().equals(config.fieldName())) { + throw new IllegalArgumentException( + "Vector field name \"" + + field.name() + + "\" does not match configured field \"" + + config.fieldName() + + "\""); + } + if (fieldType.vectorDimension() != config.dimensions()) { + throw new IllegalArgumentException( + "Vector field dimension " + + fieldType.vectorDimension() + + " does not match configured dimension " + + config.dimensions()); + } + if (fieldType.vectorEncoding() != VectorEncoding.FLOAT32) { + throw new IllegalArgumentException("CAGRA/HNSW bulk indexing requires FLOAT32 vectors"); + } + if (fieldType.vectorSimilarityFunction() != config.similarity()) { + throw new IllegalArgumentException( + "Vector field similarity " + + fieldType.vectorSimilarityFunction() + + " does not match configured similarity " + + config.similarity()); + } + } + if (vectorFields != 1) { + throw new IllegalArgumentException( + "Each bulk-indexed document must contain exactly one vector field; found " + + vectorFields); + } + return fields; + } + /** * 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 @@ -173,8 +236,22 @@ public void close() throws IOException { } throw mismatch; } - try (writer) { - writer.commit(); + try (Closeable writerCloser = this::closeUnderlyingWriter) { + long commitStartedAt = CagraHnswBuildMetrics.start(); + try { + writer.commit(); + } finally { + config.metrics().stop("bulk writer commit wall [CPU+GPU+DISK]", commitStartedAt); + } + } + } + + private void closeUnderlyingWriter() throws IOException { + long closeStartedAt = CagraHnswBuildMetrics.start(); + try { + writer.close(); + } finally { + config.metrics().stop("bulk writer close [DISK]", closeStartedAt); } } @@ -273,12 +350,9 @@ 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(); - } + FbinFileMetadata metadata = FbinFileMetadata.read(fbinPath); + int total = metadata.rows(); + int dim = metadata.dimensions(); if (dim != config.dimensions()) { throw new IllegalArgumentException( "fbinPath dimension (" @@ -297,6 +371,140 @@ public static void indexFbin( } } + /** + * Builds a normal, self-contained Lucene index while CAGRA and the flat-vector writer consume a + * shared read-only mapping of {@code fbinPath}. This avoids decoding and copying every row during + * document ingest while preserving a conventional {@code .vec} file in the finished index. + */ + public static void indexMappedFbin(Path fbinPath, Config config) throws Exception { + indexMappedFbin(fbinPath, config, null); + } + + /** As {@link #indexMappedFbin(Path, Config)}, with a callback for non-vector fields. */ + public static void indexMappedFbin(Path fbinPath, Config config, FieldCallback callback) + throws Exception { + Objects.requireNonNull(fbinPath, "fbinPath"); + Objects.requireNonNull(config, "config"); + requireSingleSegmentMappedBuild(config); + try (MappedFbinDataset mapped = MappedFbinDataset.map(fbinPath)) { + validateMappedShape(mapped.rows(), mapped.dimensions(), config); + buildBorrowed( + config, callback, BulkIndexingContext.mapped(mapped.dataset(), config.metrics())); + } + } + + /** + * Builds a non-self-contained Lucene index whose exact-vector scoring reads from the immutable + * FBIN represented by {@code source}. The caller must retain {@code source} through every reader + * lifetime and replicate the referenced FBIN together with the Lucene directory. + */ + public static void indexImmutableFbin( + ExternalFbinFileRegistry.Registration source, Config config, ExternalFbinOptions options) + throws Exception { + indexImmutableFbin(source, config, options, null); + } + + /** As {@link #indexImmutableFbin(ExternalFbinFileRegistry.Registration, Config, + * ExternalFbinOptions)}, with a callback for non-vector fields. */ + public static void indexImmutableFbin( + ExternalFbinFileRegistry.Registration source, + Config config, + ExternalFbinOptions options, + FieldCallback callback) + throws Exception { + Objects.requireNonNull(source, "source"); + Objects.requireNonNull(config, "config"); + Objects.requireNonNull(options, "options"); + requireSingleSegmentMappedBuild(config); + FbinFileMetadata metadata = FbinFileMetadata.read(source.path()); + int rows = metadata.rows(); + int dimensions = metadata.dimensions(); + validateMappedShape(rows, dimensions, config); + try (ImmutableExternalFbinDataset dataset = source.map(0, rows)) { + buildBorrowed( + config, callback, BulkIndexingContext.external(dataset, options, config.metrics())); + } + } + + private static void requireSingleSegmentMappedBuild(Config config) { + if (config.targetDirectory() == null) { + throw new IllegalArgumentException("targetDirectory is required for a one-shot bulk build"); + } + if (config.numSegments() != 1 || config.overlapped()) { + throw new IllegalArgumentException( + "Mapped FBIN bulk builds currently require segments(1, false)"); + } + } + + private static void validateMappedShape(int rows, int dimensions, Config config) { + if (rows <= 0 || dimensions != config.dimensions()) { + throw new IllegalArgumentException( + "FBIN shape " + + rows + + " x " + + dimensions + + " does not match configured dimension " + + config.dimensions()); + } + } + + private static void buildBorrowed( + Config config, FieldCallback callback, BulkIndexingContext context) throws Exception { + long ingestStartedAt = CagraHnswBuildMetrics.start(); + try (Directory directory = FSDirectory.open(config.targetDirectory())) { + IndexWriterConfig writerConfig = + new IndexWriterConfig().setOpenMode(IndexWriterConfig.OpenMode.CREATE); + CagraHnswBulkIndexWriter writer = + new CagraHnswBulkIndexWriter(directory, writerConfig, config, context); + try { + addBorrowedDocuments(writer, config, callback, context.exactVectorCount()); + config.metrics().stop("borrowed document ingest [CPU]", ingestStartedAt); + writer.close(); + } catch (Throwable failure) { + try { + writer.abort(); + } catch (Throwable cleanupFailure) { + failure.addSuppressed(cleanupFailure); + } + throw failure; + } + } + } + + private static void addBorrowedDocuments( + CagraHnswBulkIndexWriter writer, Config config, FieldCallback callback, int count) + throws IOException { + float[] placeholder = new float[config.dimensions()]; + if (callback == null) { + Document document = new Document(); + StringField idField = + config.idFieldName() == null + ? null + : new StringField(config.idFieldName(), "0", Field.Store.YES); + if (idField != null) { + document.add(idField); + } + document.add(new KnnFloatVectorField(config.fieldName(), placeholder, config.similarity())); + for (int id = 0; id < count; id++) { + if (idField != null) { + idField.setStringValue(Integer.toString(id)); + } + writer.addDocument(document); + } + return; + } + + for (int id = 0; id < count; id++) { + Document document = new Document(); + if (config.idFieldName() != null) { + document.add(new StringField(config.idFieldName(), Integer.toString(id), Field.Store.YES)); + } + document.add(new KnnFloatVectorField(config.fieldName(), placeholder, config.similarity())); + callback.addFields(document, id); + writer.addDocument(document); + } + } + /** * 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 @@ -518,6 +726,7 @@ public static final class Config { private final int numSegments; private final boolean overlapped; private final int pipelineDepth; + private final CagraHnswBuildMetrics metrics; private Config(Builder b) { this.fieldName = b.fieldName; @@ -529,6 +738,7 @@ private Config(Builder b) { this.numSegments = b.numSegments; this.overlapped = b.overlapped; this.pipelineDepth = b.pipelineDepth; + this.metrics = b.metrics; } public String fieldName() { @@ -571,6 +781,11 @@ public int pipelineDepth() { return pipelineDepth; } + /** Per-build metrics accumulator shared by all segments in this configuration. */ + public CagraHnswBuildMetrics metrics() { + return metrics; + } + public static Builder builder() { return new Builder(); } @@ -586,6 +801,7 @@ public static final class Builder { private int numSegments = 1; private boolean overlapped = false; private int pipelineDepth = 2; + private CagraHnswBuildMetrics metrics = new CagraHnswBuildMetrics(); /** Sets the vector field name and dimensionality; required. */ public Builder field(String fieldName, int dimensions, VectorSimilarityFunction similarity) { @@ -647,12 +863,34 @@ public Builder pipelineDepth(int pipelineDepth) { return this; } + /** Supplies a caller-owned metrics accumulator for this build. */ + public Builder metrics(CagraHnswBuildMetrics metrics) { + this.metrics = Objects.requireNonNull(metrics, "metrics"); + 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"); + CuvsDistanceType expectedMetric = + switch (similarity) { + case EUCLIDEAN -> CuvsDistanceType.L2Expanded; + case DOT_PRODUCT, MAXIMUM_INNER_PRODUCT -> CuvsDistanceType.InnerProduct; + case COSINE -> CuvsDistanceType.CosineExpanded; + }; + if (graphBuildParams.getCuvsDistanceType() != expectedMetric) { + throw new IllegalArgumentException( + "Graph metric " + + graphBuildParams.getCuvsDistanceType() + + " does not match field similarity " + + similarity + + " (expected " + + expectedMetric + + ")"); + } return new Config(this); } } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java index 906a310ede..3917879970 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java @@ -49,6 +49,7 @@ import org.apache.lucene.store.ReadAdvice; import org.apache.lucene.util.Bits; import org.apache.lucene.util.IOUtils; +import org.apache.lucene.util.VectorUtil; import org.apache.lucene.util.hnsw.IntToIntFunction; /** @@ -446,8 +447,19 @@ private interface FloatToFloatFunction { * @return an instance of the FloatToFloatFunction */ private static FloatToFloatFunction getScoreNormalizationFunc(VectorSimilarityFunction sim) { - // TODO: check for different similarities - return score -> (1f / (1f + score)); + return cuvsValue -> toLuceneScore(sim, cuvsValue); + } + + /** Convert a public cuVS search value into the score contract used by Lucene's collectors. */ + static float toLuceneScore(VectorSimilarityFunction sim, float cuvsValue) { + return switch (sim) { + case EUCLIDEAN -> 1.0f / (1.0f + cuvsValue); + // Public cuVS InnerProduct results contain the natural dot product. + case DOT_PRODUCT -> Math.max((1.0f + cuvsValue) / 2.0f, 0.0f); + // cuVS CosineExpanded reports 1-cos(a,b), while Lucene scales cosine to [0, 1]. + case COSINE -> Math.max(1.0f - cuvsValue / 2.0f, 0.0f); + case MAXIMUM_INNER_PRODUCT -> VectorUtil.scaleMaxInnerProductScore(cuvsValue); + }; } /** diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinBuildValidation.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinBuildValidation.java new file mode 100644 index 0000000000..1e9bdb13a1 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinBuildValidation.java @@ -0,0 +1,26 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +/** Controls validation and sequential read-ahead for an immutable external FBIN build. */ +public enum ExternalFbinBuildValidation { + /** + * Trusts that the registered file was previously verified and is protected by immutable, + * content-addressed storage. Performs structural checks but no full payload scan. + */ + TRUSTED_IMMUTABLE, + + /** + * Sequentially scans the referenced payload alongside graph construction to provide controlled + * read-ahead, but trusts the precomputed digest. + */ + PREFETCH, + + /** + * Sequentially hashes the complete FBIN alongside graph construction and fails the build if its + * SHA-256 differs. This also acts as read-ahead, but may extend the critical path on slow storage. + */ + VERIFY_SHA256 +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinFileRegistry.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinFileRegistry.java new file mode 100644 index 0000000000..0d99f8451f --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinFileRegistry.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 java.io.Closeable; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; + +/** + * Process-local allowlist that resolves content-addressed external FBIN references. + * + *

Lucene recreates codecs through no-argument SPI constructors, so writer-time configuration is + * not available when a process later opens an index. Every indexing or search process must register + * the immutable local file for each referenced content ID before opening the index. Registrations + * are reference counted and should remain alive for at least the complete reader lifetime. + * + *

Registration validates the path and content-ID syntax but deliberately does not hash the file. + * Use {@link ExternalFbinBuildValidation#VERIFY_SHA256} while building, or run Lucene integrity + * checks, unless the registered storage already enforces the supplied content identity. + */ +public final class ExternalFbinFileRegistry { + + private static final Object LOCK = new Object(); + private static final Map ENTRIES = new HashMap<>(); + + private ExternalFbinFileRegistry() {} + + /** Registers an allowlisted local file for a complete-file SHA-256 content identity. */ + public static Registration register(Path path, String sha256Hex) throws IOException { + Objects.requireNonNull(path, "path"); + byte[] digest = ExternalFbinReference.parseSha256(sha256Hex); + String contentId = ExternalFbinReference.contentId(digest); + Path realPath = path.toRealPath(); + if (!Files.isRegularFile(realPath) || !Files.isReadable(realPath)) { + throw new IOException("External FBIN must be a readable regular file: " + realPath); + } + + synchronized (LOCK) { + Entry existing = ENTRIES.get(contentId); + if (existing == null) { + ENTRIES.put(contentId, new Entry(realPath, 1)); + } else if (existing.path.equals(realPath)) { + existing.references++; + } else { + throw new IllegalStateException( + "Content ID " + + contentId + + " is already registered to " + + existing.path + + "; refusing ambiguous replacement with " + + realPath); + } + } + return new Registration(realPath, contentId); + } + + static Path resolve(ExternalFbinReference reference) throws IOException { + synchronized (LOCK) { + Entry entry = ENTRIES.get(reference.contentId()); + if (entry == null) { + throw new IOException( + "No allowlisted external FBIN is registered for " + + reference.contentId() + + ". Register it with ExternalFbinFileRegistry before opening the index."); + } + return entry.path; + } + } + + static void clearForTests() { + synchronized (LOCK) { + ENTRIES.clear(); + } + } + + private static final class Entry { + private final Path path; + private int references; + + private Entry(Path path, int references) { + this.path = path; + this.references = references; + } + } + + /** A scoped registry lease. Closing the final lease removes the path from the allowlist. */ + public static final class Registration implements Closeable { + + private final Path path; + private final String contentId; + private boolean closed; + + private Registration(Path path, String contentId) { + this.path = path; + this.contentId = contentId; + } + + /** Creates a path-independent reference to a contiguous row range in the registered FBIN. */ + public synchronized ExternalFbinReference reference(int firstRow, int rowCount) + throws IOException { + ensureOpen(); + return ExternalFbinReference.fromFile(path, contentId, firstRow, rowCount); + } + + /** + * Maps a registered row range and couples its native build dataset to the persisted reference. + */ + public synchronized ImmutableExternalFbinDataset map(int firstRow, int rowCount) + throws IOException { + ensureOpen(); + return ImmutableExternalFbinDataset.map( + ExternalFbinReference.fromFile(path, contentId, firstRow, rowCount)); + } + + private void ensureOpen() { + if (closed) { + throw new IllegalStateException("External FBIN registration is closed"); + } + } + + public String contentId() { + return contentId; + } + + public Path path() { + return path; + } + + @Override + public void close() { + synchronized (this) { + if (closed) { + return; + } + closed = true; + } + synchronized (LOCK) { + Entry entry = ENTRIES.get(contentId); + if (entry != null && entry.path.equals(path)) { + entry.references--; + if (entry.references == 0) { + ENTRIES.remove(contentId); + } + } + } + } + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinFlatVectorsReader.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinFlatVectorsReader.java new file mode 100644 index 0000000000..a9c267a79c --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinFlatVectorsReader.java @@ -0,0 +1,312 @@ +/* + * 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.file.Path; +import org.apache.lucene.codecs.CodecUtil; +import org.apache.lucene.codecs.hnsw.DefaultFlatVectorScorer; +import org.apache.lucene.codecs.hnsw.FlatVectorsReader; +import org.apache.lucene.codecs.lucene95.OffHeapFloatVectorValues.DenseOffHeapVectorValues; +import org.apache.lucene.index.ByteVectorValues; +import org.apache.lucene.index.CorruptIndexException; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexFileNames; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.store.ChecksumIndexInput; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.store.MMapDirectory; +import org.apache.lucene.store.ReadAdvice; +import org.apache.lucene.util.IOUtils; +import org.apache.lucene.util.RamUsageEstimator; +import org.apache.lucene.util.hnsw.RandomVectorScorer; + +/** Reads dense float32 vectors from a registered immutable FBIN instead of a Lucene {@code .vec}. */ +final class ExternalFbinFlatVectorsReader extends FlatVectorsReader { + + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(ExternalFbinFlatVectorsReader.class); + + private final FieldInfo fieldInfo; + private final ExternalFbinReference reference; + private final Directory externalDirectory; + private final IndexInput sourceInput; + private final IndexInput payloadInput; + + ExternalFbinFlatVectorsReader(SegmentReadState state) throws IOException { + this(state, MMapDirectory.DEFAULT_MAX_CHUNK_SIZE); + } + + ExternalFbinFlatVectorsReader(SegmentReadState state, long maxMmapChunkSize) throws IOException { + super(DefaultFlatVectorScorer.INSTANCE); + Descriptor descriptor = readDescriptor(state); + fieldInfo = descriptor.fieldInfo(); + reference = descriptor.reference(); + + Path sourcePath = ExternalFbinIO.validateAndResolve(reference); + Directory newDirectory = null; + IndexInput newSource = null; + IndexInput newPayload = null; + boolean success = false; + try { + Path parent = sourcePath.getParent(); + if (parent == null) { + throw new IOException("External FBIN path has no parent directory: " + sourcePath); + } + newDirectory = new MMapDirectory(parent, maxMmapChunkSize); + IOContext context = state.context.withReadAdvice(ReadAdvice.RANDOM); + newSource = newDirectory.openInput(sourcePath.getFileName().toString(), context); + validateOpenedSource(newSource, reference); + newPayload = + newSource.slice( + "external-fbin-payload", + reference.payloadOffset(), + reference.payloadLength(), + ReadAdvice.RANDOM); + success = true; + } finally { + if (!success) { + IOUtils.closeWhileHandlingException(newPayload, newSource, newDirectory); + } + } + externalDirectory = newDirectory; + sourceInput = newSource; + payloadInput = newPayload; + } + + static boolean hasExternalMarker(SegmentReadState state) throws IOException { + String segmentMarker = + state.segmentInfo.getAttribute( + ExternalFbinReferenceWriter.segmentAttribute(state.segmentSuffix)); + if (segmentMarker != null + && !ExternalFbinReferenceWriter.SEGMENT_ATTRIBUTE_VALUE.equals(segmentMarker)) { + throw new CorruptIndexException( + "Unsupported external FBIN segment marker version " + segmentMarker, + state.segmentInfo.name); + } + return segmentMarker != null; + } + + private static Descriptor readDescriptor(SegmentReadState state) throws IOException { + String fileName = + IndexFileNames.segmentFileName( + state.segmentInfo.name, state.segmentSuffix, ExternalFbinReferenceWriter.EXTENSION); + try (ChecksumIndexInput input = state.directory.openChecksumInput(fileName)) { + Throwable prior = null; + try { + CodecUtil.checkIndexHeader( + input, + ExternalFbinReferenceWriter.CODEC_NAME, + ExternalFbinReferenceWriter.VERSION_START, + ExternalFbinReferenceWriter.VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + int fieldNumber = input.readInt(); + if (fieldNumber < 0) { + throw new CorruptIndexException("External FBIN descriptor contains no field", input); + } + FieldInfo field = state.fieldInfos.fieldInfo(fieldNumber); + if (field == null) { + throw new CorruptIndexException( + "External FBIN descriptor has unknown field number " + fieldNumber, input); + } + String fieldName = input.readString(); + String encodingName = input.readString(); + String similarityName = input.readString(); + int dimensions = input.readVInt(); + int rows = input.readInt(); + int fileRows = input.readInt(); + long firstRow = input.readLong(); + long fileLength = input.readLong(); + long payloadOffset = input.readLong(); + long payloadLength = input.readLong(); + byte[] sha256 = new byte[ExternalFbinReference.SHA256_BYTES]; + input.readBytes(sha256, 0, sha256.length); + if (input.readInt() != -1) { + throw new CorruptIndexException( + "External FBIN descriptor contains more than one field", input); + } + + final VectorEncoding encoding; + final VectorSimilarityFunction similarity; + try { + encoding = VectorEncoding.valueOf(encodingName); + similarity = VectorSimilarityFunction.valueOf(similarityName); + } catch (IllegalArgumentException e) { + throw new CorruptIndexException( + "Invalid vector metadata in external FBIN descriptor", input, e); + } + if (!field.name.equals(fieldName) + || encoding != VectorEncoding.FLOAT32 + || field.getVectorEncoding() != encoding + || field.getVectorSimilarityFunction() != similarity + || field.getVectorDimension() != dimensions) { + throw new CorruptIndexException( + "External FBIN descriptor does not match FieldInfo for " + field.name, input); + } + if (rows != state.segmentInfo.maxDoc()) { + throw new CorruptIndexException( + "External FBIN rows " + + rows + + " do not match dense segment maxDoc " + + state.segmentInfo.maxDoc(), + input); + } + + ExternalFbinReference reference; + try { + reference = + ExternalFbinReference.fromDescriptor( + sha256, fileLength, payloadOffset, payloadLength, rows, dimensions); + } catch (ArithmeticException | IllegalArgumentException e) { + throw new CorruptIndexException("Invalid external FBIN range metadata", input, e); + } + long rowBytes = Math.multiplyExact((long) dimensions, Float.BYTES); + long derivedFileRows = + (reference.fileLength() - ExternalFbinReference.HEADER_BYTES) / rowBytes; + if (fileRows <= 0 + || fileRows != derivedFileRows + || firstRow != reference.firstRow() + || firstRow + rows > fileRows) { + throw new CorruptIndexException("Inconsistent external FBIN row metadata", input); + } + return new Descriptor(field, reference); + } catch (Throwable t) { + prior = t; + throw t; + } finally { + CodecUtil.checkFooter(input, prior); + } + } + } + + private static void validateOpenedSource(IndexInput source, ExternalFbinReference reference) + throws IOException { + if (source.length() != reference.fileLength()) { + throw new CorruptIndexException( + "External FBIN length changed: expected " + + reference.fileLength() + + " but got " + + source.length(), + source); + } + byte[] headerBytes = new byte[(int) ExternalFbinReference.HEADER_BYTES]; + source.seek(0L); + source.readBytes(headerBytes, 0, headerBytes.length); + ByteBuffer header = ByteBuffer.wrap(headerBytes).order(ByteOrder.LITTLE_ENDIAN); + int fileRows = header.getInt(); + int dimensions = header.getInt(); + if (fileRows <= 0 || dimensions <= 0) { + throw new CorruptIndexException( + "External FBIN header contains a non-positive shape: " + fileRows + " x " + dimensions, + source); + } + final long expectedLength; + try { + long rowBytes = Math.multiplyExact((long) dimensions, Float.BYTES); + expectedLength = + Math.addExact( + ExternalFbinReference.HEADER_BYTES, Math.multiplyExact((long) fileRows, rowBytes)); + } catch (ArithmeticException e) { + throw new CorruptIndexException( + "External FBIN header shape overflows its file length", source, e); + } + if (dimensions != reference.dimensions() + || expectedLength != source.length() + || reference.firstRow() + reference.rows() > fileRows) { + throw new CorruptIndexException( + "External FBIN header or shape does not match the segment descriptor", source); + } + } + + private void requireField(String field, VectorEncoding expectedEncoding) { + if (!fieldInfo.name.equals(field)) { + throw new IllegalArgumentException("field=\"" + field + "\" not found"); + } + if (fieldInfo.getVectorEncoding() != expectedEncoding) { + throw new IllegalArgumentException( + "field=\"" + + field + + "\" is encoded as " + + fieldInfo.getVectorEncoding() + + ", expected " + + expectedEncoding); + } + } + + @Override + public FloatVectorValues getFloatVectorValues(String field) throws IOException { + requireField(field, VectorEncoding.FLOAT32); + return new DenseOffHeapVectorValues( + reference.dimensions(), + reference.rows(), + payloadInput.clone(), + Math.multiplyExact(reference.dimensions(), Float.BYTES), + vectorScorer, + fieldInfo.getVectorSimilarityFunction()); + } + + @Override + public ByteVectorValues getByteVectorValues(String field) { + requireField(field, VectorEncoding.BYTE); + throw new AssertionError("External FBIN fields cannot use byte encoding"); + } + + @Override + public RandomVectorScorer getRandomVectorScorer(String field, float[] target) throws IOException { + requireField(field, VectorEncoding.FLOAT32); + return vectorScorer.getRandomVectorScorer( + fieldInfo.getVectorSimilarityFunction(), getFloatVectorValues(field), target); + } + + @Override + public RandomVectorScorer getRandomVectorScorer(String field, byte[] target) { + requireField(field, VectorEncoding.BYTE); + throw new AssertionError("External FBIN fields cannot use byte encoding"); + } + + @Override + public void checkIntegrity() throws IOException { + IndexInput clone = sourceInput.clone(); + // Lucene IndexInput clones are non-owning. Multi-chunk MemorySegmentIndexInput clones share + // the owner's segment array, so closing a clone would invalidate the live reader. + validateOpenedSource(clone, reference); + ExternalFbinIO.verifySha256(sourceInput, reference); + } + + @Override + public FlatVectorsReader getMergeInstance() { + try { + payloadInput.updateReadAdvice(ReadAdvice.SEQUENTIAL); + } catch (IOException e) { + throw new IllegalStateException("Unable to prepare external FBIN reader for merge", e); + } + return this; + } + + @Override + public void finishMerge() throws IOException { + payloadInput.updateReadAdvice(ReadAdvice.RANDOM); + } + + @Override + public long ramBytesUsed() { + return SHALLOW_SIZE; + } + + @Override + public void close() throws IOException { + IOUtils.close(payloadInput, sourceInput, externalDirectory); + } + + private record Descriptor(FieldInfo fieldInfo, ExternalFbinReference reference) {} +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinIO.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinIO.java new file mode 100644 index 0000000000..ea15c224af --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinIO.java @@ -0,0 +1,148 @@ +/* + * 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.channels.FileChannel; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; +import java.util.Objects; +import java.util.function.LongConsumer; +import org.apache.lucene.index.CorruptIndexException; +import org.apache.lucene.store.IndexInput; + +final class ExternalFbinIO { + + private static final int BUFFER_BYTES = 8 << 20; + + private ExternalFbinIO() {} + + static Path validateAndResolve(ExternalFbinReference reference) throws IOException { + Path path = ExternalFbinFileRegistry.resolve(reference); + ExternalFbinReference actual = + ExternalFbinReference.fromFile( + path, reference.sha256Hex(), Math.toIntExact(reference.firstRow()), reference.rows()); + if (!actual.equals(reference)) { + throw new IOException( + "Registered FBIN metadata does not match persisted reference for " + + reference.contentId() + + ": expected " + + reference + + " but found " + + actual); + } + return path; + } + + static long prefetch(ExternalFbinReference reference) throws IOException { + return prefetch(reference, ignored -> {}); + } + + static long prefetch(ExternalFbinReference reference, LongConsumer progress) throws IOException { + Path path = validateAndResolve(reference); + return scan(path, reference.payloadOffset(), reference.payloadLength(), null, progress); + } + + static long verifySha256(ExternalFbinReference reference) throws IOException { + return verifySha256(reference, ignored -> {}); + } + + static long verifySha256(ExternalFbinReference reference, LongConsumer progress) + throws IOException { + Path path = validateAndResolve(reference); + MessageDigest digest = newSha256(); + long scanned = scan(path, 0L, reference.fileLength(), digest, progress); + byte[] actual = digest.digest(); + if (!MessageDigest.isEqual(reference.sha256(), actual)) { + throw new IOException( + "External FBIN SHA-256 mismatch for " + + path + + ": expected " + + reference.sha256Hex() + + " but got " + + java.util.HexFormat.of().formatHex(actual)); + } + return scanned; + } + + static void verifySha256(IndexInput source, ExternalFbinReference reference) throws IOException { + if (source.length() != reference.fileLength()) { + throw new CorruptIndexException( + "External FBIN length changed: expected " + + reference.fileLength() + + " but opened " + + source.length(), + source); + } + MessageDigest digest = newSha256(); + byte[] buffer = new byte[BUFFER_BYTES]; + // This is a non-owning clone. In particular, closing a multi-chunk MemorySegmentIndexInput + // clone clears the segment array shared with its live owner. + IndexInput clone = source.clone(); + clone.seek(0L); + long remaining = clone.length(); + while (remaining > 0) { + int length = (int) Math.min(buffer.length, remaining); + clone.readBytes(buffer, 0, length); + digest.update(buffer, 0, length); + remaining -= length; + } + byte[] actual = digest.digest(); + if (!MessageDigest.isEqual(reference.sha256(), actual)) { + throw new CorruptIndexException( + "External FBIN SHA-256 mismatch: expected " + + reference.sha256Hex() + + " but got " + + java.util.HexFormat.of().formatHex(actual), + source); + } + } + + private static long scan( + Path path, long offset, long length, MessageDigest optionalDigest, LongConsumer progress) + throws IOException { + Objects.requireNonNull(progress, "progress"); + ByteBuffer buffer = ByteBuffer.allocateDirect((int) Math.min(BUFFER_BYTES, length)); + long position = offset; + long remaining = length; + long scanned = 0L; + try (FileChannel channel = FileChannel.open(path, StandardOpenOption.READ)) { + while (remaining > 0) { + if (Thread.currentThread().isInterrupted()) { + throw new IOException("External FBIN scan interrupted"); + } + buffer.clear(); + buffer.limit((int) Math.min(buffer.capacity(), remaining)); + int read = channel.read(buffer, position); + if (read < 0) { + throw new IOException("Unexpected EOF while scanning external FBIN " + path); + } + if (read == 0) { + throw new IOException("Unable to make progress while scanning external FBIN " + path); + } + if (optionalDigest != null) { + buffer.flip(); + optionalDigest.update(buffer); + } + position += read; + remaining -= read; + scanned += read; + progress.accept(scanned); + } + } + return length; + } + + private static MessageDigest newSha256() { + try { + return MessageDigest.getInstance("SHA-256"); + } catch (NoSuchAlgorithmException e) { + throw new AssertionError("Every Java runtime must provide SHA-256", e); + } + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinOptions.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinOptions.java new file mode 100644 index 0000000000..5e69e8f545 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinOptions.java @@ -0,0 +1,27 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import java.util.Objects; + +/** Validation and scheduling policy for an immutable external-FBIN bulk build. */ +public record ExternalFbinOptions(ExternalFbinBuildValidation validation, long scanHeadStartBytes) { + + public ExternalFbinOptions { + Objects.requireNonNull(validation, "validation"); + if (scanHeadStartBytes < 0L) { + throw new IllegalArgumentException("scanHeadStartBytes must be non-negative"); + } + if (validation == ExternalFbinBuildValidation.TRUSTED_IMMUTABLE && scanHeadStartBytes != 0L) { + throw new IllegalArgumentException( + "scanHeadStartBytes requires PREFETCH or VERIFY_SHA256 validation"); + } + } + + /** Uses complete-file SHA-256 verification with no artificial graph-start delay. */ + public static ExternalFbinOptions verifySha256() { + return new ExternalFbinOptions(ExternalFbinBuildValidation.VERIFY_SHA256, 0L); + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinReference.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinReference.java new file mode 100644 index 0000000000..9821698227 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinReference.java @@ -0,0 +1,265 @@ +/* + * 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.EOFException; +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.Arrays; +import java.util.HexFormat; +import java.util.Objects; + +/** + * A path-independent reference to a dense row-major float32 slice in an immutable FBIN file. + * + *

The complete-file SHA-256 digest is the content identity. The path is deliberately excluded so + * an index cannot request arbitrary host files and can be relocated independently of the source. A + * process must register an allowlisted local path for this content identity through {@link + * ExternalFbinFileRegistry} before opening an index that references it. + */ +public final class ExternalFbinReference { + + static final long HEADER_BYTES = 2L * Integer.BYTES; + static final int SHA256_BYTES = 32; + private static final HexFormat HEX = HexFormat.of(); + + private final byte[] sha256; + private final long fileLength; + private final long payloadOffset; + private final long payloadLength; + private final int rows; + private final int dimensions; + + private ExternalFbinReference( + byte[] sha256, + long fileLength, + long payloadOffset, + long payloadLength, + int rows, + int dimensions) { + this.sha256 = validateSha256(sha256); + if (fileLength <= HEADER_BYTES) { + throw new IllegalArgumentException("FBIN file length must include a non-empty payload"); + } + if (rows <= 0 || dimensions <= 0) { + throw new IllegalArgumentException("rows and dimensions must be positive"); + } + long rowBytes = Math.multiplyExact((long) dimensions, Float.BYTES); + long expectedPayloadLength = Math.multiplyExact((long) rows, rowBytes); + if (payloadLength != expectedPayloadLength) { + throw new IllegalArgumentException( + "payloadLength must equal rows * dimensions * 4; expected " + + expectedPayloadLength + + " but got " + + payloadLength); + } + if (payloadOffset < HEADER_BYTES + || Math.floorMod(payloadOffset - HEADER_BYTES, rowBytes) != 0) { + throw new IllegalArgumentException("payloadOffset must be aligned to an FBIN row boundary"); + } + if (Math.addExact(payloadOffset, payloadLength) > fileLength) { + throw new IllegalArgumentException("Referenced payload extends beyond the FBIN file"); + } + this.fileLength = fileLength; + this.payloadOffset = payloadOffset; + this.payloadLength = payloadLength; + this.rows = rows; + this.dimensions = dimensions; + } + + /** + * Reads and validates an FBIN header and creates a reference to a contiguous row range. + * + *

This method does not hash the file. {@code sha256Hex} must be a previously established digest + * of the complete FBIN file. + */ + public static ExternalFbinReference fromFile( + Path path, String sha256Hex, int firstRow, int rowCount) throws IOException { + Objects.requireNonNull(path, "path"); + if (firstRow < 0 || rowCount <= 0) { + throw new IllegalArgumentException("firstRow must be non-negative and rowCount positive"); + } + + try (FileChannel channel = FileChannel.open(path, StandardOpenOption.READ)) { + ByteBuffer header = ByteBuffer.allocate((int) HEADER_BYTES).order(ByteOrder.LITTLE_ENDIAN); + readFully(channel, header, 0L); + header.flip(); + int fileRows = header.getInt(); + int dimensions = header.getInt(); + if (fileRows <= 0 || dimensions <= 0) { + throw new IOException( + "Invalid FBIN header in " + path + ": " + fileRows + " x " + dimensions); + } + if ((long) firstRow + rowCount > fileRows) { + throw new IllegalArgumentException( + "Requested row range [" + + firstRow + + ", " + + ((long) firstRow + rowCount) + + ") exceeds FBIN row count " + + fileRows); + } + + long rowBytes = Math.multiplyExact((long) dimensions, Float.BYTES); + long expectedFileLength = + Math.addExact(HEADER_BYTES, Math.multiplyExact((long) fileRows, rowBytes)); + long actualFileLength = channel.size(); + if (actualFileLength != expectedFileLength) { + throw new IOException( + "FBIN file length " + + actualFileLength + + " does not match header-derived length " + + expectedFileLength + + " for " + + path); + } + + long payloadOffset = + Math.addExact(HEADER_BYTES, Math.multiplyExact((long) firstRow, rowBytes)); + long payloadLength = Math.multiplyExact((long) rowCount, rowBytes); + return new ExternalFbinReference( + parseSha256(sha256Hex), + actualFileLength, + payloadOffset, + payloadLength, + rowCount, + dimensions); + } + } + + static ExternalFbinReference fromDescriptor( + byte[] sha256, + long fileLength, + long payloadOffset, + long payloadLength, + int rows, + int dimensions) { + return new ExternalFbinReference( + sha256, fileLength, payloadOffset, payloadLength, rows, dimensions); + } + + static byte[] parseSha256(String sha256Hex) { + Objects.requireNonNull(sha256Hex, "sha256Hex"); + String normalized = + sha256Hex.regionMatches(true, 0, "sha256:", 0, "sha256:".length()) + ? sha256Hex.substring("sha256:".length()) + : sha256Hex; + if (normalized.length() != SHA256_BYTES * 2) { + throw new IllegalArgumentException("SHA-256 must contain exactly 64 hexadecimal characters"); + } + try { + return validateSha256(HEX.parseHex(normalized)); + } catch (IllegalArgumentException e) { + throw new IllegalArgumentException("Invalid SHA-256 hexadecimal value", e); + } + } + + static String contentId(byte[] sha256) { + return "sha256:" + HEX.formatHex(validateSha256(sha256)); + } + + private static byte[] validateSha256(byte[] sha256) { + Objects.requireNonNull(sha256, "sha256"); + if (sha256.length != SHA256_BYTES) { + throw new IllegalArgumentException("SHA-256 must contain exactly 32 bytes"); + } + return sha256.clone(); + } + + private static void readFully(FileChannel channel, ByteBuffer target, long position) + throws IOException { + long current = position; + while (target.hasRemaining()) { + int read = channel.read(target, current); + if (read < 0) { + throw new EOFException("Unexpected EOF while reading FBIN header"); + } + if (read == 0) { + throw new IOException("Unable to make progress while reading FBIN header"); + } + current += read; + } + } + + public String contentId() { + return contentId(sha256); + } + + public String sha256Hex() { + return HEX.formatHex(sha256); + } + + public byte[] sha256() { + return sha256.clone(); + } + + public long fileLength() { + return fileLength; + } + + public long payloadOffset() { + return payloadOffset; + } + + public long payloadLength() { + return payloadLength; + } + + public int rows() { + return rows; + } + + public int dimensions() { + return dimensions; + } + + public long firstRow() { + return (payloadOffset - HEADER_BYTES) / ((long) dimensions * Float.BYTES); + } + + @Override + public boolean equals(Object other) { + if (this == other) { + return true; + } + if (!(other instanceof ExternalFbinReference that)) { + return false; + } + return fileLength == that.fileLength + && payloadOffset == that.payloadOffset + && payloadLength == that.payloadLength + && rows == that.rows + && dimensions == that.dimensions + && Arrays.equals(sha256, that.sha256); + } + + @Override + public int hashCode() { + int result = Arrays.hashCode(sha256); + result = 31 * result + Long.hashCode(fileLength); + result = 31 * result + Long.hashCode(payloadOffset); + result = 31 * result + Long.hashCode(payloadLength); + result = 31 * result + rows; + result = 31 * result + dimensions; + return result; + } + + @Override + public String toString() { + return "ExternalFbinReference[contentId=" + + contentId() + + ", firstRow=" + + firstRow() + + ", rows=" + + rows + + ", dimensions=" + + dimensions + + "]"; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinReferenceWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinReferenceWriter.java new file mode 100644 index 0000000000..56dfe33f8e --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinReferenceWriter.java @@ -0,0 +1,116 @@ +/* + * 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; +import org.apache.lucene.codecs.CodecUtil; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.IndexFileNames; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.util.IOUtils; + +final class ExternalFbinReferenceWriter implements Closeable { + + static final String EXTENSION = "vefr"; + static final String CODEC_NAME = "CuVSExternalFbinReference"; + static final int VERSION_START = 0; + static final int VERSION_CURRENT = 0; + static final String SEGMENT_ATTRIBUTE_PREFIX = "cuvs.external_fbin."; + static final String SEGMENT_ATTRIBUTE_VALUE = "1"; + + private final IndexOutput output; + private boolean wroteField; + private boolean finished; + + ExternalFbinReferenceWriter(SegmentWriteState state) throws IOException { + String previous = + state.segmentInfo.putAttribute( + segmentAttribute(state.segmentSuffix), SEGMENT_ATTRIBUTE_VALUE); + if (previous != null && !SEGMENT_ATTRIBUTE_VALUE.equals(previous)) { + throw new IllegalStateException( + "Segment already has an incompatible external FBIN marker: " + previous); + } + String fileName = fileName(state); + IndexOutput newOutput = null; + boolean success = false; + try { + newOutput = state.directory.createOutput(fileName, state.context); + CodecUtil.writeIndexHeader( + newOutput, CODEC_NAME, VERSION_CURRENT, state.segmentInfo.getId(), state.segmentSuffix); + success = true; + } finally { + if (!success) { + IOUtils.closeWhileHandlingException(newOutput); + } + } + output = newOutput; + } + + static String fileName(SegmentWriteState state) { + return IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, EXTENSION); + } + + static String segmentAttribute(String segmentSuffix) { + return SEGMENT_ATTRIBUTE_PREFIX + segmentSuffix; + } + + void writeField(FieldInfo field, ExternalFbinReference reference) throws IOException { + if (finished) { + throw new IllegalStateException("External FBIN reference writer is already finished"); + } + if (wroteField) { + throw new UnsupportedOperationException( + "Immutable external FBIN mode supports exactly one vector field per segment"); + } + if (field.getVectorEncoding() != VectorEncoding.FLOAT32) { + throw new IllegalArgumentException("External FBIN references require FLOAT32 vectors"); + } + if (field.getVectorDimension() != reference.dimensions()) { + throw new IllegalArgumentException( + "Field dimension does not match the external FBIN reference"); + } + + long rowBytes = Math.multiplyExact((long) reference.dimensions(), Float.BYTES); + long fullPayloadBytes = reference.fileLength() - ExternalFbinReference.HEADER_BYTES; + if (Math.floorMod(fullPayloadBytes, rowBytes) != 0) { + throw new IllegalArgumentException("FBIN file length is not an exact number of rows"); + } + int fileRows = Math.toIntExact(fullPayloadBytes / rowBytes); + + output.writeInt(field.number); + output.writeString(field.name); + output.writeString(field.getVectorEncoding().name()); + output.writeString(field.getVectorSimilarityFunction().name()); + output.writeVInt(reference.dimensions()); + output.writeInt(reference.rows()); + output.writeInt(fileRows); + output.writeLong(reference.firstRow()); + output.writeLong(reference.fileLength()); + output.writeLong(reference.payloadOffset()); + output.writeLong(reference.payloadLength()); + output.writeBytes(reference.sha256(), ExternalFbinReference.SHA256_BYTES); + wroteField = true; + } + + void finish() throws IOException { + if (finished) { + throw new IllegalStateException("External FBIN reference writer is already finished"); + } + if (!wroteField) { + throw new IllegalStateException("No external FBIN field was written"); + } + finished = true; + output.writeInt(-1); + CodecUtil.writeFooter(output); + } + + @Override + public void close() throws IOException { + output.close(); + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinScanCoordinator.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinScanCoordinator.java new file mode 100644 index 0000000000..fddadb0253 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinScanCoordinator.java @@ -0,0 +1,283 @@ +/* + * 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.io.InterruptedIOException; +import java.util.Objects; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; +import java.util.function.LongConsumer; + +/** Coordinates a bounded validation lead followed by concurrent FBIN scanning and graph build. */ +final class ExternalFbinScanCoordinator { + + @FunctionalInterface + interface ScanAction { + long scan(LongConsumer progress) throws IOException; + } + + @FunctionalInterface + interface BuildAction { + void run() throws Throwable; + } + + private final long requestedLeadBytes; + private final long requiredScanBytes; + private final long totalScanBytes; + private final String scanStage; + private final CagraHnswBuildMetrics metrics; + private final AtomicLong scannedBytes = new AtomicLong(); + private final CompletableFuture leadReached = new CompletableFuture<>(); + private final ExecutorService executor; + private final Future completion; + private final AtomicBoolean started = new AtomicBoolean(); + + static ExternalFbinScanCoordinator start( + ExternalFbinReference reference, + ExternalFbinBuildValidation validation, + long requestedLeadBytes, + CagraHnswBuildMetrics metrics) { + Objects.requireNonNull(reference, "reference"); + Objects.requireNonNull(validation, "validation"); + if (validation == ExternalFbinBuildValidation.TRUSTED_IMMUTABLE) { + throw new IllegalArgumentException("TRUSTED_IMMUTABLE has no external FBIN scan"); + } + + long totalScanBytes = + validation == ExternalFbinBuildValidation.VERIFY_SHA256 + ? reference.fileLength() + : reference.payloadLength(); + // SHA-256 must consume the complete file in order. For a nonzero first row, account for the + // prefix before the selected payload so the requested number of selected bytes is actually hot. + long requiredScanBytes = + requestedLeadBytes == 0L + ? 0L + : validation == ExternalFbinBuildValidation.VERIFY_SHA256 + ? Math.addExact(reference.payloadOffset(), requestedLeadBytes) + : requestedLeadBytes; + String scanStage = + validation == ExternalFbinBuildValidation.VERIFY_SHA256 + ? "external fbin SHA-256 [DISK+CPU]" + : "external fbin prefetch [DISK]"; + ScanAction scanAction = + validation == ExternalFbinBuildValidation.VERIFY_SHA256 + ? progress -> ExternalFbinIO.verifySha256(reference, progress) + : progress -> ExternalFbinIO.prefetch(reference, progress); + return new ExternalFbinScanCoordinator( + requestedLeadBytes, requiredScanBytes, totalScanBytes, scanStage, scanAction, metrics); + } + + static ExternalFbinScanCoordinator createForTests( + long requestedLeadBytes, long requiredScanBytes, long totalScanBytes, ScanAction scanAction) { + return new ExternalFbinScanCoordinator( + requestedLeadBytes, + requiredScanBytes, + totalScanBytes, + null, + scanAction, + new CagraHnswBuildMetrics()); + } + + private ExternalFbinScanCoordinator( + long requestedLeadBytes, + long requiredScanBytes, + long totalScanBytes, + String scanStage, + ScanAction scanAction, + CagraHnswBuildMetrics metrics) { + if (requestedLeadBytes < 0L + || requiredScanBytes < requestedLeadBytes + || requiredScanBytes > totalScanBytes + || totalScanBytes <= 0L) { + throw new IllegalArgumentException( + "Invalid external FBIN scan bounds: requested=" + + requestedLeadBytes + + ", required=" + + requiredScanBytes + + ", total=" + + totalScanBytes); + } + this.requestedLeadBytes = requestedLeadBytes; + this.requiredScanBytes = requiredScanBytes; + this.totalScanBytes = totalScanBytes; + this.scanStage = scanStage; + this.metrics = Objects.requireNonNull(metrics, "metrics"); + if (requiredScanBytes == 0L) { + leadReached.complete(0L); + } + + executor = + Executors.newSingleThreadExecutor( + task -> { + Thread thread = new Thread(task, "cuvs-external-fbin-validator"); + thread.setDaemon(false); + return thread; + }); + completion = executor.submit(() -> scan(Objects.requireNonNull(scanAction, "scanAction"))); + } + + private long scan(ScanAction scanAction) throws IOException { + long start = CagraHnswBuildMetrics.start(); + try { + long scanned = scanAction.scan(this::recordProgress); + if (scanned != totalScanBytes) { + throw new IOException( + "External FBIN scan reported " + scanned + " bytes; expected " + totalScanBytes); + } + recordProgress(scanned); + // Do not release a full-scan lead until final verification (including digest comparison) has + // succeeded. Partial leads intentionally allow a later validation failure to abort the build. + leadReached.complete(scanned); + return scanned; + } catch (IOException | RuntimeException | Error failure) { + leadReached.completeExceptionally(failure); + throw failure; + } finally { + if (scanStage != null) { + metrics.stop(scanStage, start, scannedBytes.get()); + } + } + } + + private void recordProgress(long current) { + long previous = scannedBytes.get(); + if (current < previous || current > totalScanBytes) { + throw new IllegalStateException( + "External FBIN scan progress is invalid: previous=" + + previous + + ", current=" + + current + + ", total=" + + totalScanBytes); + } + scannedBytes.set(current); + if (requiredScanBytes != totalScanBytes && current >= requiredScanBytes) { + leadReached.complete(current); + } + } + + void runAfterHeadStart(BuildAction buildAction) throws IOException { + Objects.requireNonNull(buildAction, "buildAction"); + if (!started.compareAndSet(false, true)) { + throw new IllegalStateException("External FBIN scan coordinator has already been used"); + } + + Throwable failure = null; + boolean interrupted = false; + boolean buildStarted = false; + long postLeadStart = 0L; + try { + long headStartTimer = CagraHnswBuildMetrics.start(); + try { + leadReached.get(); + } catch (InterruptedException e) { + interrupted = true; + failure = interruptedFailure("Interrupted while awaiting external FBIN scan head start", e); + } catch (ExecutionException e) { + failure = e.getCause(); + } finally { + if (requestedLeadBytes != 0L) { + metrics.stop("external fbin scan head-start [DISK]", headStartTimer, scannedBytes.get()); + } + } + + if (failure == null && Thread.interrupted()) { + interrupted = true; + failure = + interruptedFailure("Interrupted before building with an external FBIN scan", null); + } + + if (failure == null) { + buildStarted = true; + if (requestedLeadBytes != 0L) { + metrics.addCounter("external fbin requested head-start bytes", requestedLeadBytes); + metrics.addCounter("external fbin required scan bytes", requiredScanBytes); + metrics.addCounter("external fbin scan bytes at cagra start", scannedBytes.get()); + postLeadStart = CagraHnswBuildMetrics.start(); + } + try { + buildAction.run(); + } catch (Throwable buildFailure) { + failure = buildFailure; + } finally { + if (requestedLeadBytes != 0L) { + metrics.addCounter("external fbin scan bytes at cagra end", scannedBytes.get()); + } + } + } + + if (Thread.interrupted()) { + interrupted = true; + failure = + combine( + failure, + interruptedFailure("Interrupted while building with an external FBIN scan", null)); + } + if (failure != null) { + completion.cancel(true); + } + + try { + completion.get(); + } catch (InterruptedException e) { + interrupted = true; + completion.cancel(true); + failure = + combine( + failure, + interruptedFailure("Interrupted while awaiting external FBIN validation", e)); + } catch (ExecutionException e) { + failure = combine(failure, e.getCause()); + } catch (CancellationException cancellation) { + if (failure == null) { + failure = cancellation; + } + } + } finally { + if (failure != null) { + completion.cancel(true); + executor.shutdownNow(); + } else { + executor.shutdown(); + } + executor.close(); + if (requestedLeadBytes != 0L && buildStarted) { + metrics.stop("external fbin post-lead overlap wall [GPU+DISK]", postLeadStart); + } + if (interrupted) { + Thread.currentThread().interrupt(); + } + } + + if (failure != null) { + Utils.handleThrowable(failure); + } + } + + private static Throwable combine(Throwable primary, Throwable additional) { + if (primary == null) { + return additional; + } + if (additional != null && additional != primary) { + primary.addSuppressed(additional); + } + return primary; + } + + private static InterruptedIOException interruptedFailure(String message, Throwable cause) { + InterruptedIOException failure = new InterruptedIOException(message); + if (cause != null) { + failure.initCause(cause); + } + return failure; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFloat32Dataset.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFloat32Dataset.java new file mode 100644 index 0000000000..fe356b40cf --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFloat32Dataset.java @@ -0,0 +1,138 @@ +/* + * 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 com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.spi.CuVSProvider; +import java.lang.foreign.MemorySegment; +import java.nio.ByteOrder; +import java.util.Objects; + +/** + * A dense float32 dataset backed by caller-owned native memory. + * + *

The memory is wrapped, not copied. The caller must keep the segment's scope alive until index + * writing has completed. A shared scope is required because the accelerated writer may read the + * dataset concurrently while serializing Lucene's flat-vector file. + * + *

This is an expert-only input for the unsorted, single-segment native-buffering path. Row + * {@code i} must correspond exactly to Lucene vector ordinal and document ID {@code i}; gaps, + * reordered documents, and multiple vector fields are rejected by the writer. The caller is also + * responsible for ensuring that every raw vector component is finite; unlike ordinary Lucene field + * ingestion, this path deliberately does not scan the payload one float at a time. + */ +public final class ExternalFloat32Dataset { + + private final MemorySegment memorySegment; + private final CuVSHostMatrix matrix; + private final int rows; + private final int dimensions; + + private ExternalFloat32Dataset( + MemorySegment memorySegment, CuVSHostMatrix matrix, int rows, int dimensions) { + this.memorySegment = memorySegment; + this.matrix = matrix; + this.rows = rows; + this.dimensions = dimensions; + } + + /** + * Wraps a contiguous native-memory region containing row-major float32 values without copying. + */ + public static ExternalFloat32Dataset fromMemorySegment( + MemorySegment memorySegment, int rows, int dimensions) { + Objects.requireNonNull(memorySegment, "memorySegment"); + if (rows <= 0) { + throw new IllegalArgumentException("rows must be positive"); + } + if (dimensions <= 0) { + throw new IllegalArgumentException("dimensions must be positive"); + } + if (ByteOrder.nativeOrder() != ByteOrder.LITTLE_ENDIAN) { + throw new UnsupportedOperationException( + "External float32 datasets currently require a little-endian host"); + } + if (!memorySegment.isNative()) { + throw new IllegalArgumentException("memorySegment must be backed by native memory"); + } + if (!memorySegment.isReadOnly()) { + throw new IllegalArgumentException( + "memorySegment must be read-only while the graph and flat writers consume it"); + } + if (Math.floorMod(memorySegment.address(), Float.BYTES) != 0) { + throw new IllegalArgumentException("memorySegment address must be float-aligned"); + } + if (!memorySegment.scope().isAlive()) { + throw new IllegalArgumentException("memorySegment scope is not alive"); + } + Thread accessProbe = Thread.ofPlatform().unstarted(() -> {}); + if (!memorySegment.isAccessibleBy(accessProbe)) { + throw new IllegalArgumentException( + "memorySegment must have a shared scope so the flat writer can access it"); + } + + long expectedBytes = + Math.multiplyExact(Math.multiplyExact((long) rows, dimensions), Float.BYTES); + if (memorySegment.byteSize() != expectedBytes) { + throw new IllegalArgumentException( + "memorySegment byte size (" + + memorySegment.byteSize() + + ") must equal rows * dimensions * 4 (" + + expectedBytes + + ")"); + } + + final CuVSMatrix wrapped; + try { + wrapped = + (CuVSMatrix) + CuVSProvider.provider() + .newNativeMatrixBuilder() + .invokeExact(memorySegment, rows, dimensions, CuVSMatrix.DataType.FLOAT); + } catch (Throwable t) { + if (t instanceof Error error) { + throw error; + } + if (t instanceof RuntimeException runtimeException) { + throw runtimeException; + } + throw new IllegalStateException("Unable to wrap the external float32 dataset", t); + } + if (!(wrapped instanceof CuVSHostMatrix hostMatrix)) { + wrapped.close(); + throw new IllegalStateException("Native-memory dataset factory did not return a host matrix"); + } + if (hostMatrix.size() != rows + || hostMatrix.columns() != dimensions + || hostMatrix.dataType() != CuVSMatrix.DataType.FLOAT) { + hostMatrix.close(); + throw new IllegalStateException( + "Native-memory dataset factory returned an unexpected matrix shape or data type"); + } + return new ExternalFloat32Dataset(memorySegment, hostMatrix, rows, dimensions); + } + + public MemorySegment memorySegment() { + return memorySegment; + } + + public CuVSHostMatrix matrix() { + return matrix; + } + + public int rows() { + return rows; + } + + public int dimensions() { + return dimensions; + } + + @Override + public String toString() { + return "ExternalFloat32Dataset[rows=" + rows + ", dimensions=" + dimensions + "]"; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FbinFileMetadata.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FbinFileMetadata.java new file mode 100644 index 0000000000..9ed667197e --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FbinFileMetadata.java @@ -0,0 +1,61 @@ +/* + * 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.EOFException; +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; + +/** Synchronously reads and validates an FBIN header without prefetching any vector payload. */ +record FbinFileMetadata(int rows, int dimensions, long fileBytes) { + + private static final int HEADER_BYTES = 2 * Integer.BYTES; + + static FbinFileMetadata read(Path path) throws IOException { + try (FileChannel channel = FileChannel.open(path, StandardOpenOption.READ)) { + ByteBuffer header = ByteBuffer.allocate(HEADER_BYTES).order(ByteOrder.LITTLE_ENDIAN); + long position = 0L; + while (header.hasRemaining()) { + int read = channel.read(header, position); + if (read < 0) { + throw new EOFException("Unexpected EOF while reading FBIN header from " + path); + } + if (read == 0) { + throw new IOException("Unable to make progress while reading FBIN header from " + path); + } + position += read; + } + header.flip(); + int rows = header.getInt(); + int dimensions = header.getInt(); + if (rows <= 0 || dimensions <= 0) { + throw new IOException( + "FBIN header must contain positive rows and dimensions: " + rows + " x " + dimensions); + } + final long expectedBytes; + try { + expectedBytes = + Math.addExact( + HEADER_BYTES, + Math.multiplyExact(Math.multiplyExact((long) rows, dimensions), Float.BYTES)); + } catch (ArithmeticException overflow) { + throw new IOException("FBIN header shape overflows its file length", overflow); + } + long actualBytes = channel.size(); + if (actualBytes != expectedBytes) { + throw new IOException( + "FBIN length " + + actualBytes + + " does not match header-derived length " + + expectedBytes); + } + return new FbinFileMetadata(rows, dimensions, actualBytes); + } + } +} 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 2007bc5c30..5bcaad4d2c 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 @@ -12,6 +12,7 @@ import com.nvidia.cuvs.RowView; import java.io.IOException; import java.util.ArrayList; +import java.util.Arrays; import java.util.List; import java.util.concurrent.Callable; import java.util.concurrent.ExecutorService; @@ -189,16 +190,20 @@ public NeighborArray getNeighbors(int level, int node) { int[] nodes = layerNodes.get(level - 1); NeighborArray[] neighbors = layerNeighbors.get(level - 1); - // Find the index of this node in the layer - for (int i = 0; i < nodes.length; i++) { - if (nodes[i] == node) { - return neighbors[i]; - } + int ordinal = findUpperLayerOrdinal(nodes, node); + if (ordinal >= 0) { + return neighbors[ordinal]; } } return null; } + /** Returns the ordinal of {@code node} in the sorted upper-layer IDs, or {@code -1}. */ + static int findUpperLayerOrdinal(int[] sortedNodeIds, int node) { + int ordinal = Arrays.binarySearch(sortedNodeIds, node); + return ordinal >= 0 ? ordinal : -1; + } + // Implementation of abstract methods from HnswGraph private int currentNode = -1; private int currentLevel = -1; diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ImmutableExternalFbinDataset.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ImmutableExternalFbinDataset.java new file mode 100644 index 0000000000..54c9ac3a7c --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ImmutableExternalFbinDataset.java @@ -0,0 +1,104 @@ +/* + * 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.lang.foreign.Arena; +import java.lang.foreign.MemorySegment; +import java.nio.channels.FileChannel; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; + +/** + * An owned read-only mapping whose native build bytes and persisted external reference are created + * from the same registered FBIN range. + * + *

Instances are created by {@link ExternalFbinFileRegistry.Registration#map(int, int)}. This + * coupled type prevents accidentally building a graph from one same-shaped native dataset while + * persisting a reference to another. Closing it releases both the cuVS matrix wrapper and mapping; + * callers must keep it open until the index writer has finished. + */ +public final class ImmutableExternalFbinDataset implements AutoCloseable { + + private final ExternalFbinReference reference; + private final Arena arena; + private final ExternalFloat32Dataset dataset; + private boolean closed; + + static ImmutableExternalFbinDataset map(ExternalFbinReference reference) throws IOException { + Path sourcePath = ExternalFbinIO.validateAndResolve(reference); + Arena arena = Arena.ofShared(); + try { + MemorySegment payload; + try (FileChannel channel = FileChannel.open(sourcePath, StandardOpenOption.READ)) { + long actualLength = channel.size(); + if (actualLength != reference.fileLength()) { + throw new IOException( + "External FBIN length changed before mapping: expected " + + reference.fileLength() + + " but got " + + actualLength); + } + payload = + channel.map( + FileChannel.MapMode.READ_ONLY, + reference.payloadOffset(), + reference.payloadLength(), + arena); + } + ExternalFloat32Dataset dataset = + ExternalFloat32Dataset.fromMemorySegment( + payload, reference.rows(), reference.dimensions()); + return new ImmutableExternalFbinDataset(reference, arena, dataset); + } catch (IOException | RuntimeException | Error failure) { + arena.close(); + throw failure; + } + } + + private ImmutableExternalFbinDataset( + ExternalFbinReference reference, Arena arena, ExternalFloat32Dataset dataset) { + this.reference = reference; + this.arena = arena; + this.dataset = dataset; + } + + ExternalFloat32Dataset dataset() { + ensureOpen(); + return dataset; + } + + public ExternalFbinReference reference() { + ensureOpen(); + return reference; + } + + public int rows() { + return reference.rows(); + } + + public int dimensions() { + return reference.dimensions(); + } + + private synchronized void ensureOpen() { + if (closed) { + throw new IllegalStateException("Immutable external FBIN dataset is closed"); + } + } + + @Override + public synchronized void close() { + if (closed) { + return; + } + closed = true; + try { + dataset.matrix().close(); + } finally { + arena.close(); + } + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ImmutableExternalFbinOutput.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ImmutableExternalFbinOutput.java new file mode 100644 index 0000000000..8c34b3d8e5 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ImmutableExternalFbinOutput.java @@ -0,0 +1,69 @@ +/* + * 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 org.apache.lucene.index.DocsWithFieldSet; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.SegmentWriteState; + +/** Persists an FBIN descriptor while validation overlaps GPU graph construction. */ +final class ImmutableExternalFbinOutput implements BorrowedDatasetOutput { + + private final ExternalFbinReferenceWriter referenceOutput; + private final ExternalFbinReference reference; + private final ExternalFbinOptions options; + private final CagraHnswBuildMetrics metrics; + + ImmutableExternalFbinOutput(SegmentWriteState state, BulkIndexingContext context) + throws IOException { + this.referenceOutput = new ExternalFbinReferenceWriter(state); + this.reference = context.reference(); + this.options = context.externalOptions(); + this.metrics = context.metrics(); + } + + @Override + public void writeField( + FieldInfo field, + ExternalFloat32Dataset dataset, + int maxDoc, + DocsWithFieldSet docsWithField, + AcceleratedHnswGraphOutput graphOutput) + throws IOException { + if (docsWithField.cardinality() != maxDoc) { + throw new IllegalStateException( + "Immutable external FBIN requires one vector per document; maxDoc=" + + maxDoc + + ", vectors=" + + docsWithField.cardinality()); + } + referenceOutput.writeField(field, reference); + if (options.validation() == ExternalFbinBuildValidation.TRUSTED_IMMUTABLE) { + graphOutput.writeBorrowedField(field, dataset.matrix()); + return; + } + + long overlapStartedAt = CagraHnswBuildMetrics.start(); + try { + ExternalFbinScanCoordinator coordinator = + ExternalFbinScanCoordinator.start( + reference, options.validation(), options.scanHeadStartBytes(), metrics); + coordinator.runAfterHeadStart(() -> graphOutput.writeBorrowedField(field, dataset.matrix())); + } finally { + metrics.stop("external reference overlap wall [GPU+DISK]", overlapStartedAt); + } + } + + @Override + public void finish() throws IOException { + referenceOutput.finish(); + } + + @Override + public void close() throws IOException { + referenceOutput.close(); + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java index f95487e953..4454dc5ba5 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java @@ -70,8 +70,16 @@ public Lucene101AcceleratedHNSWCodec(AcceleratedHNSWParams acceleratedHNSWParams */ Lucene101AcceleratedHNSWCodec(AcceleratedHNSWParams acceleratedHNSWParams, int numInputVectors) throws Exception { + this( + acceleratedHNSWParams, + BulkIndexingContext.nativeBuffered(numInputVectors, new CagraHnswBuildMetrics())); + } + + /** Bulk-only constructor carrying storage state that generic codec callers cannot configure. */ + Lucene101AcceleratedHNSWCodec( + AcceleratedHNSWParams acceleratedHNSWParams, BulkIndexingContext context) throws Exception { this(NAME, LuceneProvider.getCodec("101")); - initializeFormat(acceleratedHNSWParams, numInputVectors); + initializeFormat(acceleratedHNSWParams, context); } /** @@ -79,7 +87,7 @@ public Lucene101AcceleratedHNSWCodec(AcceleratedHNSWParams acceleratedHNSWParams * with an instance of {@link AcceleratedHNSWParams} with default parameter values. */ private void initializeFormatDefaultValues() { - initializeFormat(new AcceleratedHNSWParams.Builder().build(), 0); + initializeFormat(new AcceleratedHNSWParams.Builder().build(), null); } /** @@ -90,8 +98,17 @@ private void initializeFormatDefaultValues() { * flat buffer (0 = disabled, the default heap-buffered path) */ private void initializeFormat(AcceleratedHNSWParams acceleratedHNSWParams, int numInputVectors) { + initializeFormat( + acceleratedHNSWParams, + numInputVectors == 0 + ? null + : BulkIndexingContext.nativeBuffered(numInputVectors, new CagraHnswBuildMetrics())); + } + + private void initializeFormat( + AcceleratedHNSWParams acceleratedHNSWParams, BulkIndexingContext context) { try { - format = new Lucene99AcceleratedHNSWVectorsFormat(acceleratedHNSWParams, numInputVectors); + format = new Lucene99AcceleratedHNSWVectorsFormat(acceleratedHNSWParams, context); setKnnFormat(format); } catch (LibraryException ex) { log.log( 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 59e012427d..36e73c3c47 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 @@ -15,9 +15,11 @@ import org.apache.lucene.codecs.KnnVectorsWriter; import org.apache.lucene.codecs.hnsw.DefaultFlatVectorScorer; import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; +import org.apache.lucene.codecs.hnsw.FlatVectorsReader; import org.apache.lucene.index.SegmentReadState; import org.apache.lucene.index.SegmentWriteState; import org.apache.lucene.search.TaskExecutor; +import org.apache.lucene.util.IOUtils; /** * cuVS based KnnVectorsFormat for indexing on GPU and searching on the CPU. @@ -31,7 +33,7 @@ public class Lucene99AcceleratedHNSWVectorsFormat extends KnnVectorsFormat { private static final FlatVectorsFormat FLAT_VECTORS_FORMAT; private static final int MAX_DIMENSIONS = 4096; private final AcceleratedHNSWParams acceleratedHNSWParams; - private final int numInputVectors; + private final BulkIndexingContext bulkContext; static final String HNSW_META_CODEC_NAME = "Lucene99HnswVectorsFormatMeta"; static final String HNSW_META_CODEC_EXT = "vem"; @@ -82,9 +84,18 @@ public Lucene99AcceleratedHNSWVectorsFormat(AcceleratedHNSWParams acceleratedHNS */ Lucene99AcceleratedHNSWVectorsFormat( AcceleratedHNSWParams acceleratedHNSWParams, int numInputVectors) { + this( + acceleratedHNSWParams, + numInputVectors == 0 + ? null + : BulkIndexingContext.nativeBuffered(numInputVectors, new CagraHnswBuildMetrics())); + } + + Lucene99AcceleratedHNSWVectorsFormat( + AcceleratedHNSWParams acceleratedHNSWParams, BulkIndexingContext bulkContext) { super("Lucene99AcceleratedHNSWVectorsFormat"); this.acceleratedHNSWParams = acceleratedHNSWParams; - this.numInputVectors = numInputVectors; + this.bulkContext = bulkContext; } /** @@ -92,14 +103,21 @@ public Lucene99AcceleratedHNSWVectorsFormat(AcceleratedHNSWParams acceleratedHNS */ @Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { - boolean nativeMode = isSupported() && numInputVectors > 0; - if (isSupported()) { - if (nativeMode) { + boolean supported = isSupported(); + if (bulkContext != null && !supported) { + throw new IllegalStateException("CAGRA/HNSW bulk indexing requires cuVS GPU support"); + } + if (supported) { + if (bulkContext != null) { + if (bulkContext.storage() != BulkIndexingContext.Storage.NATIVE_BUFFERED) { + log.log(Level.FINE, "cuVS is supported so using the borrowed FBIN bulk writer"); + return new BorrowedDatasetHnswVectorsWriter(state, acceleratedHNSWParams, bulkContext); + } log.log(Level.FINE, "cuVS is supported so using the NativeFlatBufferedHNSWVectorsWriter"); // 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). return new NativeFlatBufferedHNSWVectorsWriter( - state, acceleratedHNSWParams, numInputVectors); + state, acceleratedHNSWParams, bulkContext.exactVectorCount(), bulkContext.metrics()); } log.log(Level.FINE, "cuVS is supported so using the Lucene99AcceleratedHNSWVectorsWriter"); var flatWriter = FLAT_VECTORS_FORMAT.fieldsWriter(state); @@ -128,10 +146,15 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException */ @Override public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { + FlatVectorsReader flatReader = null; try { - return LUCENE_PROVIDER.getLuceneHnswVectorsReaderInstance( - state, FLAT_VECTORS_FORMAT.fieldsReader(state)); + flatReader = + ExternalFbinFlatVectorsReader.hasExternalMarker(state) + ? new ExternalFbinFlatVectorsReader(state) + : FLAT_VECTORS_FORMAT.fieldsReader(state); + return LUCENE_PROVIDER.getLuceneHnswVectorsReaderInstance(state, flatReader); } catch (Exception e) { + IOUtils.closeWhileHandlingException(flatReader); throw Utils.handleThrowable(e); } } 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 cf6d19f7ca..0dea5e0f77 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 @@ -184,7 +184,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throw acceleratedHNSWParams.getHnswLayers(), params, QuantizationType.BINARY, - acceleratedHNSWParams.getWriterThreads()); + acceleratedHNSWParams.getWriterThreads(), + acceleratedHNSWParams.getHnswLayerSeed()); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index 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 aebfbea7c6..d3111dd553 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 @@ -209,7 +209,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE acceleratedHNSWParams.getHnswLayers(), params, QuantizationType.SCALAR, - acceleratedHNSWParams.getWriterThreads()); + acceleratedHNSWParams.getWriterThreads(), + acceleratedHNSWParams.getHnswLayerSeed()); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/MappedFbinDataset.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/MappedFbinDataset.java new file mode 100644 index 0000000000..29a76d44b8 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/MappedFbinDataset.java @@ -0,0 +1,117 @@ +/* + * 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.EOFException; +import java.io.IOException; +import java.lang.foreign.Arena; +import java.lang.foreign.MemorySegment; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.channels.FileChannel; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; + +/** Owned read-only mapping of a structurally validated dense float32 FBIN payload. */ +final class MappedFbinDataset implements AutoCloseable { + + private static final long HEADER_BYTES = 2L * Integer.BYTES; + + private final Arena arena; + private final ExternalFloat32Dataset dataset; + private boolean closed; + + static MappedFbinDataset map(Path path) throws IOException { + int rows; + int dimensions; + long payloadBytes; + Arena arena = Arena.ofShared(); + try { + MemorySegment payload; + try (FileChannel channel = FileChannel.open(path, StandardOpenOption.READ)) { + ByteBuffer header = ByteBuffer.allocate((int) HEADER_BYTES).order(ByteOrder.LITTLE_ENDIAN); + readFully(channel, header, 0L); + header.flip(); + rows = header.getInt(); + dimensions = header.getInt(); + if (rows <= 0 || dimensions <= 0) { + throw new IOException( + "FBIN header must contain positive rows and dimensions: " + + rows + + " x " + + dimensions); + } + payloadBytes = Math.multiplyExact(Math.multiplyExact((long) rows, dimensions), Float.BYTES); + long expectedBytes = Math.addExact(HEADER_BYTES, payloadBytes); + if (channel.size() != expectedBytes) { + throw new IOException( + "FBIN length " + + channel.size() + + " does not match header-derived length " + + expectedBytes); + } + payload = channel.map(FileChannel.MapMode.READ_ONLY, HEADER_BYTES, payloadBytes, arena); + } + ExternalFloat32Dataset dataset = + ExternalFloat32Dataset.fromMemorySegment(payload, rows, dimensions); + return new MappedFbinDataset(arena, dataset); + } catch (IOException | RuntimeException | Error failure) { + arena.close(); + throw failure; + } + } + + private static void readFully(FileChannel channel, ByteBuffer target, long position) + throws IOException { + long current = position; + while (target.hasRemaining()) { + int read = channel.read(target, current); + if (read < 0) { + throw new EOFException("Unexpected EOF while reading FBIN header"); + } + if (read == 0) { + throw new IOException("Unable to make progress while reading FBIN header"); + } + current += read; + } + } + + private MappedFbinDataset(Arena arena, ExternalFloat32Dataset dataset) { + this.arena = arena; + this.dataset = dataset; + } + + ExternalFloat32Dataset dataset() { + ensureOpen(); + return dataset; + } + + int rows() { + return dataset.rows(); + } + + int dimensions() { + return dataset.dimensions(); + } + + private void ensureOpen() { + if (closed) { + throw new IllegalStateException("Mapped FBIN dataset is closed"); + } + } + + @Override + public void close() { + if (closed) { + return; + } + closed = true; + try { + dataset.matrix().close(); + } finally { + arena.close(); + } + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/MappedSelfContainedOutput.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/MappedSelfContainedOutput.java new file mode 100644 index 0000000000..89b3f7822e --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/MappedSelfContainedOutput.java @@ -0,0 +1,118 @@ +/* + * 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.io.InterruptedIOException; +import java.util.concurrent.CancellationException; +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.DocsWithFieldSet; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.SegmentWriteState; + +/** Writes a normal Lucene flat-vector file while graph construction consumes the same mapping. */ +final class MappedSelfContainedOutput implements BorrowedDatasetOutput { + + private final NativeFlatVectorsWriter flatOutput; + private final CagraHnswBuildMetrics metrics; + + MappedSelfContainedOutput(SegmentWriteState state, CagraHnswBuildMetrics metrics) + throws IOException { + this.flatOutput = new NativeFlatVectorsWriter(state, metrics); + this.metrics = metrics; + } + + @Override + public void writeField( + FieldInfo field, + ExternalFloat32Dataset dataset, + int maxDoc, + DocsWithFieldSet docsWithField, + AcceleratedHnswGraphOutput graphOutput) + throws IOException { + long overlapStartedAt = CagraHnswBuildMetrics.start(); + try (ExecutorService executor = + Executors.newSingleThreadExecutor( + task -> Thread.ofPlatform().name("cuvs-mapped-flat-writer").unstarted(task))) { + Future flatWrite = + executor.submit( + () -> { + long startedAt = CagraHnswBuildMetrics.start(); + flatOutput.writeField( + field, dataset.matrix(), dataset.memorySegment(), maxDoc, docsWithField); + metrics.stop( + "mapped flat-write worker [DISK]", + startedAt, + dataset.memorySegment().byteSize()); + return null; + }); + + Throwable failure = null; + try { + graphOutput.writeBorrowedField(field, dataset.matrix()); + } catch (Throwable graphFailure) { + failure = graphFailure; + } + failure = combine(failure, await(flatWrite)); + if (failure != null) { + throw Utils.handleThrowable(failure); + } + } finally { + metrics.stop("mapped flush overlap wall [GPU+DISK]", overlapStartedAt); + } + } + + private static Throwable await(Future future) { + boolean interrupted = false; + Throwable failure = null; + while (true) { + try { + future.get(); + break; + } catch (InterruptedException e) { + interrupted = true; + } catch (ExecutionException e) { + failure = e.getCause(); + break; + } catch (CancellationException e) { + failure = e; + break; + } + } + if (interrupted) { + Thread.currentThread().interrupt(); + InterruptedIOException interruption = + new InterruptedIOException("Interrupted while awaiting mapped flat-vector output"); + if (failure != null) { + interruption.addSuppressed(failure); + } + return interruption; + } + return failure; + } + + private static Throwable combine(Throwable primary, Throwable additional) { + if (primary == null) { + return additional; + } + if (additional != null && additional != primary) { + primary.addSuppressed(additional); + } + return primary; + } + + @Override + public void finish() throws IOException { + flatOutput.finish(); + } + + @Override + public void close() throws IOException { + flatOutput.close(); + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatBufferedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatBufferedHNSWVectorsWriter.java index 4f52d882d7..0a92f8600e 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatBufferedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatBufferedHNSWVectorsWriter.java @@ -10,6 +10,7 @@ import static org.apache.lucene.util.RamUsageEstimator.shallowSizeOfInstance; import com.nvidia.cuvs.CuVSHostMatrix; +import java.io.Closeable; import java.io.IOException; import java.util.ArrayList; import java.util.List; @@ -52,6 +53,15 @@ final class NativeFlatBufferedHNSWVectorsWriter extends KnnVectorsWriter { NativeFlatBufferedHNSWVectorsWriter( SegmentWriteState state, AcceleratedHNSWParams acceleratedHNSWParams, int numInputVectors) throws IOException { + this(state, acceleratedHNSWParams, numInputVectors, new CagraHnswBuildMetrics()); + } + + NativeFlatBufferedHNSWVectorsWriter( + SegmentWriteState state, + AcceleratedHNSWParams acceleratedHNSWParams, + int numInputVectors, + CagraHnswBuildMetrics metrics) + throws IOException { super(); if (numInputVectors <= 0) { throw new IllegalArgumentException("numInputVectors must be > 0, got " + numInputVectors); @@ -67,8 +77,8 @@ final class NativeFlatBufferedHNSWVectorsWriter extends KnnVectorsWriter { try { // In hint mode we own the flat files; the Lucene flat writer must be absent to avoid opening // the same .vec/.vemf outputs. - nativeFlat = new NativeFlatVectorsWriter(state); - graphOutput = new AcceleratedHnswGraphOutput(state, acceleratedHNSWParams); + nativeFlat = new NativeFlatVectorsWriter(state, metrics); + graphOutput = new AcceleratedHnswGraphOutput(state, acceleratedHNSWParams, metrics); success = true; printInfoStream(infoStream, COMPONENT, "NativeFlatBufferedHNSWVectorsWriter is initialized"); } finally { @@ -168,7 +178,14 @@ public void close() throws IOException { try { releaseAllNativeBuffers(); } finally { - IOUtils.close(graphOutput, nativeFlat); + closeOutputsAndResources(graphOutput, nativeFlat); + } + } + + static void closeOutputsAndResources(Closeable... outputs) throws IOException { + try { + IOUtils.close(outputs); + } finally { closeCuVSResourcesInstance(); } } 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 index 4b37686590..3a23c657d4 100644 --- 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 @@ -17,6 +17,7 @@ import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.IndexFileNames; import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.store.DataOutput; import org.apache.lucene.store.IndexOutput; import org.apache.lucene.util.IOUtils; @@ -82,12 +83,27 @@ final class NativeFlatVectorsWriter implements Closeable { // Byte granularity for a single writeBytes call; bounds the transient encode buffer. private static final int CHUNK_BYTES = 1 << 18; // 256 KiB + private static final int RAW_CHUNK_BYTES = + Integer.getInteger("cuvs.rawFlatChunkBytes", 1 << 23); // 8 MiB + + static { + if (RAW_CHUNK_BYTES <= 0) { + throw new IllegalArgumentException("cuvs.rawFlatChunkBytes must be positive"); + } + } private final IndexOutput meta; private final IndexOutput vectorData; + private final CagraHnswBuildMetrics metrics; private boolean finished; NativeFlatVectorsWriter(SegmentWriteState state) throws IOException { + this(state, new CagraHnswBuildMetrics()); + } + + NativeFlatVectorsWriter(SegmentWriteState state, CagraHnswBuildMetrics metrics) + throws IOException { + this.metrics = metrics; String metaFileName = IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, META_EXTENSION); String vectorDataFileName = @@ -127,15 +143,63 @@ final class NativeFlatVectorsWriter implements Closeable { void writeField( FieldInfo field, CuVSHostMatrix matrix, int maxDoc, DocsWithFieldSet docsWithField) throws IOException { + writeField(field, matrix, null, maxDoc, docsWithField); + } + + void writeField( + FieldInfo field, + CuVSHostMatrix matrix, + MemorySegment rawFloat32Data, + 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); + if (rawFloat32Data == null) { + writeFloat32Vectors(matrix, count, dim); + } else { + long expectedBytes = Math.multiplyExact(Math.multiplyExact((long) count, dim), Float.BYTES); + if (rawFloat32Data.byteSize() != expectedBytes) { + throw new IllegalArgumentException( + "Raw float32 byte size must equal vector count * dimension * 4"); + } + copyRawFloat32Vectors(rawFloat32Data, expectedBytes, vectorData, metrics); + } long vectorDataLength = vectorData.getFilePointer() - vectorDataOffset; writeMeta(field, maxDoc, count, vectorDataOffset, vectorDataLength, docsWithField); } + static void copyRawFloat32Vectors( + MemorySegment source, long byteCount, DataOutput output, CagraHnswBuildMetrics metrics) + throws IOException { + if (byteCount < 0L || byteCount > source.byteSize()) { + throw new IllegalArgumentException("Invalid raw float32 byte count: " + byteCount); + } + byte[] chunk = new byte[(int) Math.min(RAW_CHUNK_BYTES, Math.max(1L, byteCount))]; + MemorySegment chunkSegment = MemorySegment.ofArray(chunk); + long offset = 0L; + long stagingCopyNanos = 0L; + long outputWriteNanos = 0L; + long chunks = 0L; + while (offset < byteCount) { + int length = (int) Math.min(chunk.length, byteCount - offset); + long startedAt = CagraHnswBuildMetrics.start(); + MemorySegment.copy(source, offset, chunkSegment, 0L, length); + stagingCopyNanos += System.nanoTime() - startedAt; + + startedAt = CagraHnswBuildMetrics.start(); + output.writeBytes(chunk, 0, length); + outputWriteNanos += System.nanoTime() - startedAt; + offset += length; + chunks++; + } + metrics.record("mapped flat staging-copy [MMAP+CPU]", stagingCopyNanos, byteCount); + metrics.record("mapped flat IndexOutput-write [PAGECACHE]", outputWriteNanos, byteCount); + metrics.addCounter("mapped flat chunks", chunks); + } + private void writeFloat32Vectors(CuVSHostMatrix matrix, int count, int dim) throws IOException { int rowBytes = dim * Float.BYTES; int chunkRows = Math.max(1, CHUNK_BYTES / rowBytes); diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java index 03b5633a87..33a1521054 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java @@ -10,6 +10,7 @@ import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_CUVS_DISTANCE_TYPE; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_GRAPH_DEGREE; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_HNSW_LAYERS; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_HNSW_LAYER_SEED; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_INT_GRAPH_DEGREE; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_MAX_CONN; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_NN_DESCENT_NUM_ITERATIONS; @@ -64,6 +65,15 @@ public void testAcceleratedHNSWParamsDefaultValues() { assertEquals(DEFAULT_STRATEGY, params.getStrategy()); assertEquals(DEFAULT_CUVS_DISTANCE_TYPE, params.getCuvsDistanceType()); assertEquals(DEFAULT_NN_DESCENT_NUM_ITERATIONS, params.getNNDescentNumIterations()); + assertEquals(DEFAULT_HNSW_LAYER_SEED, params.getHnswLayerSeed()); + } + + @Test + public void testAcceleratedHNSWLayerSeed() { + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder().withHnswLayerSeed(-123456789L).build(); + + assertEquals(-123456789L, params.getHnswLayerSeed()); } @Test diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWUtils.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWUtils.java new file mode 100644 index 0000000000..7e75bd4059 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWUtils.java @@ -0,0 +1,61 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.nvidia.cuvs.lucene; + +import java.util.Arrays; +import java.util.List; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.junit.Test; + +public class TestAcceleratedHNSWUtils extends LuceneTestCase { + + @Test + public void testHnswMUsesCeilingAndNeverReturnsZero() { + assertEquals(1, AcceleratedHNSWUtils.hnswM(0)); + assertEquals(1, AcceleratedHNSWUtils.hnswM(1)); + assertEquals(1, AcceleratedHNSWUtils.hnswM(2)); + assertEquals(2, AcceleratedHNSWUtils.hnswM(3)); + assertEquals(32, AcceleratedHNSWUtils.hnswM(63)); + assertEquals(32, AcceleratedHNSWUtils.hnswM(64)); + assertEquals(33, AcceleratedHNSWUtils.hnswM(65)); + } + + @Test + public void testUpperLayerSizesUseCeilingAndStopNaturally() { + List layers = AcceleratedHNSWUtils.selectUpperLayerNodes(10_000, 99, 32, 44L); + assertEquals(2, layers.size()); + assertEquals(313, layers.get(0).length); + assertEquals(10, layers.get(1).length); + + List capped = AcceleratedHNSWUtils.selectUpperLayerNodes(10_000, 2, 32, 44L); + assertEquals(1, capped.size()); + assertEquals(313, capped.get(0).length); + } + + @Test + public void testUpperLayerSamplingIsDeterministicAndNested() { + List first = AcceleratedHNSWUtils.selectUpperLayerNodes(10_000, 99, 32, 987654L); + List repeat = AcceleratedHNSWUtils.selectUpperLayerNodes(10_000, 99, 32, 987654L); + List differentSeed = AcceleratedHNSWUtils.selectUpperLayerNodes(10_000, 99, 32, 987655L); + + assertEquals(first.size(), repeat.size()); + for (int i = 0; i < first.size(); i++) { + assertArrayEquals(first.get(i), repeat.get(i)); + } + assertFalse(Arrays.equals(first.get(0), differentSeed.get(0))); + + for (int node : first.get(1)) { + assertTrue(Arrays.binarySearch(first.get(0), node) >= 0); + } + } + + @Test + public void testUpperLayersRequireStrictShrink() { + assertTrue(AcceleratedHNSWUtils.selectUpperLayerNodes(1_000, 99, 1, 44L).isEmpty()); + assertTrue(AcceleratedHNSWUtils.selectUpperLayerNodes(32, 99, 32, 44L).isEmpty()); + assertTrue(AcceleratedHNSWUtils.selectUpperLayerNodes(10_000, 1, 32, 44L).isEmpty()); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBulkExternalFbinApiSurface.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBulkExternalFbinApiSurface.java new file mode 100644 index 0000000000..f16670709c --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBulkExternalFbinApiSurface.java @@ -0,0 +1,62 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static java.lang.reflect.Modifier.isPublic; + +import java.lang.reflect.Method; +import java.nio.file.Path; +import java.util.Arrays; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.junit.Test; + +/** Keeps external-FBIN build controls on the owned bulk API rather than the generic codec API. */ +public class TestBulkExternalFbinApiSurface extends LuceneTestCase { + + @Test + public void testGenericAcceleratedParamsDoNotExposeBulkStorageControls() { + for (Method method : AcceleratedHNSWParams.Builder.class.getMethods()) { + String name = method.getName().toLowerCase(); + assertFalse("generic params must not expose external FBIN: " + method, name.contains("fbin")); + assertFalse( + "generic params must not expose external datasets: " + method, name.contains("external")); + assertFalse( + "generic params must not expose an exact bulk count: " + method, + name.contains("numinputvectors")); + } + } + + @Test + public void testExternalBuildIsOnlyExposedByBulkWriter() throws Exception { + Method external = + CagraHnswBulkIndexWriter.class.getMethod( + "indexImmutableFbin", + ExternalFbinFileRegistry.Registration.class, + CagraHnswBulkIndexWriter.Config.class, + ExternalFbinOptions.class); + Method mapped = + CagraHnswBulkIndexWriter.class.getMethod( + "indexMappedFbin", Path.class, CagraHnswBulkIndexWriter.Config.class); + + assertTrue(isPublic(external.getModifiers())); + assertTrue(isPublic(mapped.getModifiers())); + assertFalse( + Arrays.stream(Lucene101AcceleratedHNSWCodec.class.getConstructors()) + .flatMap(constructor -> Arrays.stream(constructor.getParameterTypes())) + .anyMatch( + type -> + type == ExternalFbinOptions.class + || type == ImmutableExternalFbinDataset.class + || type == ExternalFloat32Dataset.class)); + } + + @Test + public void testExternalOptionsValidateHeadStart() { + assertThrows( + IllegalArgumentException.class, + () -> new ExternalFbinOptions(ExternalFbinBuildValidation.VERIFY_SHA256, -1L)); + assertThrows(NullPointerException.class, () -> new ExternalFbinOptions(null, 0L)); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBulkFbinIndexStorage.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBulkFbinIndexStorage.java new file mode 100644 index 0000000000..1220c244eb --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBulkFbinIndexStorage.java @@ -0,0 +1,232 @@ +/* + * 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.index.VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT; + +import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; +import com.nvidia.cuvs.spi.CuVSProvider; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.channels.FileChannel; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.security.MessageDigest; +import java.util.Arrays; +import java.util.HexFormat; +import java.util.Random; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +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.SuppressFileSystems; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.junit.After; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Test; + +/** End-to-end storage and lifecycle coverage for the two borrowed-FBIN bulk build modes. */ +@SuppressFileSystems("*") +@SuppressSysoutChecks(bugUrl = "") +public class TestBulkFbinIndexStorage extends LuceneTestCase { + + private static final String ID_FIELD = "id"; + private static final String VECTOR_FIELD = "vector"; + private static final int ROWS = 300; + private static final int DIMENSIONS = 32; + + @BeforeClass + public static void beforeClass() { + try { + CuVSProvider.provider().enableRMMAsyncMemory(); + } catch (UnsupportedOperationException unsupported) { + assumeTrue("cuVS not supported: " + unsupported.getMessage(), false); + } + } + + @Before + public void requireCuvs() { + assumeTrue("cuVS not supported", isSupported()); + } + + @After + public void clearRegistry() { + ExternalFbinFileRegistry.clearForTests(); + } + + @Test + public void testMappedFbinWritesSelfContainedMIPIndex() throws Exception { + Path root = createTempDir("mapped-fbin-index"); + Path fbin = writeFbin(root, createMIPVectors()); + Path index = root.resolve("index"); + + CagraHnswBulkIndexWriter.indexMappedFbin(fbin, config(index)); + + assertContainsExtension(index, ".vec"); + assertContainsExtension(index, ".vemf"); + assertOmitsExtension(index, ".vefr"); + Files.delete(fbin); + assertMIPIndexSearchable(index); + } + + @Test + public void testImmutableExternalFbinRequiresRegistrationAndDetectsMutation() throws Exception { + Path root = createTempDir("external-fbin-index"); + float[][] vectors = createMIPVectors(); + Path fbin = writeFbin(root, vectors); + Path index = root.resolve("index"); + String digest = sha256(fbin); + + try (ExternalFbinFileRegistry.Registration registration = + ExternalFbinFileRegistry.register(fbin, digest)) { + CagraHnswBulkIndexWriter.indexImmutableFbin(registration, config(index), trustedImmutable()); + + assertContainsExtension(index, ".vefr"); + assertContainsExtension(index, ".vex"); + assertOmitsExtension(index, ".vec"); + assertOmitsExtension(index, ".vemf"); + assertMIPIndexSearchable(index); + + try (Directory directory = FSDirectory.open(index); + DirectoryReader reader = DirectoryReader.open(directory)) { + FloatVectorValues values = getOnlyLeafReader(reader).getFloatVectorValues(VECTOR_FIELD); + assertArrayEquals(vectors[0], values.vectorValue(0), 0.0f); + getOnlyLeafReader(reader).checkIntegrity(); + } + } + + try (Directory directory = FSDirectory.open(index)) { + IOException missingRegistration = + expectThrows(IOException.class, () -> DirectoryReader.open(directory)); + assertTrue( + missingRegistration.getMessage(), + missingRegistration.getMessage().contains("No allowlisted external FBIN")); + } + + corruptPayload(fbin); + try (ExternalFbinFileRegistry.Registration ignored = + ExternalFbinFileRegistry.register(fbin, digest); + Directory directory = FSDirectory.open(index); + DirectoryReader reader = DirectoryReader.open(directory)) { + IOException corruptSource = + expectThrows(IOException.class, () -> getOnlyLeafReader(reader).checkIntegrity()); + assertTrue( + corruptSource.getMessage(), corruptSource.getMessage().contains("SHA-256 mismatch")); + } + } + + @Test + public void testDigestMismatchDoesNotPublishCommit() throws Exception { + Path root = createTempDir("external-fbin-validation"); + Path fbin = writeFbin(root, createMIPVectors()); + Path index = root.resolve("index"); + String digest = sha256(fbin); + corruptPayload(fbin); + + try (ExternalFbinFileRegistry.Registration registration = + ExternalFbinFileRegistry.register(fbin, digest)) { + IOException mismatch = + expectThrows( + IOException.class, + () -> + CagraHnswBulkIndexWriter.indexImmutableFbin( + registration, config(index), ExternalFbinOptions.verifySha256())); + assertTrue(mismatch.getMessage(), mismatch.getMessage().contains("SHA-256 mismatch")); + } + + try (Directory directory = FSDirectory.open(index)) { + assertFalse( + "A failed external-FBIN verification must not publish a commit", + DirectoryReader.indexExists(directory)); + } + } + + private static CagraHnswBulkIndexWriter.Config config(Path index) { + AcceleratedHNSWParams graphBuild = + new AcceleratedHNSWParams.Builder() + .withCuvsDistanceType(CuvsDistanceType.InnerProduct) + .build(); + return CagraHnswBulkIndexWriter.Config.builder() + .field(VECTOR_FIELD, DIMENSIONS, MAXIMUM_INNER_PRODUCT) + .idField(ID_FIELD) + .graphBuild(graphBuild) + .segments(1, false) + .targetDirectory(index) + .build(); + } + + private static ExternalFbinOptions trustedImmutable() { + return new ExternalFbinOptions(ExternalFbinBuildValidation.TRUSTED_IMMUTABLE, 0L); + } + + private static void assertMIPIndexSearchable(Path index) throws Exception { + try (Directory directory = FSDirectory.open(index); + DirectoryReader reader = DirectoryReader.open(directory)) { + assertEquals(1, reader.leaves().size()); + float[] query = new float[DIMENSIONS]; + query[0] = 1.0f; + IndexSearcher searcher = new IndexSearcher(reader); + TopDocs hits = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, query, 10), 10); + + assertEquals(10, hits.scoreDocs.length); + assertEquals(0, hits.scoreDocs[0].doc); + assertEquals(101.0f, hits.scoreDocs[0].score, 0.0f); + assertEquals("0", searcher.storedFields().document(hits.scoreDocs[0].doc).get(ID_FIELD)); + } + } + + private static float[][] createMIPVectors() { + float[][] vectors = new float[ROWS][DIMENSIONS]; + vectors[0][0] = 100.0f; + Random random = new Random(8675309L); + for (int row = 1; row < vectors.length; row++) { + for (int dimension = 0; dimension < DIMENSIONS; dimension++) { + vectors[row][dimension] = random.nextFloat() * 2.0f - 1.0f; + } + } + return vectors; + } + + private static Path writeFbin(Path root, float[][] vectors) throws IOException { + Path fbin = root.resolve("vectors.fbin"); + TestUtils.writeFbin(fbin, vectors); + return fbin; + } + + private static String sha256(Path file) throws Exception { + MessageDigest digest = MessageDigest.getInstance("SHA-256"); + digest.update(Files.readAllBytes(file)); + return HexFormat.of().formatHex(digest.digest()); + } + + private static void corruptPayload(Path fbin) throws IOException { + try (FileChannel channel = FileChannel.open(fbin, StandardOpenOption.WRITE)) { + channel.write(ByteBuffer.wrap(new byte[] {42}), ExternalFbinReference.HEADER_BYTES + 7L); + } + } + + private static void assertContainsExtension(Path index, String extension) throws IOException { + try (Directory directory = FSDirectory.open(index)) { + assertTrue( + "Expected " + extension + " in " + index, + Arrays.stream(directory.listAll()).anyMatch(name -> name.endsWith(extension))); + } + } + + private static void assertOmitsExtension(Path index, String extension) throws IOException { + try (Directory directory = FSDirectory.open(index)) { + assertFalse( + "Did not expect " + extension + " in " + index, + Arrays.stream(directory.listAll()).anyMatch(name -> name.endsWith(extension))); + } + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBuildMetricsConcurrency.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBuildMetricsConcurrency.java new file mode 100644 index 0000000000..8437e69dff --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBuildMetricsConcurrency.java @@ -0,0 +1,97 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.atomic.AtomicBoolean; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.junit.Test; + +public class TestCagraHnswBuildMetricsConcurrency extends LuceneTestCase { + + @Test + public void testGaugeAllowsRepeatedIdenticalValues() { + CagraHnswBuildMetrics metrics = new CagraHnswBuildMetrics(); + + metrics.setGauge("effective graph degree", 32L); + metrics.setGauge("effective graph degree", 32L); + + assertEquals(32L, metrics.snapshot().get("gauge/effective graph degree").longValue()); + } + + @Test + public void testGaugeRejectsConflictingValues() { + CagraHnswBuildMetrics metrics = new CagraHnswBuildMetrics(); + metrics.setGauge("effective graph degree", 32L); + + IllegalStateException failure = + expectThrows( + IllegalStateException.class, () -> metrics.setGauge("effective graph degree", 56L)); + + assertTrue(failure.getMessage(), failure.getMessage().contains("effective graph degree")); + assertTrue(failure.getMessage(), failure.getMessage().contains("32")); + assertTrue(failure.getMessage(), failure.getMessage().contains("56")); + assertEquals(32L, metrics.snapshot().get("gauge/effective graph degree").longValue()); + } + + @Test + public void testCounterRemainsAdditive() { + CagraHnswBuildMetrics metrics = new CagraHnswBuildMetrics(); + + metrics.addCounter("logical cagra adjacency bytes", 4_096L); + metrics.addCounter("logical cagra adjacency bytes", 8_192L); + + assertEquals( + 12_288L, metrics.snapshot().get("counter/logical cagra adjacency bytes").longValue()); + } + + @Test + public void testSnapshotWhileNewStagesAreRecorded() throws Exception { + CagraHnswBuildMetrics metrics = new CagraHnswBuildMetrics(); + CountDownLatch start = new CountDownLatch(1); + AtomicBoolean recording = new AtomicBoolean(true); + + try (ExecutorService executor = Executors.newFixedThreadPool(2)) { + Future recorder = + executor.submit( + () -> { + await(start); + try { + for (int i = 0; i < 20_000; i++) { + metrics.record("stage-" + i, i, i + 1L); + } + } finally { + recording.set(false); + } + }); + Future snapshotter = + executor.submit( + () -> { + await(start); + while (recording.get()) { + metrics.snapshot(); + } + }); + + start.countDown(); + recorder.get(); + snapshotter.get(); + } + + assertEquals(20_000, metrics.snapshot().size() / 3); + } + + private static void await(CountDownLatch latch) { + try { + latch.await(); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + throw new AssertionError("interrupted while starting concurrency test", interrupted); + } + } +} 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 index 2a299a6440..204039ff7d 100644 --- 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 @@ -8,11 +8,13 @@ import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; +import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; import com.nvidia.cuvs.spi.CuVSProvider; import java.io.File; import java.io.IOException; import java.nio.file.Path; import java.nio.file.Paths; +import java.util.Map; import java.util.Random; import java.util.UUID; import org.apache.commons.io.FileUtils; @@ -96,9 +98,13 @@ public void testSingleSegmentBuildIsSearchable() throws Exception { float[][] dataset = generateDataset(random, numDocs, dimension); TestUtils.writeFbin(fbinPath, dataset); - CagraHnswBulkIndexWriter.indexFbin(fbinPath, configFor(dimension, 1, false)); + CagraHnswBulkIndexWriter.Config config = configFor(dimension, 1, false); + CagraHnswBulkIndexWriter.indexFbin(fbinPath, config); assertSearchable(numDocs, dataset, /* expectedSegments= */ 1); + Map metrics = config.metrics().snapshot(); + assertEquals(1L, metrics.get("stage/bulk writer commit wall [CPU+GPU+DISK]/count")); + assertEquals(1L, metrics.get("stage/bulk writer close [DISK]/count")); } @Test @@ -347,6 +353,24 @@ public void testConstructorRejectsIndexSort() throws Exception { } } + @Test + public void testConfigRejectsGraphMetricThatDisagreesWithField() { + AcceleratedHNSWParams innerProductGraph = + new AcceleratedHNSWParams.Builder() + .withCuvsDistanceType(CuvsDistanceType.InnerProduct) + .build(); + + IllegalArgumentException failure = + expectThrows( + IllegalArgumentException.class, + () -> + CagraHnswBulkIndexWriter.Config.builder() + .field(VECTOR_FIELD, 8, EUCLIDEAN) + .graphBuild(innerProductGraph) + .build()); + assertTrue(failure.getMessage().contains("does not match")); + } + private CagraHnswBulkIndexWriter.Config configFor( int dimension, int numSegments, boolean overlapped) { return CagraHnswBulkIndexWriter.Config.builder() diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBulkManualFieldContract.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBulkManualFieldContract.java new file mode 100644 index 0000000000..b5894b5088 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBulkManualFieldContract.java @@ -0,0 +1,91 @@ +/* + * 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.apache.lucene.index.VectorSimilarityFunction.COSINE; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +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.IndexWriterConfig; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.junit.Test; + +/** Fail-closed field-contract tests for the public manual bulk writer API. */ +public class TestCagraHnswBulkManualFieldContract extends LuceneTestCase { + + private static final String VECTOR_FIELD = "vector"; + private static final int DIMENSIONS = 8; + + @Test + public void testRejectsWrongVectorFieldName() throws Exception { + Document document = vectorDocument("unexpected", DIMENSIONS, EUCLIDEAN); + + assertRejected(document, "field name"); + } + + @Test + public void testRejectsWrongVectorDimension() throws Exception { + Document document = vectorDocument(VECTOR_FIELD, DIMENSIONS - 1, EUCLIDEAN); + + assertRejected(document, "dimension"); + } + + @Test + public void testRejectsWrongVectorSimilarity() throws Exception { + Document document = vectorDocument(VECTOR_FIELD, DIMENSIONS, COSINE); + + assertRejected(document, "similarity"); + } + + @Test + public void testRejectsDocumentWithoutVector() throws Exception { + Document document = new Document(); + document.add(new StringField("id", "0", Field.Store.YES)); + + assertRejected(document, "exactly one vector field"); + } + + @Test + public void testRejectsDocumentWithMoreThanOneVector() throws Exception { + Document document = vectorDocument(VECTOR_FIELD, DIMENSIONS, EUCLIDEAN); + document.add(new KnnFloatVectorField(VECTOR_FIELD, new float[DIMENSIONS], EUCLIDEAN)); + + assertRejected(document, "exactly one vector field"); + } + + private static Document vectorDocument( + String fieldName, + int dimensions, + org.apache.lucene.index.VectorSimilarityFunction similarity) { + Document document = new Document(); + document.add(new KnnFloatVectorField(fieldName, new float[dimensions], similarity)); + return document; + } + + private static void assertRejected(Document document, String expectedMessage) throws Exception { + try (Directory directory = newDirectory(); + CagraHnswBulkIndexWriter writer = + new CagraHnswBulkIndexWriter(directory, new IndexWriterConfig(), config(), 1)) { + IllegalArgumentException failure; + try { + failure = expectThrows(IllegalArgumentException.class, () -> writer.addDocument(document)); + } finally { + writer.abort(); + } + assertTrue(failure.getMessage(), failure.getMessage().contains(expectedMessage)); + } + } + + private static CagraHnswBulkIndexWriter.Config config() { + return CagraHnswBulkIndexWriter.Config.builder() + .field(VECTOR_FIELD, DIMENSIONS, EUCLIDEAN) + .graphBuild(new AcceleratedHNSWParams.Builder().build()) + .build(); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVS2510GPUVectorsReaderScoreNormalization.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVS2510GPUVectorsReaderScoreNormalization.java new file mode 100644 index 0000000000..fafc58d1b2 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVS2510GPUVectorsReaderScoreNormalization.java @@ -0,0 +1,41 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.nvidia.cuvs.lucene; + +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.junit.Test; + +public class TestCuVS2510GPUVectorsReaderScoreNormalization extends LuceneTestCase { + + @Test + public void testMaximumInnerProductValueUsesLuceneScoreContract() { + assertEquals( + 13.0f, + CuVS2510GPUVectorsReader.toLuceneScore( + VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT, 12.0f), + 0.0f); + assertEquals( + 0.2f, + CuVS2510GPUVectorsReader.toLuceneScore( + VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT, -4.0f), + 0.0f); + } + + @Test + public void testOtherCuVSValuesUseTheirLuceneScoreContracts() { + assertEquals( + 0.2f, + CuVS2510GPUVectorsReader.toLuceneScore(VectorSimilarityFunction.EUCLIDEAN, 4.0f), + 0.0f); + assertEquals( + 0.75f, + CuVS2510GPUVectorsReader.toLuceneScore(VectorSimilarityFunction.DOT_PRODUCT, 0.5f), + 0.0f); + assertEquals( + 0.75f, CuVS2510GPUVectorsReader.toLuceneScore(VectorSimilarityFunction.COSINE, 0.5f), 0.0f); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestExternalFbinFlatVectorsReader.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestExternalFbinFlatVectorsReader.java new file mode 100644 index 0000000000..fec0704bbf --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestExternalFbinFlatVectorsReader.java @@ -0,0 +1,516 @@ +/* + * 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.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +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.StandardOpenOption; +import java.security.MessageDigest; +import java.util.HexFormat; +import java.util.Map; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.CodecUtil; +import org.apache.lucene.index.CorruptIndexException; +import org.apache.lucene.index.DocValuesSkipIndexType; +import org.apache.lucene.index.DocValuesType; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FieldInfos; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexFileNames; +import org.apache.lucene.index.IndexOptions; +import org.apache.lucene.index.SegmentInfo; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.VectorEncoding; +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.util.InfoStream; +import org.apache.lucene.util.StringHelper; +import org.apache.lucene.util.Version; +import org.apache.lucene.util.hnsw.RandomVectorScorer; +import org.junit.After; + +public class TestExternalFbinFlatVectorsReader extends LuceneTestCase { + + private static final int[] GENERALIZATION_DIMENSIONS = {64, 65, 128, 129, 192, 193, 1536, 2048}; + private static final long FORCED_MMAP_CHUNK_SIZE = 32L; + + @After + public void clearRegistry() { + ExternalFbinFileRegistry.clearForTests(); + } + + public void testDescriptorVectorValuesScoringAndIntegrity() throws Exception { + float[][] vectors = {{1.0f, 2.0f, 3.0f}, {4.0f, 5.0f, 6.0f}, {7.0f, 8.0f, 9.0f}}; + Path fbin = writeFbin(createTempDir().resolve("vectors.fbin"), vectors); + String digest = sha256(fbin); + + try (Directory indexDirectory = newDirectory(); + ExternalFbinFileRegistry.Registration registration = + ExternalFbinFileRegistry.register(fbin, digest)) { + ExternalFbinReference reference = registration.reference(0, vectors.length); + SegmentFixture fixture = writeDescriptor(indexDirectory, reference); + + try (ExternalFbinFlatVectorsReader reader = + new ExternalFbinFlatVectorsReader(fixture.readState(), 16L)) { + FloatVectorValues values = reader.getFloatVectorValues("vector"); + assertEquals(vectors.length, values.size()); + assertEquals(vectors[0].length, values.dimension()); + assertArrayEquals(vectors[0], values.vectorValue(0), 0.0f); + assertArrayEquals(vectors[2], values.copy().vectorValue(2), 0.0f); + + RandomVectorScorer scorer = reader.getRandomVectorScorer("vector", vectors[1]); + assertEquals(1.0f, scorer.score(1), 0.0f); + assertTrue(scorer.score(0) < scorer.score(1)); + assertEquals(1, scorer.ordToDoc(1)); + reader.checkIntegrity(); + reader.checkIntegrity(); + assertArrayEquals(vectors[2], reader.getFloatVectorValues("vector").vectorValue(2), 0.0f); + + try (FileChannel channel = FileChannel.open(fbin, StandardOpenOption.WRITE)) { + channel.write(ByteBuffer.wrap(new byte[] {17}), ExternalFbinReference.HEADER_BYTES + 3); + } + IOException failure = expectThrows(IOException.class, reader::checkIntegrity); + assertTrue(failure.getMessage().contains("SHA-256 mismatch")); + } + } + } + + public void testCorruptDescriptorChecksumIsRejected() throws Exception { + float[][] vectors = {{1.0f, 2.0f}, {3.0f, 4.0f}}; + Path fbin = writeFbin(createTempDir().resolve("vectors.fbin"), vectors); + + try (Directory indexDirectory = newDirectory(); + ExternalFbinFileRegistry.Registration registration = + ExternalFbinFileRegistry.register(fbin, sha256(fbin))) { + SegmentFixture fixture = + writeDescriptor(indexDirectory, registration.reference(0, vectors.length)); + corruptDescriptorPayload(indexDirectory, fixture.readState()); + + expectThrows( + CorruptIndexException.class, + () -> new ExternalFbinFlatVectorsReader(fixture.readState())); + } + } + + public void testDimensionBoundariesWithSlicedRowsAcrossMmapChunks() throws Exception { + for (int dimensions : GENERALIZATION_DIMENSIONS) { + float[][] vectors = createVectors(5, dimensions); + Path fbin = writeFbin(createTempDir().resolve("vectors-" + dimensions + "d.fbin"), vectors); + + try (Directory indexDirectory = newDirectory(); + ExternalFbinFileRegistry.Registration registration = + ExternalFbinFileRegistry.register(fbin, sha256(fbin))) { + ExternalFbinReference reference = registration.reference(1, 3); + SegmentFixture fixture = writeDescriptor(indexDirectory, reference); + + try (ExternalFbinFlatVectorsReader reader = + new ExternalFbinFlatVectorsReader(fixture.readState(), FORCED_MMAP_CHUNK_SIZE)) { + FloatVectorValues values = reader.getFloatVectorValues("vector"); + assertEquals(3, values.size()); + assertEquals(dimensions, values.dimension()); + assertArrayEquals(vectors[1], values.vectorValue(0), 0.0f); + assertArrayEquals(vectors[3], values.copy().vectorValue(2), 0.0f); + + RandomVectorScorer scorer = reader.getRandomVectorScorer("vector", vectors[2]); + assertEquals(1.0f, scorer.score(1), 0.0f); + assertTrue(scorer.score(0) < scorer.score(1)); + reader.checkIntegrity(); + } + } + } + } + + public void testSparseFileReadsFinalRowAcrossTwoGiBBoundary() throws Exception { + int dimensions = 64; + long rowBytes = (long) dimensions * Float.BYTES; + long twoGiB = 1L << 31; + int fileRows = Math.toIntExact((twoGiB - ExternalFbinReference.HEADER_BYTES) / rowBytes + 1L); + int finalRow = fileRows - 1; + long finalRowOffset = + ExternalFbinReference.HEADER_BYTES + Math.multiplyExact((long) finalRow, rowBytes); + assertTrue(finalRowOffset < twoGiB); + assertTrue(finalRowOffset + rowBytes > twoGiB); + + float[] expected = createVectors(1, dimensions)[0]; + Path fbin = createTempDir().resolve("sparse-over-2gib.fbin"); + writeSparseFbin(fbin, fileRows, dimensions, finalRowOffset, expected); + assertTrue(Files.size(fbin) > twoGiB); + + // Offset behavior is under test here. A synthetic content ID avoids reading the 2 GiB hole. + String syntheticContentId = "a5".repeat(ExternalFbinReference.SHA256_BYTES); + try (Directory indexDirectory = newDirectory(); + ExternalFbinFileRegistry.Registration registration = + ExternalFbinFileRegistry.register(fbin, syntheticContentId)) { + SegmentFixture fixture = writeDescriptor(indexDirectory, registration.reference(finalRow, 1)); + + try (ExternalFbinFlatVectorsReader reader = + new ExternalFbinFlatVectorsReader(fixture.readState(), 1L << 30)) { + FloatVectorValues values = reader.getFloatVectorValues("vector"); + assertEquals(1, values.size()); + assertEquals(dimensions, values.dimension()); + assertArrayEquals(expected, values.vectorValue(0), 0.0f); + } + } + } + + public void testSparseHundredMillionBy2048SliceUsesLongOffsetsAcrossMmapChunks() + throws Exception { + int fileRows = 100_000_000; + int dimensions = 2048; + long rowBytes = Math.multiplyExact((long) dimensions, Float.BYTES); + long payloadBytes = Math.multiplyExact((long) fileRows, rowBytes); + long fileLength = Math.addExact(ExternalFbinReference.HEADER_BYTES, payloadBytes); + long mmapChunkSize = 1L << 30; + assertEquals(819_200_000_000L, payloadBytes); + assertEquals(819_200_000_008L, fileLength); + + long finalChunkStart = Math.floorDiv(fileLength - 1L, mmapChunkSize) * mmapChunkSize; + long crossedChunkBoundary = finalChunkStart - 2L * mmapChunkSize; + int firstRow = + Math.toIntExact( + Math.floorDiv(crossedChunkBoundary - ExternalFbinReference.HEADER_BYTES, rowBytes)); + int rowCount = Math.toIntExact((long) fileRows - firstRow); + int finalRow = fileRows - 1; + long firstRowOffset = + Math.addExact( + ExternalFbinReference.HEADER_BYTES, Math.multiplyExact((long) firstRow, rowBytes)); + long finalRowOffset = + Math.addExact( + ExternalFbinReference.HEADER_BYTES, Math.multiplyExact((long) finalRow, rowBytes)); + assertTrue(firstRowOffset < crossedChunkBoundary); + assertTrue(firstRowOffset + rowBytes > crossedChunkBoundary); + assertTrue(Math.multiplyExact((long) rowCount, rowBytes) > Integer.MAX_VALUE); + + float[] firstExpected = createVector(dimensions, 0.25f); + float[] finalExpected = createVector(dimensions, -0.5f); + Path fbin = createTempDir().resolve("sparse-100m-2048d.fbin"); + writeSparseFbinRows( + fbin, + fileRows, + dimensions, + new SparseRow(firstRowOffset, firstExpected), + new SparseRow(finalRowOffset, finalExpected)); + assertEquals(fileLength, Files.size(fbin)); + + // The synthetic identity deliberately avoids hashing or otherwise reading the sparse hole. + String syntheticContentId = "c7".repeat(ExternalFbinReference.SHA256_BYTES); + try (Directory indexDirectory = newDirectory(); + ExternalFbinFileRegistry.Registration registration = + ExternalFbinFileRegistry.register(fbin, syntheticContentId)) { + ExternalFbinReference reference = registration.reference(firstRow, rowCount); + assertEquals(fileLength, reference.fileLength()); + assertEquals(firstRow, reference.firstRow()); + assertEquals(rowCount, reference.rows()); + assertEquals(dimensions, reference.dimensions()); + assertEquals(firstRowOffset, reference.payloadOffset()); + assertEquals(Math.multiplyExact((long) rowCount, rowBytes), reference.payloadLength()); + SegmentFixture fixture = writeDescriptor(indexDirectory, reference); + + try (ExternalFbinFlatVectorsReader reader = + new ExternalFbinFlatVectorsReader(fixture.readState(), mmapChunkSize)) { + FloatVectorValues values = reader.getFloatVectorValues("vector"); + assertEquals(rowCount, values.size()); + assertEquals(dimensions, values.dimension()); + assertArrayEquals(firstExpected, values.vectorValue(0), 0.0f); + + int finalOrdinal = rowCount - 1; + long finalComponentOffset = + Math.addExact( + Math.multiplyExact((long) finalOrdinal, rowBytes), + Math.multiplyExact((long) dimensions - 1L, Float.BYTES)); + assertTrue(finalComponentOffset > Integer.MAX_VALUE); + float[] finalActual = values.copy().vectorValue(finalOrdinal); + assertEquals(finalExpected[0], finalActual[0], 0.0f); + assertEquals(finalExpected[dimensions - 1], finalActual[dimensions - 1], 0.0f); + assertArrayEquals(finalExpected, finalActual, 0.0f); + } + } + } + + public void testReaderRequiresRegisteredContentId() throws Exception { + float[][] vectors = {{1.0f, 2.0f}, {3.0f, 4.0f}}; + Path fbin = writeFbin(createTempDir().resolve("vectors.fbin"), vectors); + String digest = sha256(fbin); + + try (Directory indexDirectory = newDirectory()) { + SegmentFixture fixture; + try (ExternalFbinFileRegistry.Registration registration = + ExternalFbinFileRegistry.register(fbin, digest)) { + fixture = writeDescriptor(indexDirectory, registration.reference(0, vectors.length)); + } + + IOException failure = + expectThrows( + IOException.class, () -> new ExternalFbinFlatVectorsReader(fixture.readState())); + assertTrue(failure.getMessage().contains("No allowlisted external FBIN")); + } + } + + public void testIntegrityReportsOverflowingMutableHeaderAsCorruption() throws Exception { + float[][] vectors = {{1.0f, 2.0f}, {3.0f, 4.0f}}; + Path fbin = writeFbin(createTempDir().resolve("vectors.fbin"), vectors); + String digest = sha256(fbin); + + try (Directory indexDirectory = newDirectory(); + ExternalFbinFileRegistry.Registration registration = + ExternalFbinFileRegistry.register(fbin, digest)) { + SegmentFixture fixture = + writeDescriptor(indexDirectory, registration.reference(0, vectors.length)); + try (ExternalFbinFlatVectorsReader reader = + new ExternalFbinFlatVectorsReader(fixture.readState())) { + ByteBuffer corruptHeader = + ByteBuffer.allocate((int) ExternalFbinReference.HEADER_BYTES) + .order(ByteOrder.LITTLE_ENDIAN) + .putInt(Integer.MAX_VALUE) + .putInt(Integer.MAX_VALUE); + corruptHeader.flip(); + try (FileChannel channel = FileChannel.open(fbin, StandardOpenOption.WRITE)) { + while (corruptHeader.hasRemaining()) { + channel.write(corruptHeader, corruptHeader.position()); + } + } + expectThrows(CorruptIndexException.class, reader::checkIntegrity); + } + } + } + + public void testExternalMarkerIsScopedToSegmentSuffix() throws Exception { + float[][] vectors = {{1.0f, 2.0f}, {3.0f, 4.0f}}; + Path fbin = writeFbin(createTempDir().resolve("vectors.fbin"), vectors); + String digest = sha256(fbin); + + try (Directory indexDirectory = newDirectory(); + ExternalFbinFileRegistry.Registration registration = + ExternalFbinFileRegistry.register(fbin, digest)) { + SegmentFixture fixture = + writeDescriptor(indexDirectory, registration.reference(0, vectors.length)); + SegmentReadState base = fixture.readState(); + base.segmentInfo.putAttribute( + ExternalFbinReferenceWriter.segmentAttribute("external"), + ExternalFbinReferenceWriter.SEGMENT_ATTRIBUTE_VALUE); + + assertTrue( + ExternalFbinFlatVectorsReader.hasExternalMarker(new SegmentReadState(base, "external"))); + assertFalse( + ExternalFbinFlatVectorsReader.hasExternalMarker( + new SegmentReadState(base, "conventional"))); + } + } + + public void testOrphanDescriptorDoesNotSelectExternalReader() throws Exception { + float[][] vectors = {{1.0f, 2.0f}, {3.0f, 4.0f}}; + Path fbin = writeFbin(createTempDir().resolve("vectors.fbin"), vectors); + + try (Directory indexDirectory = newDirectory(); + ExternalFbinFileRegistry.Registration registration = + ExternalFbinFileRegistry.register(fbin, sha256(fbin))) { + SegmentFixture fixture = + writeDescriptor(indexDirectory, registration.reference(0, vectors.length)); + SegmentReadState state = fixture.readState(); + state.segmentInfo.putAttribute(ExternalFbinReferenceWriter.segmentAttribute(""), null); + assertFalse(ExternalFbinFlatVectorsReader.hasExternalMarker(state)); + + IOException failure = + expectThrows( + IOException.class, + () -> new Lucene99AcceleratedHNSWVectorsFormat().fieldsReader(state)); + assertFalse( + "An uncommitted descriptor must not override the committed segment marker: " + failure, + failure.getMessage().contains("marker/descriptor mismatch")); + } + } + + private static SegmentFixture writeDescriptor( + Directory directory, ExternalFbinReference reference) throws Exception { + FieldInfo fieldInfo = + new FieldInfo( + "vector", + 0, + false, + false, + false, + IndexOptions.NONE, + DocValuesType.NONE, + DocValuesSkipIndexType.NONE, + -1L, + Map.of(), + 0, + 0, + 0, + reference.dimensions(), + VectorEncoding.FLOAT32, + EUCLIDEAN, + false, + false); + FieldInfos fieldInfos = new FieldInfos(new FieldInfo[] {fieldInfo}); + SegmentInfo segmentInfo = + new SegmentInfo( + directory, + Version.LATEST, + Version.LATEST, + "_0", + reference.rows(), + false, + false, + Codec.getDefault(), + Map.of(), + StringHelper.randomId(), + Map.of(), + null); + segmentInfo.putAttribute( + ExternalFbinReferenceWriter.segmentAttribute(""), + ExternalFbinReferenceWriter.SEGMENT_ATTRIBUTE_VALUE); + SegmentWriteState writeState = + new SegmentWriteState( + InfoStream.NO_OUTPUT, directory, segmentInfo, fieldInfos, null, IOContext.DEFAULT); + try (ExternalFbinReferenceWriter writer = new ExternalFbinReferenceWriter(writeState)) { + writer.writeField(fieldInfo, reference); + writer.finish(); + } + return new SegmentFixture( + new SegmentReadState(directory, segmentInfo, fieldInfos, IOContext.DEFAULT)); + } + + private static void corruptDescriptorPayload(Directory directory, SegmentReadState state) + throws IOException { + String fileName = + IndexFileNames.segmentFileName( + state.segmentInfo.name, state.segmentSuffix, ExternalFbinReferenceWriter.EXTENSION); + byte[] bytes; + try (IndexInput input = directory.openInput(fileName, IOContext.DEFAULT)) { + bytes = new byte[Math.toIntExact(input.length())]; + input.readBytes(bytes, 0, bytes.length); + } + int lastDigestByte = bytes.length - CodecUtil.footerLength() - Integer.BYTES - 1; + bytes[lastDigestByte] ^= 1; + String replacement; + try (IndexOutput output = + directory.createTempOutput( + "corrupt", ExternalFbinReferenceWriter.EXTENSION, IOContext.DEFAULT)) { + replacement = output.getName(); + output.writeBytes(bytes, bytes.length); + } + directory.deleteFile(fileName); + directory.rename(replacement, fileName); + } + + private static Path writeFbin(Path file, float[][] vectors) throws IOException { + int rows = vectors.length; + int dimensions = vectors[0].length; + ByteBuffer data = + ByteBuffer.allocate(2 * Integer.BYTES + rows * dimensions * Float.BYTES) + .order(ByteOrder.LITTLE_ENDIAN); + data.putInt(rows); + data.putInt(dimensions); + for (float[] vector : vectors) { + for (float value : vector) { + data.putFloat(value); + } + } + data.flip(); + try (FileChannel channel = + FileChannel.open(file, StandardOpenOption.CREATE_NEW, StandardOpenOption.WRITE)) { + while (data.hasRemaining()) { + channel.write(data); + } + } + return file; + } + + private static void writeSparseFbin( + Path file, int rows, int dimensions, long finalRowOffset, float[] finalRow) + throws IOException { + ByteBuffer header = + ByteBuffer.allocate((int) ExternalFbinReference.HEADER_BYTES) + .order(ByteOrder.LITTLE_ENDIAN) + .putInt(rows) + .putInt(dimensions); + header.flip(); + ByteBuffer row = ByteBuffer.allocate(dimensions * Float.BYTES).order(ByteOrder.LITTLE_ENDIAN); + for (float value : finalRow) { + row.putFloat(value); + } + row.flip(); + + try (FileChannel channel = + FileChannel.open(file, StandardOpenOption.CREATE_NEW, StandardOpenOption.WRITE)) { + writeFully(channel, header, 0L); + writeFully(channel, row, finalRowOffset); + } + } + + private static void writeSparseFbinRows( + Path file, int rows, int dimensions, SparseRow... sparseRows) throws IOException { + ByteBuffer header = + ByteBuffer.allocate((int) ExternalFbinReference.HEADER_BYTES) + .order(ByteOrder.LITTLE_ENDIAN) + .putInt(rows) + .putInt(dimensions); + header.flip(); + + try (FileChannel channel = + FileChannel.open(file, StandardOpenOption.CREATE_NEW, StandardOpenOption.WRITE)) { + writeFully(channel, header, 0L); + for (SparseRow sparseRow : sparseRows) { + assertEquals(dimensions, sparseRow.values().length); + ByteBuffer row = + ByteBuffer.allocate(dimensions * Float.BYTES).order(ByteOrder.LITTLE_ENDIAN); + for (float value : sparseRow.values()) { + row.putFloat(value); + } + row.flip(); + writeFully(channel, row, sparseRow.offset()); + } + } + } + + private static void writeFully(FileChannel channel, ByteBuffer source, long position) + throws IOException { + while (source.hasRemaining()) { + int written = channel.write(source, position + source.position()); + if (written == 0) { + throw new IOException("Unable to make progress writing test FBIN"); + } + } + } + + private static float[][] createVectors(int rows, int dimensions) { + float[][] vectors = new float[rows][dimensions]; + for (int row = 0; row < rows; row++) { + for (int dimension = 0; dimension < dimensions; dimension++) { + vectors[row][dimension] = row * 0.25f + dimension * 0.03125f; + } + } + return vectors; + } + + private static float[] createVector(int dimensions, float base) { + float[] vector = new float[dimensions]; + for (int dimension = 0; dimension < dimensions; dimension++) { + vector[dimension] = base + dimension * 0.03125f; + } + return vector; + } + + private static String sha256(Path file) throws Exception { + MessageDigest digest = MessageDigest.getInstance("SHA-256"); + digest.update(Files.readAllBytes(file)); + return HexFormat.of().formatHex(digest.digest()); + } + + private record SegmentFixture(SegmentReadState readState) {} + + private record SparseRow(long offset, float[] values) {} +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestExternalFbinReference.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestExternalFbinReference.java new file mode 100644 index 0000000000..bc1fac3582 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestExternalFbinReference.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.Files; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.security.MessageDigest; +import java.util.ArrayList; +import java.util.HexFormat; +import java.util.List; +import java.util.Map; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.util.IOUtils; +import org.junit.After; + +public class TestExternalFbinReference extends LuceneTestCase { + + @After + public void clearRegistry() { + ExternalFbinFileRegistry.clearForTests(); + } + + public void testRegisteredReferenceRoundTripAndVerification() throws Exception { + Path file = writeFbin(createTempDir().resolve("vectors.fbin"), 4, 3); + String sha256 = sha256(file); + + try (ExternalFbinFileRegistry.Registration registration = + ExternalFbinFileRegistry.register(file, sha256)) { + ExternalFbinReference reference = registration.reference(1, 2); + + assertEquals("sha256:" + sha256, reference.contentId()); + assertEquals(1L, reference.firstRow()); + assertEquals(2, reference.rows()); + assertEquals(3, reference.dimensions()); + assertEquals(8L + 3L * Float.BYTES, reference.payloadOffset()); + assertEquals(2L * 3 * Float.BYTES, reference.payloadLength()); + assertEquals(file.toRealPath(), ExternalFbinIO.validateAndResolve(reference)); + assertEquals(reference.fileLength(), ExternalFbinIO.verifySha256(reference)); + + List prefetchProgress = new ArrayList<>(); + assertEquals( + reference.payloadLength(), ExternalFbinIO.prefetch(reference, prefetchProgress::add)); + assertFalse(prefetchProgress.isEmpty()); + assertEquals( + reference.payloadLength(), prefetchProgress.get(prefetchProgress.size() - 1).longValue()); + + List verifyProgress = new ArrayList<>(); + assertEquals( + reference.fileLength(), ExternalFbinIO.verifySha256(reference, verifyProgress::add)); + assertFalse(verifyProgress.isEmpty()); + assertEquals( + reference.fileLength(), verifyProgress.get(verifyProgress.size() - 1).longValue()); + } + } + + public void testScanHeadStartRequiresScanningValidationAndFitsPayload() throws Exception { + Path root = Files.createTempDirectory("cuvs-external-fbin-params-"); + try { + Path file = writeFbin(root.resolve("vectors.fbin"), 4, 3); + String sha256 = sha256(file); + + try (ExternalFbinFileRegistry.Registration registration = + ExternalFbinFileRegistry.register(file, sha256); + ImmutableExternalFbinDataset dataset = registration.map(0, 4)) { + IllegalArgumentException trustedFailure = + expectThrows( + IllegalArgumentException.class, + () -> new ExternalFbinOptions(ExternalFbinBuildValidation.TRUSTED_IMMUTABLE, 1L)); + assertTrue(trustedFailure.getMessage().contains("PREFETCH or VERIFY_SHA256")); + + ExternalFbinOptions oversized = + new ExternalFbinOptions( + ExternalFbinBuildValidation.PREFETCH, dataset.reference().payloadLength() + 1L); + IllegalArgumentException oversizedFailure = + expectThrows( + IllegalArgumentException.class, + () -> + BulkIndexingContext.external(dataset, oversized, new CagraHnswBuildMetrics())); + assertTrue(oversizedFailure.getMessage().contains("exceeds")); + } + } finally { + IOUtils.rm(root); + } + } + + public void testVerifyHeadStartAccountsForNonzeroFirstRow() throws Exception { + Path file = writeFbin(createTempDir().resolve("vectors.fbin"), 4, 3); + try (ExternalFbinFileRegistry.Registration registration = + ExternalFbinFileRegistry.register(file, sha256(file))) { + ExternalFbinReference reference = registration.reference(1, 2); + long requestedPayloadLead = (long) reference.dimensions() * Float.BYTES; + CagraHnswBuildMetrics metrics = new CagraHnswBuildMetrics(); + ExternalFbinScanCoordinator coordinator = + ExternalFbinScanCoordinator.start( + reference, ExternalFbinBuildValidation.VERIFY_SHA256, requestedPayloadLead, metrics); + + coordinator.runAfterHeadStart(() -> {}); + + Map snapshot = metrics.snapshot(); + assertEquals( + requestedPayloadLead, + snapshot.get("counter/external fbin requested head-start bytes").longValue()); + assertEquals( + reference.payloadOffset() + requestedPayloadLead, + snapshot.get("counter/external fbin required scan bytes").longValue()); + } + } + + public void testMissingRegistrationFailsResolution() throws Exception { + Path file = writeFbin(createTempDir().resolve("vectors.fbin"), 2, 2); + String sha256 = sha256(file); + ExternalFbinReference reference = ExternalFbinReference.fromFile(file, sha256, 0, 2); + + IOException failure = + expectThrows(IOException.class, () -> ExternalFbinIO.validateAndResolve(reference)); + assertTrue(failure.getMessage().contains("No allowlisted external FBIN")); + } + + public void testDigestMismatchFailsVerification() throws Exception { + Path file = writeFbin(createTempDir().resolve("vectors.fbin"), 3, 2); + String sha256 = sha256(file); + + try (ExternalFbinFileRegistry.Registration registration = + ExternalFbinFileRegistry.register(file, sha256)) { + ExternalFbinReference reference = registration.reference(0, 3); + try (FileChannel channel = + FileChannel.open(file, StandardOpenOption.WRITE, StandardOpenOption.READ)) { + channel.write(ByteBuffer.wrap(new byte[] {99}), ExternalFbinReference.HEADER_BYTES + 1); + } + + IOException failure = + expectThrows(IOException.class, () -> ExternalFbinIO.verifySha256(reference)); + assertTrue(failure.getMessage().contains("SHA-256 mismatch")); + } + } + + public void testTruncatedFbinIsRejected() throws Exception { + Path file = writeFbin(createTempDir().resolve("vectors.fbin"), 3, 2); + String sha256 = sha256(file); + try (FileChannel channel = FileChannel.open(file, StandardOpenOption.WRITE)) { + channel.truncate(channel.size() - 1); + } + + IOException failure = + expectThrows(IOException.class, () -> ExternalFbinReference.fromFile(file, sha256, 0, 3)); + assertTrue(failure.getMessage().contains("does not match header-derived length")); + } + + public void testAmbiguousContentRegistrationIsRejected() throws Exception { + Path directory = createTempDir(); + Path first = writeFbin(directory.resolve("first.fbin"), 2, 2); + Path second = Files.copy(first, directory.resolve("second.fbin")); + String sha256 = sha256(first); + + try (ExternalFbinFileRegistry.Registration ignored = + ExternalFbinFileRegistry.register(first, sha256)) { + IllegalStateException failure = + expectThrows( + IllegalStateException.class, () -> ExternalFbinFileRegistry.register(second, sha256)); + assertTrue(failure.getMessage().contains("refusing ambiguous replacement")); + } + } + + public void testRegistrationLeasesAreReferenceCounted() throws Exception { + Path file = writeFbin(createTempDir().resolve("vectors.fbin"), 2, 2); + String sha256 = sha256(file); + ExternalFbinFileRegistry.Registration first = ExternalFbinFileRegistry.register(file, sha256); + ExternalFbinFileRegistry.Registration second = ExternalFbinFileRegistry.register(file, sha256); + ExternalFbinReference reference = first.reference(0, 2); + + first.close(); + assertEquals(file.toRealPath(), ExternalFbinIO.validateAndResolve(reference)); + second.close(); + + IOException failure = + expectThrows(IOException.class, () -> ExternalFbinIO.validateAndResolve(reference)); + assertTrue(failure.getMessage().contains("No allowlisted external FBIN")); + } + + private static Path writeFbin(Path file, int rows, int dimensions) throws IOException { + ByteBuffer data = + ByteBuffer.allocate(2 * Integer.BYTES + rows * dimensions * Float.BYTES) + .order(ByteOrder.LITTLE_ENDIAN); + data.putInt(rows); + data.putInt(dimensions); + for (int row = 0; row < rows; row++) { + for (int dimension = 0; dimension < dimensions; dimension++) { + data.putFloat(row * 10.0f + dimension); + } + } + data.flip(); + try (FileChannel channel = + FileChannel.open(file, StandardOpenOption.CREATE_NEW, StandardOpenOption.WRITE)) { + while (data.hasRemaining()) { + channel.write(data); + } + } + return file; + } + + private static String sha256(Path file) throws Exception { + MessageDigest digest = MessageDigest.getInstance("SHA-256"); + digest.update(Files.readAllBytes(file)); + return HexFormat.of().formatHex(digest.digest()); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestExternalFbinScanCoordinator.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestExternalFbinScanCoordinator.java new file mode 100644 index 0000000000..da4b5e6a62 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestExternalFbinScanCoordinator.java @@ -0,0 +1,200 @@ +/* + * 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.io.InterruptedIOException; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import org.apache.lucene.tests.util.LuceneTestCase; + +public class TestExternalFbinScanCoordinator extends LuceneTestCase { + + public void testZeroLeadStartsBuildWithoutWaitingForScanProgress() throws Exception { + CountDownLatch buildStarted = new CountDownLatch(1); + ExternalFbinScanCoordinator coordinator = + ExternalFbinScanCoordinator.createForTests( + 0L, + 0L, + 10L, + progress -> { + await(buildStarted); + progress.accept(10L); + return 10L; + }); + + coordinator.runAfterHeadStart(buildStarted::countDown); + } + + public void testZeroLeadPreInterruptedCallerDoesNotStartBuild() throws Exception { + CountDownLatch releaseScan = new CountDownLatch(1); + AtomicBoolean buildStarted = new AtomicBoolean(); + ExternalFbinScanCoordinator coordinator = + ExternalFbinScanCoordinator.createForTests( + 0L, + 0L, + 16L, + progress -> { + try { + releaseScan.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException("scan interrupted", e); + } + progress.accept(16L); + return 16L; + }); + + Thread.currentThread().interrupt(); + try { + expectThrows( + InterruptedIOException.class, + () -> coordinator.runAfterHeadStart(() -> buildStarted.set(true))); + assertFalse(buildStarted.get()); + assertTrue(Thread.currentThread().isInterrupted()); + } finally { + releaseScan.countDown(); + Thread.interrupted(); + } + } + + public void testBuildStartsOnlyAfterRequiredProgress() throws Exception { + AtomicBoolean buildStarted = new AtomicBoolean(); + CountDownLatch releaseScanner = new CountDownLatch(1); + ExternalFbinScanCoordinator coordinator = + ExternalFbinScanCoordinator.createForTests( + 8L, + 8L, + 10L, + progress -> { + progress.accept(4L); + assertFalse(buildStarted.get()); + progress.accept(8L); + await(releaseScanner); + progress.accept(10L); + return 10L; + }); + + coordinator.runAfterHeadStart( + () -> { + buildStarted.set(true); + releaseScanner.countDown(); + }); + assertTrue(buildStarted.get()); + } + + public void testFailureBeforeLeadPreventsBuild() throws Exception { + AtomicBoolean buildStarted = new AtomicBoolean(); + ExternalFbinScanCoordinator coordinator = + ExternalFbinScanCoordinator.createForTests( + 8L, + 8L, + 10L, + progress -> { + progress.accept(4L); + throw new IOException("scan failed before lead"); + }); + + IOException failure = + expectThrows( + IOException.class, () -> coordinator.runAfterHeadStart(() -> buildStarted.set(true))); + assertEquals("scan failed before lead", failure.getMessage()); + assertFalse(buildStarted.get()); + } + + public void testFailureAfterLeadFailsCompletedBuild() throws Exception { + CountDownLatch buildStarted = new CountDownLatch(1); + ExternalFbinScanCoordinator coordinator = + ExternalFbinScanCoordinator.createForTests( + 8L, + 8L, + 10L, + progress -> { + progress.accept(8L); + await(buildStarted); + throw new IOException("scan failed after lead"); + }); + + IOException failure = + expectThrows( + IOException.class, () -> coordinator.runAfterHeadStart(buildStarted::countDown)); + assertEquals("scan failed after lead", failure.getMessage()); + assertEquals(0L, buildStarted.getCount()); + } + + public void testBuildFailureCancelsScanner() throws Exception { + CountDownLatch scannerWaiting = new CountDownLatch(1); + AtomicBoolean scannerInterrupted = new AtomicBoolean(); + ExternalFbinScanCoordinator coordinator = + ExternalFbinScanCoordinator.createForTests( + 1L, + 1L, + 10L, + progress -> { + progress.accept(1L); + scannerWaiting.countDown(); + try { + new CountDownLatch(1).await(); + throw new AssertionError("Scanner was not cancelled"); + } catch (InterruptedException expected) { + scannerInterrupted.set(true); + throw new IOException("scanner cancelled", expected); + } + }); + + IOException failure = + expectThrows( + IOException.class, + () -> + coordinator.runAfterHeadStart( + () -> { + assertTrue(scannerWaiting.await(10, TimeUnit.SECONDS)); + throw new IOException("graph failed"); + })); + assertEquals("graph failed", failure.getMessage()); + assertTrue(scannerInterrupted.get()); + } + + public void testCallerInterruptionCancelsScanAndPreservesInterrupt() throws Exception { + CountDownLatch scannerStarted = new CountDownLatch(1); + AtomicBoolean scannerInterrupted = new AtomicBoolean(); + ExternalFbinScanCoordinator coordinator = + ExternalFbinScanCoordinator.createForTests( + 1L, + 1L, + 10L, + progress -> { + scannerStarted.countDown(); + try { + new CountDownLatch(1).await(); + throw new AssertionError("Scanner was not cancelled"); + } catch (InterruptedException expected) { + scannerInterrupted.set(true); + throw new IOException("scanner cancelled", expected); + } + }); + + assertTrue(scannerStarted.await(10, TimeUnit.SECONDS)); + Thread.currentThread().interrupt(); + try { + expectThrows( + InterruptedIOException.class, () -> coordinator.runAfterHeadStart(() -> fail("build"))); + assertTrue(Thread.currentThread().isInterrupted()); + assertTrue(scannerInterrupted.get()); + } finally { + Thread.interrupted(); + } + } + + private static void await(CountDownLatch latch) throws IOException { + try { + assertTrue(latch.await(10, TimeUnit.SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException("Interrupted while awaiting test coordination", e); + } + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestExternalFloat32Dataset.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestExternalFloat32Dataset.java new file mode 100644 index 0000000000..f8a569b830 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestExternalFloat32Dataset.java @@ -0,0 +1,36 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import java.lang.foreign.Arena; +import java.lang.foreign.MemorySegment; +import org.apache.lucene.tests.util.LuceneTestCase; + +public class TestExternalFloat32Dataset extends LuceneTestCase { + + public void testRejectsWritableMemory() { + try (Arena arena = Arena.ofShared()) { + MemorySegment writable = arena.allocate(2L * 3 * Float.BYTES, Float.BYTES); + + expectThrows( + IllegalArgumentException.class, + () -> ExternalFloat32Dataset.fromMemorySegment(writable, 2, 3)); + } + } + + public void testRejectsMisalignedMemory() { + try (Arena arena = Arena.ofShared()) { + MemorySegment misaligned = + arena + .allocate(2L * 3 * Float.BYTES + 1, Float.BYTES) + .asSlice(1, 2L * 3 * Float.BYTES) + .asReadOnly(); + + expectThrows( + IllegalArgumentException.class, + () -> ExternalFloat32Dataset.fromMemorySegment(misaligned, 2, 3)); + } + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestFbinFileMetadata.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestFbinFileMetadata.java new file mode 100644 index 0000000000..a3eb8d57cc --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestFbinFileMetadata.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 java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import org.apache.lucene.tests.util.LuceneTestCase; + +public class TestFbinFileMetadata extends LuceneTestCase { + + public void testReadsShapeWithoutAStreamingSource() throws Exception { + Path file = createTempDir().resolve("vectors.fbin"); + TestUtils.writeFbin(file, new float[][] {{1f, 2f}, {3f, 4f}, {5f, 6f}}); + + FbinFileMetadata metadata = FbinFileMetadata.read(file); + + assertEquals(3, metadata.rows()); + assertEquals(2, metadata.dimensions()); + assertEquals(Files.size(file), metadata.fileBytes()); + } + + public void testRejectsTrailingBytes() throws Exception { + Path file = createTempDir().resolve("vectors.fbin"); + TestUtils.writeFbin(file, new float[][] {{1f, 2f}}); + Files.write(file, new byte[] {1}, java.nio.file.StandardOpenOption.APPEND); + + IOException failure = expectThrows(IOException.class, () -> FbinFileMetadata.read(file)); + assertTrue(failure.getMessage().contains("does not match")); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGPUBuiltHnswGraph.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGPUBuiltHnswGraph.java new file mode 100644 index 0000000000..290abbe287 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGPUBuiltHnswGraph.java @@ -0,0 +1,49 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import org.apache.lucene.tests.util.LuceneTestCase; +import org.junit.Test; + +public class TestGPUBuiltHnswGraph extends LuceneTestCase { + + @Test + public void findsUpperLayerOrdinalInSortedNodeIds() { + int[] nodeIds = {3, 11, 27, 42, 81}; + + assertEquals(0, GPUBuiltHnswGraph.findUpperLayerOrdinal(nodeIds, 3)); + assertEquals(2, GPUBuiltHnswGraph.findUpperLayerOrdinal(nodeIds, 27)); + assertEquals(4, GPUBuiltHnswGraph.findUpperLayerOrdinal(nodeIds, 81)); + } + + @Test + public void returnsMissingOrdinalForUnknownUpperLayerNode() { + int[] nodeIds = {3, 11, 27, 42, 81}; + + assertEquals(-1, GPUBuiltHnswGraph.findUpperLayerOrdinal(nodeIds, 2)); + assertEquals(-1, GPUBuiltHnswGraph.findUpperLayerOrdinal(nodeIds, 28)); + assertEquals(-1, GPUBuiltHnswGraph.findUpperLayerOrdinal(nodeIds, 82)); + } + + @Test(timeout = 5000L) + public void repeatedUpperLayerLookupsDoNotUseLinearScans() { + int[] nodeIds = new int[1 << 20]; + for (int i = 0; i < nodeIds.length; i++) { + nodeIds[i] = i * 2; + } + + // These near-tail lookups require about 100 billion comparisons with a linear scan, while a + // binary search performs about two million. The wide timeout keeps ordinary CI noise + // irrelevant. + long ordinalSum = 0; + for (int i = 0; i < 100_000; i++) { + int expectedOrdinal = nodeIds.length - 1 - (i & 1023); + int ordinal = GPUBuiltHnswGraph.findUpperLayerOrdinal(nodeIds, nodeIds[expectedOrdinal]); + assertEquals(expectedOrdinal, ordinal); + ordinalSum += ordinal; + } + assertTrue(ordinalSum > 0); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMappedFbinDataset.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMappedFbinDataset.java new file mode 100644 index 0000000000..73a742936c --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMappedFbinDataset.java @@ -0,0 +1,64 @@ +/* + * 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.Files; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressFileSystems; +import org.junit.Test; + +@SuppressFileSystems("*") +public class TestMappedFbinDataset extends LuceneTestCase { + + @Test + public void testMapsDenseFloatPayloadWithoutCopying() throws Exception { + Path file = createTempDir().resolve("vectors.fbin"); + writeFbin(file, 2, 3, new float[] {1, 2, 3, 4, 5, 6}); + + try (MappedFbinDataset mapped = MappedFbinDataset.map(file)) { + assertEquals(2, mapped.rows()); + assertEquals(3, mapped.dimensions()); + float[] second = new float[3]; + mapped.dataset().matrix().getRow(1).toArray(second); + assertArrayEquals(new float[] {4, 5, 6}, second, 0.0f); + } + } + + @Test + public void testRejectsTrailingOrTruncatedPayload() throws Exception { + Path trailing = createTempDir().resolve("trailing.fbin"); + writeFbin(trailing, 1, 2, new float[] {1, 2, 3}); + assertThrows(IOException.class, () -> MappedFbinDataset.map(trailing)); + + Path truncated = createTempDir().resolve("truncated.fbin"); + writeFbin(truncated, 2, 2, new float[] {1, 2, 3}); + assertThrows(IOException.class, () -> MappedFbinDataset.map(truncated)); + } + + private static void writeFbin(Path path, int rows, int dimensions, float[] values) + throws IOException { + ByteBuffer data = + ByteBuffer.allocate(2 * Integer.BYTES + values.length * Float.BYTES) + .order(ByteOrder.LITTLE_ENDIAN); + data.putInt(rows).putInt(dimensions); + for (float value : values) { + data.putFloat(value); + } + data.flip(); + try (FileChannel channel = + FileChannel.open(path, StandardOpenOption.CREATE_NEW, StandardOpenOption.WRITE)) { + while (data.hasRemaining()) { + channel.write(data); + } + } + assertTrue(Files.isRegularFile(path)); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferedCleanup.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferedCleanup.java new file mode 100644 index 0000000000..4ee42dc4b4 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferedCleanup.java @@ -0,0 +1,49 @@ +/* + * 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.CuVSResources; +import java.io.Closeable; +import java.io.IOException; +import java.lang.reflect.Proxy; +import java.util.concurrent.atomic.AtomicBoolean; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.junit.Test; + +public class TestNativeFlatBufferedCleanup extends LuceneTestCase { + + @Test + public void testThreadResourcesCloseWhenOutputCloseFails() { + AtomicBoolean resourcesClosed = new AtomicBoolean(); + CuVSResources resources = resourcesThatRecordClose(resourcesClosed); + ThreadLocalCuVSResourcesProvider.setCuVSResourcesInstance(resources); + Closeable failingOutput = + () -> { + throw new IOException("output close failed"); + }; + + IOException failure = + expectThrows( + IOException.class, + () -> NativeFlatBufferedHNSWVectorsWriter.closeOutputsAndResources(failingOutput)); + + assertEquals("output close failed", failure.getMessage()); + assertTrue("cuVS resources were not closed", resourcesClosed.get()); + } + + private static CuVSResources resourcesThatRecordClose(AtomicBoolean closed) { + return (CuVSResources) + Proxy.newProxyInstance( + CuVSResources.class.getClassLoader(), + new Class[] {CuVSResources.class}, + (proxy, method, arguments) -> { + if (method.getName().equals("close")) { + closed.set(true); + return null; + } + throw new UnsupportedOperationException(method.getName()); + }); + } +} From 472918658b5fe2bac1b47453ac9cc6af44f9c33d Mon Sep 17 00:00:00 2001 From: EC2 Default User Date: Fri, 11 Sep 2026 04:45:54 +0000 Subject: [PATCH 2/3] Support merge-free multi-segment FBIN bulk builds (cuvs-issue-2592) --- java/cuvs-lucene/README.md | 28 +- .../lucene/AcceleratedHnswGraphOutput.java | 2 + .../cuvs/lucene/CagraHnswBuildMetrics.java | 28 +- .../cuvs/lucene/CagraHnswBulkIndexWriter.java | 303 +++++++++++--- .../cuvs/lucene/ExternalFloat32Dataset.java | 10 +- .../nvidia/cuvs/lucene/MappedFbinDataset.java | 76 ++-- .../TestBulkFbinMultiSegmentStorage.java | 396 ++++++++++++++++++ .../TestCagraHnswBuildMetricsConcurrency.java | 14 +- .../cuvs/lucene/TestMappedFbinDataset.java | 27 ++ 9 files changed, 762 insertions(+), 122 deletions(-) create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBulkFbinMultiSegmentStorage.java diff --git a/java/cuvs-lucene/README.md b/java/cuvs-lucene/README.md index 2a192abebd..0d80f08953 100644 --- a/java/cuvs-lucene/README.md +++ b/java/cuvs-lucene/README.md @@ -159,18 +159,21 @@ The streaming/native-buffered mode has three entry shapes: overlapped multi-segment pipeline in which host ingest can overlap a serialized GPU build. `FieldCallback` can add non-vector fields in either one-shot form. -Both mapped modes currently require `segments(1, false)`. They also require exactly one dense -`FLOAT32` vector field, one vector per document, and row `i` to correspond to Lucene document and -vector ordinal `i`. +Both mapped modes support one or more contiguous, independently built segments. With overlap +disabled, slices are built sequentially; with overlap enabled, host-side preparation can run in a +bounded pipeline while GPU commits remain serialized. The slices are published together without a +Lucene vector merge. They require exactly one dense `FLOAT32` vector field, one vector per document, +and preserve each source row's absolute ID across segment-local document ordinals. ### Ownership and sharing All bulk forms own their internal `IndexWriter`, disable index sorting, automatic flushes, compound files, and merges, and commit only after the promised vector count has been written. Closing a manual writer with too few vectors rolls it back; adding too many vectors is rejected. A failed -segment is rolled back. A sequential multi-segment build can already have committed earlier -segments when a later segment fails, so callers must treat and replace that target as a failed -build. +segment is rolled back. Multi-segment mapped builds use temporary per-slice indexes and do not +publish the final target until every slice succeeds. Streaming sequential builds append directly, +so they can already have committed earlier segments when a later segment fails; callers must treat +and replace that target as a failed build. The generic, streaming, and mapped self-contained results can be moved by copying the Lucene index directory. The immutable external-FBIN result must be shipped as two artifacts: the Lucene index @@ -226,7 +229,9 @@ larger than the referenced payload. Opening an external-FBIN index verifies the descriptor checksum, field metadata, registered path, header, range, and file length, but does not hash the full file. Lucene's `checkIntegrity()` hashes -the complete external FBIN and compares it with the persisted SHA-256. +the complete external FBIN and compares it with the persisted SHA-256. Because integrity checking +is a per-segment Lucene API, a K-segment external index currently hashes that shared FBIN K times; +this cost does not occur during normal reader opening or search. ### Metrics @@ -240,10 +245,13 @@ Metric keys use three namespaces: graph build, graph conversion and output, mapped flat output, external scan/overlap, and bulk writer commit/close stages. The commit wall includes the flush and its nested GPU/output work; nested stage durations must not be added to it. -- `counter/` reports values such as logical adjacency bytes, mapped flat chunks, and external - scan progress at CAGRA start and end. +- `counter/` reports values such as logical adjacency bytes, mapped flat chunks, the number + of segments using each effective CAGRA algorithm, and external scan progress at CAGRA start and + end. - `gauge/` records effective graph-build parameters, including graph degrees, writer threads, - NN-descent iterations, and IVF-PQ dimensions, lists, probes, and k-means iterations when used. + NN-descent iterations, and IVF-PQ dimensions, lists, probes, and k-means iterations when used. If + independent segments resolve different heuristic values, the scalar key is replaced by explicit + `gauge//min` and `gauge//max` keys. ### Why these controls are bulk-only diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHnswGraphOutput.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHnswGraphOutput.java index db9dbb9366..60dc3aef6d 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHnswGraphOutput.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHnswGraphOutput.java @@ -209,6 +209,8 @@ private void writeField(FieldInfo fieldInfo, CuVSMatrix dataset, boolean closeDa } private void recordEffectiveParams(CagraIndexParams params) { + metrics.addCounter( + "effective cagra algorithm/" + params.getCagraGraphBuildAlgo() + " segments", 1L); metrics.setGauge("effective graph degree", params.getGraphDegree()); metrics.setGauge("effective intermediate graph degree", params.getIntermediateGraphDegree()); metrics.setGauge("effective writer threads", params.getNumWriterThreads()); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraHnswBuildMetrics.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraHnswBuildMetrics.java index 297e54a63b..07a0fbef75 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraHnswBuildMetrics.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraHnswBuildMetrics.java @@ -15,7 +15,7 @@ public final class CagraHnswBuildMetrics { private final Map stages = new ConcurrentHashMap<>(); private final Map counters = new ConcurrentHashMap<>(); - private final Map gauges = new ConcurrentHashMap<>(); + private final Map gauges = new ConcurrentHashMap<>(); static long start() { return System.nanoTime(); @@ -45,14 +45,15 @@ void addCounter(String name, long value) { counters.computeIfAbsent(name, ignored -> new LongAdder()).add(value); } - /** Records a configuration value that must remain identical across every segment in a build. */ + /** Records an effective configuration value, retaining its range across independent segments. */ void setGauge(String name, long value) { Objects.requireNonNull(name, "name"); - Long existing = gauges.putIfAbsent(name, value); - if (existing != null && existing.longValue() != value) { - throw new IllegalStateException( - "Gauge \"" + name + "\" changed from " + existing + " to " + value); - } + gauges.compute( + name, + (ignored, current) -> + current == null + ? new GaugeRange(value, value) + : new GaugeRange(Math.min(current.min(), value), Math.max(current.max(), value))); } /** Returns a stable machine-readable copy; later measurements do not mutate it. */ @@ -75,7 +76,16 @@ public Map snapshot() { .forEach(name -> result.put("counter/" + name, counters.get(name).sum())); gauges.keySet().stream() .sorted() - .forEach(name -> result.put("gauge/" + name, gauges.get(name))); + .forEach( + name -> { + GaugeRange range = gauges.get(name); + if (range.min() == range.max()) { + result.put("gauge/" + name, range.min()); + } else { + result.put("gauge/" + name + "/min", range.min()); + result.put("gauge/" + name + "/max", range.max()); + } + }); return Map.copyOf(result); } @@ -89,4 +99,6 @@ StageMeasurement plus(long additionalNanos, long additionalBytes) { return new StageMeasurement(nanos + additionalNanos, count + 1L, bytes + additionalBytes); } } + + private record GaugeRange(long min, long max) {} } 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 index 161bea8cd2..f1503acc06 100644 --- 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 @@ -7,12 +7,14 @@ import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; import java.io.Closeable; import java.io.IOException; +import java.io.InterruptedIOException; 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.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; @@ -40,11 +42,11 @@ * *

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. + * writer requires (one controlled flush per segment, no index sort or merges, 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: * @@ -52,11 +54,11 @@ *

  • 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. + *
  • One-shot, one or many segments: {@link #indexFbin}, {@link #indexMappedFbin}, {@link + * #indexImmutableFbin}, and {@link #build(VectorSource, Config)} own the loop for you. They + * optionally split the input into {@code numSegments} partitions, build one independent + * graph per partition, and combine them without merging. 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 @@ -93,7 +95,7 @@ public final class CagraHnswBulkIndexWriter implements Closeable { *

    {@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)}. + * fields only matter to the one-shot build methods. */ public CagraHnswBulkIndexWriter( Directory directory, IndexWriterConfig conf, Config config, int exactVectorCount) @@ -119,7 +121,7 @@ private CagraHnswBulkIndexWriter( 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" + + " buffering requires an unsorted segment build); leave" + " IndexWriterConfig.indexSort unset"); } this.config = config; @@ -272,11 +274,12 @@ void abort() throws IOException { } /** - * 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. + * Callback invoked once per row by the one-shot build methods, right after the id and vector + * fields have been added to {@code document} and right before it is added to the index. This lets + * the caller attach additional fields (metadata) per vector. Overlapped builds may invoke the + * callback concurrently, so callback implementations must be thread-safe. Not used by the + * manual, direct-instance API, where the caller already builds the whole {@link Document} + * themselves. */ @FunctionalInterface public interface FieldCallback { @@ -385,11 +388,36 @@ public static void indexMappedFbin(Path fbinPath, Config config, FieldCallback c throws Exception { Objects.requireNonNull(fbinPath, "fbinPath"); Objects.requireNonNull(config, "config"); - requireSingleSegmentMappedBuild(config); + requireMappedBuildTarget(config); + if (config.numSegments() > 1) { + FbinFileMetadata metadata = FbinFileMetadata.read(fbinPath); + validateMappedShape(metadata.rows(), metadata.dimensions(), config); + List slices = sliceEvenly(metadata.rows(), config.numSegments()); + buildBorrowedSegments( + config, + slices, + (segmentDirectory, firstRow, rowCount, gpuPermit) -> { + try (MappedFbinDataset mapped = MappedFbinDataset.map(fbinPath, firstRow, rowCount)) { + buildBorrowedSegment( + segmentDirectory, + config, + callback, + BulkIndexingContext.mapped(mapped.dataset(), config.metrics()), + firstRow, + gpuPermit); + } + }, + null, + null); + return; + } try (MappedFbinDataset mapped = MappedFbinDataset.map(fbinPath)) { validateMappedShape(mapped.rows(), mapped.dimensions(), config); - buildBorrowed( - config, callback, BulkIndexingContext.mapped(mapped.dataset(), config.metrics())); + buildBorrowedSegment( + config.targetDirectory(), + config, + callback, + BulkIndexingContext.mapped(mapped.dataset(), config.metrics())); } } @@ -415,25 +443,46 @@ public static void indexImmutableFbin( Objects.requireNonNull(source, "source"); Objects.requireNonNull(config, "config"); Objects.requireNonNull(options, "options"); - requireSingleSegmentMappedBuild(config); + requireMappedBuildTarget(config); FbinFileMetadata metadata = FbinFileMetadata.read(source.path()); int rows = metadata.rows(); int dimensions = metadata.dimensions(); validateMappedShape(rows, dimensions, config); + if (config.numSegments() > 1) { + List slices = sliceEvenly(rows, config.numSegments()); + ExternalFbinOptions trustedOptions = + new ExternalFbinOptions(ExternalFbinBuildValidation.TRUSTED_IMMUTABLE, 0L); + buildBorrowedSegments( + config, + slices, + (segmentDirectory, firstRow, rowCount, gpuPermit) -> { + try (ImmutableExternalFbinDataset dataset = source.map(firstRow, rowCount)) { + buildBorrowedSegment( + segmentDirectory, + config, + callback, + BulkIndexingContext.external(dataset, trustedOptions, config.metrics()), + firstRow, + gpuPermit); + } + }, + source.reference(0, rows), + options); + return; + } try (ImmutableExternalFbinDataset dataset = source.map(0, rows)) { - buildBorrowed( - config, callback, BulkIndexingContext.external(dataset, options, config.metrics())); + buildBorrowedSegment( + config.targetDirectory(), + config, + callback, + BulkIndexingContext.external(dataset, options, config.metrics())); } } - private static void requireSingleSegmentMappedBuild(Config config) { + private static void requireMappedBuildTarget(Config config) { if (config.targetDirectory() == null) { throw new IllegalArgumentException("targetDirectory is required for a one-shot bulk build"); } - if (config.numSegments() != 1 || config.overlapped()) { - throw new IllegalArgumentException( - "Mapped FBIN bulk builds currently require segments(1, false)"); - } } private static void validateMappedShape(int rows, int dimensions, Config config) { @@ -448,18 +497,41 @@ private static void validateMappedShape(int rows, int dimensions, Config config) } } - private static void buildBorrowed( - Config config, FieldCallback callback, BulkIndexingContext context) throws Exception { + private static void buildBorrowedSegment( + Path targetDirectory, Config config, FieldCallback callback, BulkIndexingContext context) + throws Exception { + buildBorrowedSegment(targetDirectory, config, callback, context, 0, null); + } + + private static void buildBorrowedSegment( + Path targetDirectory, + Config config, + FieldCallback callback, + BulkIndexingContext context, + int idStart, + Semaphore gpuPermit) + throws Exception { long ingestStartedAt = CagraHnswBuildMetrics.start(); - try (Directory directory = FSDirectory.open(config.targetDirectory())) { + try (Directory directory = FSDirectory.open(targetDirectory)) { IndexWriterConfig writerConfig = new IndexWriterConfig().setOpenMode(IndexWriterConfig.OpenMode.CREATE); CagraHnswBulkIndexWriter writer = new CagraHnswBulkIndexWriter(directory, writerConfig, config, context); try { - addBorrowedDocuments(writer, config, callback, context.exactVectorCount()); + addBorrowedDocuments(writer, config, callback, idStart, context.exactVectorCount()); config.metrics().stop("borrowed document ingest [CPU]", ingestStartedAt); - writer.close(); + boolean permitAcquired = false; + try { + if (gpuPermit != null) { + gpuPermit.acquire(); + permitAcquired = true; + } + writer.close(); + } finally { + if (permitAcquired) { + gpuPermit.release(); + } + } } catch (Throwable failure) { try { writer.abort(); @@ -472,7 +544,11 @@ private static void buildBorrowed( } private static void addBorrowedDocuments( - CagraHnswBulkIndexWriter writer, Config config, FieldCallback callback, int count) + CagraHnswBulkIndexWriter writer, + Config config, + FieldCallback callback, + int idStart, + int count) throws IOException { float[] placeholder = new float[config.dimensions()]; if (callback == null) { @@ -485,7 +561,8 @@ private static void addBorrowedDocuments( document.add(idField); } document.add(new KnnFloatVectorField(config.fieldName(), placeholder, config.similarity())); - for (int id = 0; id < count; id++) { + for (int localId = 0; localId < count; localId++) { + int id = Math.addExact(idStart, localId); if (idField != null) { idField.setStringValue(Integer.toString(id)); } @@ -494,7 +571,8 @@ private static void addBorrowedDocuments( return; } - for (int id = 0; id < count; id++) { + for (int localId = 0; localId < count; localId++) { + int id = Math.addExact(idStart, localId); Document document = new Document(); if (config.idFieldName() != null) { document.add(new StringField(config.idFieldName(), Integer.toString(id), Field.Store.YES)); @@ -505,6 +583,122 @@ private static void addBorrowedDocuments( } } + @FunctionalInterface + private interface BorrowedSegmentBuilder { + void build(Path segmentDirectory, int firstRow, int rowCount, Semaphore gpuPermit) + throws Exception; + } + + private static void buildBorrowedSegments( + Config config, + List slices, + BorrowedSegmentBuilder segmentBuilder, + ExternalFbinReference validationReference, + ExternalFbinOptions validationOptions) + throws Exception { + Path target = config.targetDirectory().toAbsolutePath(); + Path parent = target.getParent(); + Files.createDirectories(parent); + Path temporaryRoot = + Files.createTempDirectory(parent, target.getFileName().toString() + ".bulk-"); + List segmentDirectories = new ArrayList<>(slices.size()); + for (int segment = 0; segment < slices.size(); segment++) { + segmentDirectories.add(temporaryRoot.resolve("segment-" + segment)); + } + + try { + if (validationReference == null + || validationOptions.validation() == ExternalFbinBuildValidation.TRUSTED_IMMUTABLE) { + executeBorrowedSegmentBuilds(config, slices, segmentDirectories, segmentBuilder); + } else { + long overlapStartedAt = CagraHnswBuildMetrics.start(); + try { + ExternalFbinScanCoordinator coordinator = + ExternalFbinScanCoordinator.start( + validationReference, + validationOptions.validation(), + validationOptions.scanHeadStartBytes(), + config.metrics()); + coordinator.runAfterHeadStart( + () -> + executeBorrowedSegmentBuilds(config, slices, segmentDirectories, segmentBuilder)); + } finally { + config.metrics().stop("external reference overlap wall [GPU+DISK]", overlapStartedAt); + } + } + + long combineStartedAt = CagraHnswBuildMetrics.start(); + try { + combineByHardlink(target, segmentDirectories); + } finally { + config.metrics().stop("segment combine [DISK]", combineStartedAt); + } + } finally { + deleteRecursivelyQuietly(temporaryRoot); + } + } + + private static void executeBorrowedSegmentBuilds( + Config config, + List slices, + List segmentDirectories, + BorrowedSegmentBuilder segmentBuilder) + throws Exception { + if (!config.overlapped()) { + for (int segment = 0; segment < slices.size(); segment++) { + int[] slice = slices.get(segment); + segmentBuilder.build(segmentDirectories.get(segment), slice[0], slice[1], null); + } + return; + } + + int depth = Math.min(slices.size(), config.pipelineDepth()); + Semaphore gpuPermit = new Semaphore(1); + ExecutorService pool = Executors.newFixedThreadPool(depth); + List> futures = new ArrayList<>(slices.size()); + try { + for (int segment = 0; segment < slices.size(); segment++) { + int[] slice = slices.get(segment); + Path segmentDirectory = segmentDirectories.get(segment); + futures.add( + pool.submit( + () -> { + segmentBuilder.build(segmentDirectory, slice[0], slice[1], gpuPermit); + return null; + })); + } + for (Future future : futures) { + try { + future.get(); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + InterruptedIOException failure = + new InterruptedIOException("Interrupted while awaiting a bulk segment build"); + failure.initCause(interrupted); + throw failure; + } catch (ExecutionException failure) { + rethrowSegmentFailure(failure.getCause()); + } + } + } finally { + for (Future future : futures) { + future.cancel(true); + } + pool.shutdownNow(); + pool.close(); + } + } + + private static void rethrowSegmentFailure(Throwable failure) throws Exception { + if (failure instanceof Exception exception) { + throw exception; + } + if (failure instanceof Error error) { + throw error; + } + throw new RuntimeException("Unexpected segment-build failure", failure); + } + /** * 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 @@ -665,10 +859,22 @@ private static void combineByHardlink(Path targetDir, List segDirs) throws 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); + new IndexWriterConfig() + .setMergePolicy(NoMergePolicy.INSTANCE) + .setOpenMode(IndexWriterConfig.OpenMode.CREATE); // keep only imported segments + try (Directory target = new HardlinkCopyDirectoryWrapper(FSDirectory.open(targetDir))) { + IndexWriter combiner = new IndexWriter(target, iwc); + try { + combiner.addIndexes(sources); + combiner.close(); + } catch (Throwable failure) { + try { + combiner.rollback(); + } catch (Throwable cleanupFailure) { + failure.addSuppressed(cleanupFailure); + } + throw Utils.handleThrowable(failure); + } } } finally { for (Directory s : sources) { @@ -761,22 +967,22 @@ public AcceleratedHNSWParams graphBuildParams() { return graphBuildParams; } - /** Only consulted by {@link #indexFbin} / {@link #build(VectorSource, Config)}. */ + /** Only consulted by the one-shot build methods. */ public Path targetDirectory() { return targetDirectory; } - /** Only consulted by {@link #indexFbin} / {@link #build(VectorSource, Config)}. */ + /** Only consulted by the one-shot build methods. */ public int numSegments() { return numSegments; } - /** Only consulted by {@link #indexFbin}. */ + /** Only consulted by one-shot methods that can open independent input slices. */ public boolean overlapped() { return overlapped; } - /** Only consulted by {@link #indexFbin}. */ + /** Only consulted by one-shot methods that can open independent input slices. */ public int pipelineDepth() { return pipelineDepth; } @@ -839,11 +1045,10 @@ public Builder targetDirectory(Path targetDirectory) { * 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. + * @param overlapped when {@code numSegments > 1}, builds up to {@link #pipelineDepth} slices + * concurrently instead of strictly sequentially. Supported by the FBIN entry points and + * ignored by {@link #build(VectorSource, Config)}, whose caller-owned source is + * single-consumer. */ public Builder segments(int numSegments, boolean overlapped) { if (numSegments < 1) { diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFloat32Dataset.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFloat32Dataset.java index fe356b40cf..039f4ed0f2 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFloat32Dataset.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFloat32Dataset.java @@ -18,11 +18,11 @@ * writing has completed. A shared scope is required because the accelerated writer may read the * dataset concurrently while serializing Lucene's flat-vector file. * - *

    This is an expert-only input for the unsorted, single-segment native-buffering path. Row - * {@code i} must correspond exactly to Lucene vector ordinal and document ID {@code i}; gaps, - * reordered documents, and multiple vector fields are rejected by the writer. The caller is also - * responsible for ensuring that every raw vector component is finite; unlike ordinary Lucene field - * ingestion, this path deliberately does not scan the payload one float at a time. + *

    This is an expert-only input for one unsorted native-buffered segment. Row {@code i} must + * correspond exactly to vector ordinal and segment-local document ID {@code i}; gaps, reordered + * documents, and multiple vector fields are rejected by the writer. The caller is also responsible + * for ensuring that every raw vector component is finite; unlike ordinary Lucene field ingestion, + * this path deliberately does not scan the payload one float at a time. */ public final class ExternalFloat32Dataset { diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/MappedFbinDataset.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/MappedFbinDataset.java index 29a76d44b8..1ac7f27d2f 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/MappedFbinDataset.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/MappedFbinDataset.java @@ -4,12 +4,9 @@ */ package com.nvidia.cuvs.lucene; -import java.io.EOFException; import java.io.IOException; import java.lang.foreign.Arena; import java.lang.foreign.MemorySegment; -import java.nio.ByteBuffer; -import java.nio.ByteOrder; import java.nio.channels.FileChannel; import java.nio.file.Path; import java.nio.file.StandardOpenOption; @@ -24,38 +21,50 @@ final class MappedFbinDataset implements AutoCloseable { private boolean closed; static MappedFbinDataset map(Path path) throws IOException { - int rows; - int dimensions; - long payloadBytes; + FbinFileMetadata metadata = FbinFileMetadata.read(path); + return map(path, metadata, 0, metadata.rows()); + } + + static MappedFbinDataset map(Path path, int firstRow, int rowCount) throws IOException { + FbinFileMetadata metadata = FbinFileMetadata.read(path); + return map(path, metadata, firstRow, rowCount); + } + + private static MappedFbinDataset map( + Path path, FbinFileMetadata metadata, int firstRow, int rowCount) throws IOException { + if (firstRow < 0 || rowCount <= 0) { + throw new IllegalArgumentException("firstRow must be non-negative and rowCount positive"); + } + if ((long) firstRow + rowCount > metadata.rows()) { + throw new IllegalArgumentException( + "Requested row range [" + + firstRow + + ", " + + ((long) firstRow + rowCount) + + ") exceeds FBIN row count " + + metadata.rows()); + } + + long rowBytes = Math.multiplyExact((long) metadata.dimensions(), Float.BYTES); + long payloadOffset = Math.addExact(HEADER_BYTES, Math.multiplyExact((long) firstRow, rowBytes)); + long payloadBytes = Math.multiplyExact((long) rowCount, rowBytes); Arena arena = Arena.ofShared(); try { MemorySegment payload; try (FileChannel channel = FileChannel.open(path, StandardOpenOption.READ)) { - ByteBuffer header = ByteBuffer.allocate((int) HEADER_BYTES).order(ByteOrder.LITTLE_ENDIAN); - readFully(channel, header, 0L); - header.flip(); - rows = header.getInt(); - dimensions = header.getInt(); - if (rows <= 0 || dimensions <= 0) { + if (channel.size() != metadata.fileBytes()) { throw new IOException( - "FBIN header must contain positive rows and dimensions: " - + rows - + " x " - + dimensions); - } - payloadBytes = Math.multiplyExact(Math.multiplyExact((long) rows, dimensions), Float.BYTES); - long expectedBytes = Math.addExact(HEADER_BYTES, payloadBytes); - if (channel.size() != expectedBytes) { - throw new IOException( - "FBIN length " + "FBIN length changed before mapping: expected " + + metadata.fileBytes() + + " but got " + channel.size() - + " does not match header-derived length " - + expectedBytes); + + " for " + + path); } - payload = channel.map(FileChannel.MapMode.READ_ONLY, HEADER_BYTES, payloadBytes, arena); + payload = channel.map(FileChannel.MapMode.READ_ONLY, payloadOffset, payloadBytes, arena); } ExternalFloat32Dataset dataset = - ExternalFloat32Dataset.fromMemorySegment(payload, rows, dimensions); + ExternalFloat32Dataset.fromMemorySegment(payload, rowCount, metadata.dimensions()); return new MappedFbinDataset(arena, dataset); } catch (IOException | RuntimeException | Error failure) { arena.close(); @@ -63,21 +72,6 @@ static MappedFbinDataset map(Path path) throws IOException { } } - private static void readFully(FileChannel channel, ByteBuffer target, long position) - throws IOException { - long current = position; - while (target.hasRemaining()) { - int read = channel.read(target, current); - if (read < 0) { - throw new EOFException("Unexpected EOF while reading FBIN header"); - } - if (read == 0) { - throw new IOException("Unable to make progress while reading FBIN header"); - } - current += read; - } - } - private MappedFbinDataset(Arena arena, ExternalFloat32Dataset dataset) { this.arena = arena; this.dataset = dataset; diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBulkFbinMultiSegmentStorage.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBulkFbinMultiSegmentStorage.java new file mode 100644 index 0000000000..c268575bbd --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBulkFbinMultiSegmentStorage.java @@ -0,0 +1,396 @@ +/* + * 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.index.VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT; + +import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; +import com.nvidia.cuvs.spi.CuVSProvider; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.channels.FileChannel; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.security.MessageDigest; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HexFormat; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; +import org.apache.lucene.document.StoredField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.LeafReader; +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.SuppressFileSystems; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.junit.After; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Test; + +/** End-to-end coverage for merge-free mapped and external multi-segment bulk builds. */ +@SuppressFileSystems("*") +@SuppressSysoutChecks(bugUrl = "") +public class TestBulkFbinMultiSegmentStorage extends LuceneTestCase { + + private static final String ID_FIELD = "id"; + private static final String CALLBACK_ID_FIELD = "callback_id"; + private static final String VECTOR_FIELD = "vector"; + private static final int ROWS = 302; + private static final int DIMENSIONS = 32; + private static final int SEGMENTS = 3; + + @BeforeClass + public static void beforeClass() { + try { + CuVSProvider.provider().enableRMMAsyncMemory(); + } catch (UnsupportedOperationException unsupported) { + assumeTrue("cuVS not supported: " + unsupported.getMessage(), false); + } + } + + @Before + public void requireCuvs() { + assumeTrue("cuVS not supported", isSupported()); + } + + @After + public void clearRegistry() { + ExternalFbinFileRegistry.clearForTests(); + } + + @Test + public void testMappedBuildPreservesSlicesAndGlobalIds() throws Exception { + Path root = createTempDir("mapped-multi-segment"); + float[][] vectors = createVectors(); + Path fbin = writeFbin(root, vectors); + Path index = root.resolve("index"); + List callbackIds = Collections.synchronizedList(new ArrayList<>()); + + CagraHnswBulkIndexWriter.indexMappedFbin( + fbin, config(index, new CagraHnswBuildMetrics()), callback(callbackIds)); + + assertEquals(expectedIds(), callbackIds); + assertEquals(SEGMENTS, countExtension(index, ".vec")); + assertEquals(SEGMENTS, countExtension(index, ".vemf")); + assertEquals(0, countExtension(index, ".vefr")); + Files.delete(fbin); + assertIndex(index, vectors); + } + + @Test + public void testMappedOverlappedBuildPublishesCompleteIndex() throws Exception { + Path root = createTempDir("mapped-overlapped-multi-segment"); + float[][] vectors = createVectors(); + Path fbin = writeFbin(root, vectors); + Path index = root.resolve("index"); + List callbackIds = Collections.synchronizedList(new ArrayList<>()); + + CagraHnswBulkIndexWriter.indexMappedFbin( + fbin, config(index, new CagraHnswBuildMetrics(), true), callback(callbackIds)); + + List sortedCallbackIds = new ArrayList<>(callbackIds); + Collections.sort(sortedCallbackIds); + assertEquals(expectedIds(), sortedCallbackIds); + Files.delete(fbin); + assertIndex(index, vectors); + } + + @Test + public void testExternalBuildVerifiesOnceAndPreservesSlices() throws Exception { + Path root = createTempDir("external-multi-segment"); + float[][] vectors = createVectors(); + Path fbin = writeFbin(root, vectors); + Path index = root.resolve("index"); + String digest = sha256(fbin); + CagraHnswBuildMetrics metrics = new CagraHnswBuildMetrics(); + List callbackIds = Collections.synchronizedList(new ArrayList<>()); + + try (ExternalFbinFileRegistry.Registration registration = + ExternalFbinFileRegistry.register(fbin, digest)) { + CagraHnswBulkIndexWriter.indexImmutableFbin( + registration, + config(index, metrics), + ExternalFbinOptions.verifySha256(), + callback(callbackIds)); + + assertEquals(expectedIds(), callbackIds); + assertEquals(SEGMENTS, countExtension(index, ".vefr")); + assertEquals(0, countExtension(index, ".vec")); + assertEquals(0, countExtension(index, ".vemf")); + assertIndex(index, vectors); + } + + Map snapshot = metrics.snapshot(); + assertEquals(1L, snapshot.get("stage/external fbin SHA-256 [DISK+CPU]/count").longValue()); + assertEquals( + Files.size(fbin), snapshot.get("stage/external fbin SHA-256 [DISK+CPU]/bytes").longValue()); + assertEquals(SEGMENTS, snapshot.get("stage/cagra-build [GPU]/count").longValue()); + } + + @Test + public void testExternalOverlappedBuildPrefetchesOnceAndPublishesCompleteIndex() + throws Exception { + Path root = createTempDir("external-overlapped-multi-segment"); + float[][] vectors = createVectors(); + Path fbin = writeFbin(root, vectors); + Path index = root.resolve("index"); + String digest = sha256(fbin); + CagraHnswBuildMetrics metrics = new CagraHnswBuildMetrics(); + List callbackIds = Collections.synchronizedList(new ArrayList<>()); + + try (ExternalFbinFileRegistry.Registration registration = + ExternalFbinFileRegistry.register(fbin, digest)) { + CagraHnswBulkIndexWriter.indexImmutableFbin( + registration, + config(index, metrics, true), + new ExternalFbinOptions(ExternalFbinBuildValidation.PREFETCH, 0L), + callback(callbackIds)); + + List sortedCallbackIds = new ArrayList<>(callbackIds); + Collections.sort(sortedCallbackIds); + assertEquals(expectedIds(), sortedCallbackIds); + assertIndex(index, vectors); + } + + Map snapshot = metrics.snapshot(); + assertEquals(1L, snapshot.get("stage/external fbin prefetch [DISK]/count").longValue()); + assertEquals( + Files.size(fbin) - ExternalFbinReference.HEADER_BYTES, + snapshot.get("stage/external fbin prefetch [DISK]/bytes").longValue()); + assertEquals(SEGMENTS, snapshot.get("stage/cagra-build [GPU]/count").longValue()); + } + + @Test + public void testExternalOverlappedBuildPreservesCallerInterruption() throws Exception { + Path root = createTempDir("external-overlapped-interruption"); + Path fbin = writeFbin(root, createVectors()); + Path index = root.resolve("index"); + String digest = sha256(fbin); + CountDownLatch callbackStarted = new CountDownLatch(1); + CountDownLatch releaseCallback = new CountDownLatch(1); + AtomicReference failure = new AtomicReference<>(); + AtomicBoolean interruptPreserved = new AtomicBoolean(); + + try (ExternalFbinFileRegistry.Registration registration = + ExternalFbinFileRegistry.register(fbin, digest)) { + Thread buildThread = + Thread.ofPlatform() + .start( + () -> { + try { + CagraHnswBulkIndexWriter.indexImmutableFbin( + registration, + config(index, new CagraHnswBuildMetrics(), true), + new ExternalFbinOptions(ExternalFbinBuildValidation.PREFETCH, 0L), + (document, id) -> { + callbackStarted.countDown(); + try { + releaseCallback.await(); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + throw new IOException("callback interrupted", interrupted); + } + }); + } catch (Throwable thrown) { + failure.set(thrown); + interruptPreserved.set(Thread.currentThread().isInterrupted()); + } + }); + + assertTrue(callbackStarted.await(10, java.util.concurrent.TimeUnit.SECONDS)); + buildThread.interrupt(); + buildThread.join(10_000L); + releaseCallback.countDown(); + assertFalse("interrupted build did not terminate", buildThread.isAlive()); + } finally { + releaseCallback.countDown(); + } + + assertTrue(failure.get() instanceof java.io.InterruptedIOException); + assertTrue(interruptPreserved.get()); + try (Directory directory = FSDirectory.open(index)) { + assertFalse(DirectoryReader.indexExists(directory)); + } + } + + @Test + public void testExternalDigestFailureDoesNotPublishPartialIndex() throws Exception { + Path root = createTempDir("external-multi-segment-failure"); + Path fbin = writeFbin(root, createVectors()); + Path index = root.resolve("index"); + String digest = sha256(fbin); + corruptMiddleSlice(fbin); + + try (ExternalFbinFileRegistry.Registration registration = + ExternalFbinFileRegistry.register(fbin, digest)) { + IOException mismatch = + expectThrows( + IOException.class, + () -> + CagraHnswBulkIndexWriter.indexImmutableFbin( + registration, + config(index, new CagraHnswBuildMetrics()), + ExternalFbinOptions.verifySha256())); + assertTrue(mismatch.getMessage(), mismatch.getMessage().contains("SHA-256 mismatch")); + } + + try (Directory directory = FSDirectory.open(index)) { + assertFalse(DirectoryReader.indexExists(directory)); + } + try (var children = Files.list(root)) { + assertFalse( + children.anyMatch(path -> path.getFileName().toString().startsWith("index.bulk-"))); + } + } + + @Test + public void testLaterSliceFailureDoesNotPublishPartialIndex() throws Exception { + Path root = createTempDir("mapped-multi-segment-failure"); + Path fbin = writeFbin(root, createVectors()); + Path index = root.resolve("index"); + + IOException failure = + expectThrows( + IOException.class, + () -> + CagraHnswBulkIndexWriter.indexMappedFbin( + fbin, + config(index, new CagraHnswBuildMetrics()), + (document, id) -> { + if (id == 150) { + throw new IOException("injected callback failure"); + } + })); + assertEquals("injected callback failure", failure.getMessage()); + + try (Directory directory = FSDirectory.open(index)) { + assertFalse(DirectoryReader.indexExists(directory)); + } + try (var children = Files.list(root)) { + assertFalse( + children.anyMatch(path -> path.getFileName().toString().startsWith("index.bulk-"))); + } + } + + private static CagraHnswBulkIndexWriter.FieldCallback callback(List callbackIds) { + return (document, id) -> { + callbackIds.add(id); + document.add(new StoredField(CALLBACK_ID_FIELD, id)); + }; + } + + private static CagraHnswBulkIndexWriter.Config config(Path index, CagraHnswBuildMetrics metrics) { + return config(index, metrics, false); + } + + private static CagraHnswBulkIndexWriter.Config config( + Path index, CagraHnswBuildMetrics metrics, boolean overlapped) { + AcceleratedHNSWParams graphBuild = + new AcceleratedHNSWParams.Builder() + .withCuvsDistanceType(CuvsDistanceType.InnerProduct) + .build(); + return CagraHnswBulkIndexWriter.Config.builder() + .field(VECTOR_FIELD, DIMENSIONS, MAXIMUM_INNER_PRODUCT) + .idField(ID_FIELD) + .graphBuild(graphBuild) + .segments(SEGMENTS, overlapped) + .targetDirectory(index) + .metrics(metrics) + .build(); + } + + private static void assertIndex(Path index, float[][] expectedVectors) throws Exception { + int documents = 0; + try (Directory directory = FSDirectory.open(index); + DirectoryReader reader = DirectoryReader.open(directory)) { + assertEquals(SEGMENTS, reader.leaves().size()); + for (var context : reader.leaves()) { + LeafReader leaf = context.reader(); + FloatVectorValues values = leaf.getFloatVectorValues(VECTOR_FIELD); + assertNotNull(values); + assertEquals(leaf.maxDoc(), values.size()); + for (int localDoc = 0; localDoc < leaf.maxDoc(); localDoc++) { + var stored = leaf.storedFields().document(localDoc); + int globalId = Integer.parseInt(stored.get(ID_FIELD)); + assertEquals(globalId, context.docBase + localDoc); + assertEquals(globalId, stored.getField(CALLBACK_ID_FIELD).numericValue().intValue()); + assertArrayEquals(expectedVectors[globalId], values.vectorValue(localDoc), 0.0f); + documents++; + } + leaf.checkIntegrity(); + } + assertEquals(ROWS, documents); + + float[] query = new float[DIMENSIONS]; + query[0] = 1.0f; + IndexSearcher searcher = new IndexSearcher(reader); + TopDocs hits = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, query, 10), 10); + assertEquals(10, hits.scoreDocs.length); + assertEquals( + Integer.toString(ROWS - 1), + searcher.storedFields().document(hits.scoreDocs[0].doc).get(ID_FIELD)); + } + } + + private static float[][] createVectors() { + float[][] vectors = new float[ROWS][DIMENSIONS]; + Random random = new Random(8675309L); + for (int row = 0; row < vectors.length; row++) { + for (int dimension = 0; dimension < DIMENSIONS; dimension++) { + vectors[row][dimension] = random.nextFloat() * 2.0f - 1.0f; + } + } + vectors[ROWS - 1][0] = 100.0f; + return vectors; + } + + private static List expectedIds() { + List ids = new ArrayList<>(ROWS); + for (int id = 0; id < ROWS; id++) { + ids.add(id); + } + return ids; + } + + private static Path writeFbin(Path root, float[][] vectors) throws IOException { + Path fbin = root.resolve("vectors.fbin"); + TestUtils.writeFbin(fbin, vectors); + return fbin; + } + + private static String sha256(Path file) throws Exception { + MessageDigest digest = MessageDigest.getInstance("SHA-256"); + digest.update(Files.readAllBytes(file)); + return HexFormat.of().formatHex(digest.digest()); + } + + private static void corruptMiddleSlice(Path fbin) throws IOException { + long rowBytes = (long) DIMENSIONS * Float.BYTES; + long position = ExternalFbinReference.HEADER_BYTES + rowBytes * (ROWS / 2L); + try (FileChannel channel = FileChannel.open(fbin, StandardOpenOption.WRITE)) { + channel.write(ByteBuffer.wrap(new byte[] {42}), position); + } + } + + private static long countExtension(Path index, String extension) throws IOException { + try (Directory directory = FSDirectory.open(index)) { + return Arrays.stream(directory.listAll()).filter(name -> name.endsWith(extension)).count(); + } + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBuildMetricsConcurrency.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBuildMetricsConcurrency.java index 8437e69dff..79f6c6ba9c 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBuildMetricsConcurrency.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBuildMetricsConcurrency.java @@ -25,18 +25,14 @@ public void testGaugeAllowsRepeatedIdenticalValues() { } @Test - public void testGaugeRejectsConflictingValues() { + public void testGaugeReportsRangeForDifferentSegmentValues() { CagraHnswBuildMetrics metrics = new CagraHnswBuildMetrics(); metrics.setGauge("effective graph degree", 32L); + metrics.setGauge("effective graph degree", 56L); - IllegalStateException failure = - expectThrows( - IllegalStateException.class, () -> metrics.setGauge("effective graph degree", 56L)); - - assertTrue(failure.getMessage(), failure.getMessage().contains("effective graph degree")); - assertTrue(failure.getMessage(), failure.getMessage().contains("32")); - assertTrue(failure.getMessage(), failure.getMessage().contains("56")); - assertEquals(32L, metrics.snapshot().get("gauge/effective graph degree").longValue()); + assertNull(metrics.snapshot().get("gauge/effective graph degree")); + assertEquals(32L, metrics.snapshot().get("gauge/effective graph degree/min").longValue()); + assertEquals(56L, metrics.snapshot().get("gauge/effective graph degree/max").longValue()); } @Test diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMappedFbinDataset.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMappedFbinDataset.java index 73a742936c..fcdfd4aee3 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMappedFbinDataset.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMappedFbinDataset.java @@ -32,6 +32,33 @@ public void testMapsDenseFloatPayloadWithoutCopying() throws Exception { } } + @Test + public void testMapsContiguousRowSlice() throws Exception { + Path file = createTempDir().resolve("vectors.fbin"); + writeFbin(file, 4, 3, new float[] {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); + + try (MappedFbinDataset mapped = MappedFbinDataset.map(file, 1, 2)) { + assertEquals(2, mapped.rows()); + assertEquals(3, mapped.dimensions()); + float[] first = new float[3]; + float[] last = new float[3]; + mapped.dataset().matrix().getRow(0).toArray(first); + mapped.dataset().matrix().getRow(1).toArray(last); + assertArrayEquals(new float[] {4, 5, 6}, first, 0.0f); + assertArrayEquals(new float[] {7, 8, 9}, last, 0.0f); + } + } + + @Test + public void testRejectsInvalidRowSlice() throws Exception { + Path file = createTempDir().resolve("vectors.fbin"); + writeFbin(file, 2, 3, new float[] {1, 2, 3, 4, 5, 6}); + + assertThrows(IllegalArgumentException.class, () -> MappedFbinDataset.map(file, -1, 1)); + assertThrows(IllegalArgumentException.class, () -> MappedFbinDataset.map(file, 0, 0)); + assertThrows(IllegalArgumentException.class, () -> MappedFbinDataset.map(file, 1, 2)); + } + @Test public void testRejectsTrailingOrTruncatedPayload() throws Exception { Path trailing = createTempDir().resolve("trailing.fbin"); From fcdb504034b2316dc9251da04c236584036857a5 Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Fri, 11 Sep 2026 13:38:09 +0000 Subject: [PATCH 3/3] Harden CAGRA HNSW FBIN bulk indexing - validate bulk-build ownership, counts, and external FBIN descriptors - make mapped ingestion, cancellation, and native cleanup failure-safe - report build metrics accurately and strengthen GPU regression coverage --- ci/test_lucene.sh | 27 ++ java/cuvs-lucene/README.md | 125 +++++-- .../cuvs/lucene/AcceleratedHNSWUtils.java | 34 +- .../lucene/AcceleratedHnswGraphOutput.java | 167 +++++----- .../BorrowedDatasetHnswVectorsWriter.java | 12 +- .../cuvs/lucene/CagraHnswBuildMetrics.java | 10 +- .../cuvs/lucene/CagraHnswBulkIndexWriter.java | 246 +++++++------- .../cuvs/lucene/CuVS2510GPUVectorsReader.java | 16 +- .../cuvs/lucene/ExternalFbinFileRegistry.java | 31 +- .../lucene/ExternalFbinFlatVectorsReader.java | 27 +- .../nvidia/cuvs/lucene/ExternalFbinIO.java | 16 +- .../lucene/ExternalFbinMappedInputCache.java | 146 +++++++++ .../cuvs/lucene/ExternalFbinReference.java | 24 +- .../cuvs/lucene/ExternalFloat32Dataset.java | 12 +- .../lucene/ImmutableExternalFbinDataset.java | 8 +- .../lucene/MappedSelfContainedOutput.java | 159 ++++++--- .../NativeFlatBufferedHNSWVectorsWriter.java | 5 +- .../cuvs/lucene/NativeFlatVectorsWriter.java | 24 +- .../nvidia/cuvs/lucene/CuVSTestSupport.java | 37 +++ .../TestAcceleratedThreeLayerRoundTrip.java | 216 +++++++++++++ .../TestBulkExternalFbinApiSurface.java | 15 + .../cuvs/lucene/TestBulkFbinIndexStorage.java | 12 +- .../TestBulkFbinMultiSegmentStorage.java | 45 ++- .../TestCagraHnswBuildMetricsConcurrency.java | 19 +- .../lucene/TestCagraHnswBulkIndexWriter.java | 29 +- ...estCagraHnswBulkIndexWriterValidation.java | 216 +++++++++++++ ...ializationAndSearchWithFallbackWriter.java | 3 - ...510GPUVectorsReaderScoreNormalization.java | 41 --- .../TestExternalFbinFlatVectorsReader.java | 304 +++++++++++++++++- ...MappedSelfContainedOutputCancellation.java | 190 +++++++++++ .../lucene/TestNativeFlatBufferedCleanup.java | 114 +++++++ .../TestSingleVectorHnswGraphCleanup.java | 77 +++++ 32 files changed, 1968 insertions(+), 439 deletions(-) create mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinMappedInputCache.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/CuVSTestSupport.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedThreeLayerRoundTrip.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBulkIndexWriterValidation.java delete mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVS2510GPUVectorsReaderScoreNormalization.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMappedSelfContainedOutputCancellation.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestSingleVectorHnswGraphCleanup.java diff --git a/ci/test_lucene.sh b/ci/test_lucene.sh index 5478a834ae..32df32b9e6 100755 --- a/ci/test_lucene.sh +++ b/ci/test_lucene.sh @@ -15,11 +15,38 @@ rapids-logger "Run cuvs-lucene build and tests" RAPIDS_CUDA_MAJOR="${RAPIDS_CUDA_VERSION%%.*}" export RAPIDS_CUDA_MAJOR +# This CI entry point requires critical GPU suites to execute; do not allow JNI, allocator, or +# native-library failures to turn the job into an all-skipped success. +export CUVS_LUCENE_REQUIRE_GPU_TESTS=1 # Forwards the name of the cuvs-java artifact given by the workflow. # TODO: switch to installing pre-built artifacts instead of rebuilding in test jobs # ref: https://github.com/rapidsai/cuvs/issues/868 ci/build_lucene.sh "$@" --run-java-tests +rapids-logger "Verify required cuvs-lucene GPU suites executed without skips" +REQUIRED_GPU_TESTS=( + TestCagraHnswBulkIndexWriter + TestBulkFbinIndexStorage + TestBulkFbinMultiSegmentStorage + TestAcceleratedThreeLayerRoundTrip +) +for test_class in "${REQUIRED_GPU_TESTS[@]}"; do + report="java/cuvs-lucene/target/surefire-reports/TEST-com.nvidia.cuvs.lucene.${test_class}.xml" + if [[ ! -f "${report}" ]]; then + rapids-logger "Missing required GPU test report: ${report}" + EXITCODE=1 + continue + fi + if ! grep -Eq ']*tests="[1-9][0-9]*"' "${report}"; then + rapids-logger "Required GPU suite executed no tests: ${test_class}" + EXITCODE=1 + fi + if grep -Eq ']*skipped="[1-9][0-9]*"' "${report}"; then + rapids-logger "Required GPU suite skipped tests: ${test_class}" + EXITCODE=1 + fi +done + rapids-logger "Test script exiting with value: $EXITCODE" exit ${EXITCODE} diff --git a/java/cuvs-lucene/README.md b/java/cuvs-lucene/README.md index 0d80f08953..6179c7d249 100644 --- a/java/cuvs-lucene/README.md +++ b/java/cuvs-lucene/README.md @@ -17,7 +17,12 @@ This is a project for using [cuVS](https://github.com/rapidsai/cuvs), NVIDIA's G Four codecs are currently provided: -- `Lucene101AcceleratedHNSWCodec` — GPU-accelerated HNSW build with CPU HNSW search. The on-disk format is standard Lucene HNSW, so indexes built on the GPU can be read by any stock Lucene 10.x reader. +- `Lucene101AcceleratedHNSWCodec` — GPU-accelerated HNSW build with CPU HNSW search. + Self-contained builds write standard Lucene HNSW and flat-vector data, so their indexes can be + read by a stock Lucene 10.x reader. Indexes built with the opt-in [immutable external-FBIN bulk + mode](#ownership-and-sharing) are not self-contained or readable by stock Lucene alone; they + require cuvs-lucene and a live registration lease for the exact FBIN throughout every reader + lifetime. - `LuceneAcceleratedHNSWScalarQuantizedCodec` — scalar-quantized vectors for a smaller index footprint. - `LuceneAcceleratedHNSWBinaryQuantizedCodec` — binary-quantized vectors for an even smaller index footprint. - `CuVS2510GPUSearchCodec` — GPU-accelerated HNSW build and GPU search @@ -161,9 +166,11 @@ The streaming/native-buffered mode has three entry shapes: Both mapped modes support one or more contiguous, independently built segments. With overlap disabled, slices are built sequentially; with overlap enabled, host-side preparation can run in a -bounded pipeline while GPU commits remain serialized. The slices are published together without a -Lucene vector merge. They require exactly one dense `FLOAT32` vector field, one vector per document, -and preserve each source row's absolute ID across segment-local document ordinals. +bounded pipeline while each per-segment commit/close region (including its GPU work) remains +serialized. The slices are published together without a Lucene vector merge. They require exactly +one dense `FLOAT32` vector field, one vector per document, and preserve each source row's absolute +ID across segment-local document ordinals. The configured segment count is exact and cannot exceed +the source row count. ### Ownership and sharing @@ -175,6 +182,14 @@ publish the final target until every slice succeeds. Streaming sequential builds so they can already have committed earlier segments when a later segment fails; callers must treat and replace that target as a failed build. +Every one-shot build exclusively owns and recreates its configured target as the final Lucene index +directory. The target is stored as an absolute normalized path, and each file-backed source is +resolved once to its absolute real path before work begins. These path rules do not provide a +cross-process lock: the application must prevent concurrent use of the target. A mapped source must +remain unchanged at that resolved identity for the whole synchronous build—do not write, truncate, +replace, move, or delete it until the method returns. The read-only mapping prevents writes through +this API but cannot stop another process from mutating the file. + The generic, streaming, and mapped self-contained results can be moved by copying the Lucene index directory. The immutable external-FBIN result must be shipped as two artifacts: the Lucene index directory and the exact FBIN identified by the descriptor's complete-file SHA-256. The descriptor @@ -190,13 +205,75 @@ try (ExternalFbinFileRegistry.Registration lease = } ``` -Registrations are process-local and reference counted. Keep at least one lease open for the entire -lifetime of every reader using that content ID. Do not modify, replace, relocate, or delete the -registered FBIN while a build or reader is alive. An index directory copied without its FBIN, or a -process that has not registered the local FBIN, cannot be opened. Registering a different path for -the same content ID while an existing registration is live is rejected. External-FBIN indexes also -require a cuvs-lucene runtime that understands their external-vector descriptor; they are not -readable by a stock Lucene runtime alone. +Registrations are local to the defining cuvs-lucene classloader (normally the application +classloader) and reference counted. Keep at least one lease open for the entire lifetime of every +reader using that content ID; it need not be the same lease used for the build. Do not modify, +replace, relocate, or delete the registered FBIN while a build or reader is alive. An index +directory copied without its FBIN, or an application/classloader that has not registered the local +FBIN, cannot be opened. Registering a different path for the same content ID while an existing +registration is live is rejected. External-FBIN indexes also require a cuvs-lucene runtime that +understands their external-vector descriptor; they are not readable by a stock Lucene runtime +alone. + +### Build and deploy examples + +A mapped build is self-contained after the synchronous call succeeds, so its source may then be +removed. The source must remain immutable until that point: + +```java +Path sourceFbin = Path.of("vectors.fbin"); +Path indexPath = Path.of("index-mapped"); +AcceleratedHNSWParams graph = new AcceleratedHNSWParams.Builder().build(); +CagraHnswBulkIndexWriter.Config mapped = + CagraHnswBulkIndexWriter.Config.builder() + .field("vector", 768, VectorSimilarityFunction.EUCLIDEAN) + .graphBuild(graph) + .targetDirectory(indexPath) + .segments(1, false) + .build(); + +CagraHnswBulkIndexWriter.indexMappedFbin(sourceFbin, mapped); +Files.delete(sourceFbin); // safe only after a successful return +try (Directory directory = FSDirectory.open(indexPath); + DirectoryReader reader = DirectoryReader.open(directory)) { + // Search the self-contained index. +} +``` + +An external build and every later reader need a live registration, but they may use different +leases and local paths. Deploy or relocate both artifacts only after all current readers and leases +are closed: + +```java +Path buildFbin = Path.of("vectors.fbin"); +Path buildIndex = Path.of("index-external"); +String sha256Hex = establishedCompleteFileSha256; +CagraHnswBulkIndexWriter.Config external = + CagraHnswBulkIndexWriter.Config.builder() + .field("vector", 768, VectorSimilarityFunction.EUCLIDEAN) + .graphBuild(new AcceleratedHNSWParams.Builder().build()) + .targetDirectory(buildIndex) + .segments(4, true) + .build(); + +try (ExternalFbinFileRegistry.Registration buildLease = + ExternalFbinFileRegistry.register(buildFbin, sha256Hex)) { + CagraHnswBulkIndexWriter.indexImmutableFbin( + buildLease, + external, + new ExternalFbinOptions(ExternalFbinBuildValidation.VERIFY_SHA256, 0L)); +} + +// Copy/move buildIndex and buildFbin together after closing their users. +Path deployedIndex = Path.of("/srv/search/index"); +Path deployedFbin = Path.of("/srv/search/data/vectors.fbin"); +try (ExternalFbinFileRegistry.Registration readerLease = + ExternalFbinFileRegistry.register(deployedFbin, sha256Hex); + Directory directory = FSDirectory.open(deployedIndex); + DirectoryReader reader = DirectoryReader.open(directory)) { + // Search, and optionally call reader.leaves().getFirst().reader().checkIntegrity(). +} +``` ### Validation @@ -235,22 +312,26 @@ this cost does not occur during normal reader opening or search. ### Metrics -Supply one `CagraHnswBuildMetrics` per build through `Config.Builder.metrics`. After the build, -`snapshot()` returns a stable `Map`, and `appendTo(target, prefix)` adds that -snapshot to a benchmark result map. +`CagraHnswBuildMetrics` is a caller-owned, thread-safe cumulative accumulator. Supply a fresh +instance through `Config.Builder.metrics` for each logical build when per-build measurements are +required. Reusing either the accumulator or its `Config`, including for a retry, intentionally +aggregates later measurements into the same snapshot. `snapshot()` returns a stable +`Map`, and `appendTo(target, prefix)` adds that snapshot to a benchmark result map. Metric keys use three namespaces: - `stage//seconds`, `stage//count`, and optional `stage//bytes` report aggregated graph build, graph conversion and output, mapped flat output, external scan/overlap, and bulk - writer commit/close stages. The commit wall includes the flush and its nested GPU/output work; - nested stage durations must not be added to it. -- `counter/` reports values such as logical adjacency bytes, mapped flat chunks, the number - of segments using each effective CAGRA algorithm, and external scan progress at CAGRA start and - end. -- `gauge/` records effective graph-build parameters, including graph degrees, writer threads, - NN-descent iterations, and IVF-PQ dimensions, lists, probes, and k-means iterations when used. If - independent segments resolve different heuristic values, the scalar key is replaced by explicit + writer commit/close stages. `base cagra-build [GPU]` covers only the base graph; upper-layer + CAGRA builds are included in `hnsw-convert [GPU+PCIe+CPU]`. The commit wall includes the flush + and its nested GPU/output work, so nested stage durations must not be added to it. +- `counter/` reports values such as actual base-graph adjacency bytes, mapped flat chunks, + the number of segments using each selected CAGRA algorithm, and external scan progress at CAGRA + start and end. +- `gauge/` records selected graph-build inputs, including graph degrees, writer threads, + NN-descent iterations, and IVF-PQ dimensions, lists, probes, and k-means iterations when used. + The separately reported `effective graph degree` comes from the returned base graph itself. If + independent segments record different values, the scalar key is replaced by explicit `gauge//min` and `gauge//max` keys. ### Why these controls are bulk-only 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 9cf30f02c2..80bf7c8807 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 @@ -64,19 +64,29 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens // Create adjacency list for single node with no neighbors int[][] singleNodeAdjacency = new int[][] {{-1}}; // -1 indicates no neighbors - // Create CuVSMatrix from the adjacency list - CuVSMatrix adjacencyMatrix = CuVSMatrix.ofArray(singleNodeAdjacency); - - // Create layer data for single-level graph - List layerNodes = new ArrayList<>(); - List layerAdjacencies = new ArrayList<>(); - - // Layer 0: contains all nodes (just the single node) - layerNodes.add(null); // Layer 0 contains all nodes, so we don't need to store node list - layerAdjacencies.add(adjacencyMatrix); + return createSingleVectorHnswGraph(size, dimensions, CuVSMatrix.ofArray(singleNodeAdjacency)); + } - // Create the single-layer graph - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies, 1); + /** + * Creates the single-vector graph from, and takes ownership of, {@code adjacencyMatrix}. + * + *

    {@link GPUBuiltHnswGraph}'s constructor synchronously materializes every adjacency row into + * heap-backed neighbor arrays, so the matrix can be closed before the graph is returned. + */ + static GPUBuiltHnswGraph createSingleVectorHnswGraph( + int size, int dimensions, CuVSMatrix adjacencyMatrix) throws Throwable { + try (adjacencyMatrix) { + // Create layer data for single-level graph + List layerNodes = new ArrayList<>(); + List layerAdjacencies = new ArrayList<>(); + + // Layer 0: contains all nodes (just the single node) + layerNodes.add(null); // Layer 0 contains all nodes, so we don't need to store node list + layerAdjacencies.add(adjacencyMatrix); + + // Create the single-layer graph + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies, 1); + } } /** diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHnswGraphOutput.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHnswGraphOutput.java index 60dc3aef6d..3dbb1702d0 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHnswGraphOutput.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHnswGraphOutput.java @@ -137,91 +137,108 @@ private void writeField(FieldInfo fieldInfo, CuVSMatrix dataset, boolean closeDa throws IOException { long graphOutputStart = CagraHnswBuildMetrics.start(); try { - 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()); - recordEffectiveParams(params); - long adjacencyBytes = - Math.multiplyExact( - Math.multiplyExact((long) size, params.getGraphDegree()), Integer.BYTES); - metrics.addCounter("logical cagra adjacency bytes", adjacencyBytes); - long stageStart = CagraHnswBuildMetrics.start(); - try (CagraIndex cagraIndex = - CagraIndex.newBuilder(getCuVSResourcesInstance()) - .withDataset(dataset) - .withIndexParams(params) - .build()) { - metrics.stop("cagra-build [GPU]", stageStart); - stageStart = CagraHnswBuildMetrics.start(); - CuVSMatrix adjacencyListMatrix = cagraIndex.getGraph(); - metrics.stop("graph-view [GPU]", stageStart); - int dimensions = fieldInfo.getVectorDimension(); - stageStart = CagraHnswBuildMetrics.start(); - GPUBuiltHnswGraph hnswGraph = - createMultiLayerHnswGraph( - fieldInfo, - dimensions, - adjacencyListMatrix, - dataset, - acceleratedHNSWParams.getHnswLayers(), - params, - QuantizationType.NONE, - acceleratedHNSWParams.getWriterThreads(), - acceleratedHNSWParams.getHnswLayerSeed()); - metrics.stop("hnsw-convert [PCIe+CPU]", stageStart, adjacencyBytes); - long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - stageStart = CagraHnswBuildMetrics.start(); - int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); - long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; - metrics.stop("write-graph [CPU+DISK]", stageStart, vectorIndexLength); - writeMeta( - hnswVectorIndex, - hnswMeta, - fieldInfo, - vectorIndexOffset, - vectorIndexLength, - size, - hnswGraph, - graphLevelNodeOffsets); - } - } catch (Throwable t) { - Utils.handleThrowable(t); + if (closeDataset) { + withOwnedDataset(dataset, () -> writeUnownedField(fieldInfo, dataset)); + } else { + writeUnownedField(fieldInfo, dataset); } } finally { metrics.stop("graph-output wall [CPU+GPU+DISK]", graphOutputStart); - if (closeDataset) { - dataset.close(); + } + } + + static void withOwnedDataset(CuVSMatrix dataset, IOAction action) throws IOException { + try (dataset) { + action.run(); + } + } + + @FunctionalInterface + interface IOAction { + void run() throws IOException; + } + + private void writeUnownedField(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()); + recordSelectedParams(params); + long stageStart = CagraHnswBuildMetrics.start(); + try (CagraIndex cagraIndex = + CagraIndex.newBuilder(getCuVSResourcesInstance()) + .withDataset(dataset) + .withIndexParams(params) + .build()) { + metrics.stop("base cagra-build [GPU]", stageStart); + stageStart = CagraHnswBuildMetrics.start(); + CuVSMatrix adjacencyListMatrix = cagraIndex.getGraph(); + metrics.stop("base graph-view [GPU]", stageStart); + long adjacencyBytes = + Math.multiplyExact( + Math.multiplyExact(adjacencyListMatrix.size(), adjacencyListMatrix.columns()), + Integer.BYTES); + metrics.setGauge("effective graph degree", adjacencyListMatrix.columns()); + metrics.addCounter("actual cagra adjacency bytes", adjacencyBytes); + int dimensions = fieldInfo.getVectorDimension(); + stageStart = CagraHnswBuildMetrics.start(); + GPUBuiltHnswGraph hnswGraph = + createMultiLayerHnswGraph( + fieldInfo, + dimensions, + adjacencyListMatrix, + dataset, + acceleratedHNSWParams.getHnswLayers(), + params, + QuantizationType.NONE, + acceleratedHNSWParams.getWriterThreads(), + acceleratedHNSWParams.getHnswLayerSeed()); + // Upper HNSW layers build additional CAGRA graphs inside this conversion stage. + metrics.stop("hnsw-convert [GPU+PCIe+CPU]", stageStart, adjacencyBytes); + long vectorIndexOffset = hnswVectorIndex.getFilePointer(); + stageStart = CagraHnswBuildMetrics.start(); + int[][] graphLevelNodeOffsets = + writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; + metrics.stop("write-graph [CPU+DISK]", stageStart, vectorIndexLength); + writeMeta( + hnswVectorIndex, + hnswMeta, + fieldInfo, + vectorIndexOffset, + vectorIndexLength, + size, + hnswGraph, + graphLevelNodeOffsets); } + } catch (Throwable t) { + Utils.handleThrowable(t); } } - private void recordEffectiveParams(CagraIndexParams params) { + private void recordSelectedParams(CagraIndexParams params) { metrics.addCounter( - "effective cagra algorithm/" + params.getCagraGraphBuildAlgo() + " segments", 1L); - metrics.setGauge("effective graph degree", params.getGraphDegree()); - metrics.setGauge("effective intermediate graph degree", params.getIntermediateGraphDegree()); - metrics.setGauge("effective writer threads", params.getNumWriterThreads()); - metrics.setGauge("effective nn-descent iterations", params.getNNDescentNumIterations()); + "selected cagra algorithm/" + params.getCagraGraphBuildAlgo() + " segments", 1L); + metrics.setGauge("selected graph degree", params.getGraphDegree()); + metrics.setGauge("selected intermediate graph degree", params.getIntermediateGraphDegree()); + metrics.setGauge("selected writer threads", params.getNumWriterThreads()); + metrics.setGauge("selected nn-descent iterations", params.getNNDescentNumIterations()); if (params.getCagraGraphBuildAlgo() == CagraIndexParams.CagraGraphBuildAlgo.IVF_PQ) { var ivf = params.getCuVSIvfPqParams(); - metrics.setGauge("effective ivf-pq dimensions", ivf.getIndexParams().getPqDim()); - metrics.setGauge("effective ivf-pq lists", ivf.getIndexParams().getnLists()); - metrics.setGauge("effective ivf-pq probes", ivf.getSearchParams().getnProbes()); - metrics.setGauge( - "effective ivf-pq kmeans iterations", ivf.getIndexParams().getKmeansNIters()); + metrics.setGauge("selected ivf-pq dimensions", ivf.getIndexParams().getPqDim()); + metrics.setGauge("selected ivf-pq lists", ivf.getIndexParams().getnLists()); + metrics.setGauge("selected ivf-pq probes", ivf.getSearchParams().getnProbes()); + metrics.setGauge("selected ivf-pq kmeans iterations", ivf.getIndexParams().getKmeansNIters()); } } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/BorrowedDatasetHnswVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/BorrowedDatasetHnswVectorsWriter.java index 324e31ccee..ad8e6eccf7 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/BorrowedDatasetHnswVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/BorrowedDatasetHnswVectorsWriter.java @@ -5,8 +5,9 @@ package com.nvidia.cuvs.lucene; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.printInfoStream; -import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.closeCuVSResourcesInstance; +import static com.nvidia.cuvs.lucene.NativeFlatBufferedHNSWVectorsWriter.closeOutputsAndResources; +import java.io.Closeable; import java.io.IOException; import java.util.ArrayList; import java.util.List; @@ -55,7 +56,8 @@ final class BorrowedDatasetHnswVectorsWriter extends KnnVectorsWriter { graphOutput = newGraphOutput; vectorOutput = newVectorOutput; if (!success) { - IOUtils.closeWhileHandlingException(newVectorOutput, newGraphOutput); + Closeable resourceCloser = ThreadLocalCuVSResourcesProvider::closeCuVSResourcesInstance; + IOUtils.closeWhileHandlingException(newVectorOutput, newGraphOutput, resourceCloser); } } printInfoStream(infoStream, getClass().getSimpleName(), "borrowed FBIN writer initialized"); @@ -126,11 +128,7 @@ public void finish() throws IOException { @Override public void close() throws IOException { printInfoStream(infoStream, getClass().getSimpleName(), "closing resources"); - try { - IOUtils.close(vectorOutput, graphOutput); - } finally { - closeCuVSResourcesInstance(); - } + closeOutputsAndResources(vectorOutput, graphOutput); } @Override diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraHnswBuildMetrics.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraHnswBuildMetrics.java index 07a0fbef75..0127cfbd3a 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraHnswBuildMetrics.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraHnswBuildMetrics.java @@ -10,7 +10,13 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.LongAdder; -/** Per-build, thread-safe measurements for the bulk CAGRA-to-HNSW pipeline. */ +/** + * Thread-safe cumulative measurements for the bulk CAGRA-to-HNSW pipeline. + * + *

    An instance is never reset automatically. Reusing it, or a bulk-writer configuration that + * contains it, intentionally aggregates every invocation into the same snapshot. Supply a fresh + * instance for each logical build when per-build measurements are required. + */ public final class CagraHnswBuildMetrics { private final Map stages = new ConcurrentHashMap<>(); @@ -45,7 +51,7 @@ void addCounter(String name, long value) { counters.computeIfAbsent(name, ignored -> new LongAdder()).add(value); } - /** Records an effective configuration value, retaining its range across independent segments. */ + /** Records a gauge value, retaining its range across independent segments or invocations. */ void setGauge(String name, long value) { Objects.requireNonNull(name, "name"); gauges.compute( 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 index f1503acc06..e49d7553b8 100644 --- 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 @@ -118,6 +118,14 @@ private CagraHnswBulkIndexWriter( if (exactVectorCount <= 0) { throw new IllegalArgumentException("exactVectorCount must be > 0, got " + exactVectorCount); } + if (exactVectorCount > IndexWriter.MAX_DOCS) { + throw new IllegalArgumentException( + "exactVectorCount " + + exactVectorCount + + " exceeds Lucene's index-wide document limit " + + IndexWriter.MAX_DOCS + + "; reduce the input size"); + } if (conf.getIndexSort() != null) { throw new IllegalArgumentException( "CagraHnswBulkIndexWriter does not support an index-sorted segment (native flat" @@ -128,19 +136,26 @@ private CagraHnswBulkIndexWriter( this.exactVectorCount = exactVectorCount; Codec codec = new Lucene101AcceleratedHNSWCodec(config.graphBuildParams(), bulkContext); - 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()); + IndexWriterConfig ownedConf = ownedIndexWriterConfig(conf, codec, exactVectorCount); this.writer = new IndexWriter(directory, ownedConf); } + static IndexWriterConfig ownedIndexWriterConfig( + IndexWriterConfig conf, Codec codec, int exactVectorCount) { + return new IndexWriterConfig(conf.getAnalyzer()) + .setSimilarity(conf.getSimilarity()) + .setInfoStream(conf.getInfoStream()) + .setCodec(codec) + .setUseCompoundFile(false) + .setMaxBufferedDocs(Math.addExact(exactVectorCount, 1)) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) + .setMergePolicy(NoMergePolicy.INSTANCE) + // close() below is the only publication point. If its explicit commit fails, + // IndexWriter.close() must not retry and publish a commit while reporting failure. + .setCommitOnClose(false) + .setOpenMode(conf.getOpenMode()); + } + /** * Adds one document, exactly like {@link IndexWriter#addDocument}. {@code doc} must contain * exactly one vector field whose name, dimension, and similarity match {@code config}; that field @@ -304,6 +319,7 @@ public static void build(VectorSource source, Config config, FieldCallback callb throws Exception { Objects.requireNonNull(source, "source"); Objects.requireNonNull(config, "config"); + requireOneShotTarget(config); if (config.overlapped()) { throw new IllegalArgumentException( "Config.overlapped() is not supported by build(VectorSource, Config): a VectorSource is" @@ -325,7 +341,9 @@ public static void build(VectorSource source, Config config, FieldCallback callb /** * 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. + * {@link FbinVectorSource} with {@link #DEFAULT_CHUNK_SIZE_MB}-sized prefetched chunks. The + * caller must not mutate, replace, move, or delete the resolved source until this synchronous + * build returns. */ public static void indexFbin(Path fbinPath, Config config) throws Exception { indexFbin(fbinPath, config, null, DEFAULT_CHUNK_SIZE_MB); @@ -353,7 +371,9 @@ public static void indexFbin( Path fbinPath, Config config, FieldCallback callback, int chunkSizeMB) throws Exception { Objects.requireNonNull(fbinPath, "fbinPath"); Objects.requireNonNull(config, "config"); - FbinFileMetadata metadata = FbinFileMetadata.read(fbinPath); + requireOneShotTarget(config); + Path sourcePath = resolveSourcePath(fbinPath, config); + FbinFileMetadata metadata = FbinFileMetadata.read(sourcePath); int total = metadata.rows(); int dim = metadata.dimensions(); if (dim != config.dimensions()) { @@ -366,9 +386,23 @@ public static void indexFbin( } List slices = sliceEvenly(total, config.numSegments()); if (config.overlapped() && slices.size() > 1) { - buildOverlapped(fbinPath, config, callback, chunkSizeMB, slices, dim); + buildPartitionedSegments( + config, + slices, + (segmentDirectory, firstRow, rowCount, gpuPermit) -> { + float[] scratch = new float[dim]; + try (FbinVectorSource source = + new FbinVectorSource(sourcePath, firstRow, rowCount, chunkSizeMB); + Directory directory = FSDirectory.open(segmentDirectory)) { + buildSegment( + directory, source, scratch, config, callback, 0, firstRow, rowCount, true, + gpuPermit); + } + }, + null, + null); } else { - try (FbinVectorSource source = new FbinVectorSource(fbinPath, chunkSizeMB)) { + try (FbinVectorSource source = new FbinVectorSource(sourcePath, chunkSizeMB)) { buildSequential(source, config, callback, slices); } } @@ -378,6 +412,8 @@ public static void indexFbin( * Builds a normal, self-contained Lucene index while CAGRA and the flat-vector writer consume a * shared read-only mapping of {@code fbinPath}. This avoids decoding and copying every row during * document ingest while preserving a conventional {@code .vec} file in the finished index. + * The caller must not write, truncate, replace, move, or delete the resolved source file from + * method entry until this synchronous build returns. */ public static void indexMappedFbin(Path fbinPath, Config config) throws Exception { indexMappedFbin(fbinPath, config, null); @@ -388,16 +424,17 @@ public static void indexMappedFbin(Path fbinPath, Config config, FieldCallback c throws Exception { Objects.requireNonNull(fbinPath, "fbinPath"); Objects.requireNonNull(config, "config"); - requireMappedBuildTarget(config); + requireOneShotTarget(config); + Path sourcePath = resolveSourcePath(fbinPath, config); + FbinFileMetadata metadata = FbinFileMetadata.read(sourcePath); + validateMappedShape(metadata.rows(), metadata.dimensions(), config); + List slices = sliceEvenly(metadata.rows(), config.numSegments()); if (config.numSegments() > 1) { - FbinFileMetadata metadata = FbinFileMetadata.read(fbinPath); - validateMappedShape(metadata.rows(), metadata.dimensions(), config); - List slices = sliceEvenly(metadata.rows(), config.numSegments()); - buildBorrowedSegments( + buildPartitionedSegments( config, slices, (segmentDirectory, firstRow, rowCount, gpuPermit) -> { - try (MappedFbinDataset mapped = MappedFbinDataset.map(fbinPath, firstRow, rowCount)) { + try (MappedFbinDataset mapped = MappedFbinDataset.map(sourcePath, firstRow, rowCount)) { buildBorrowedSegment( segmentDirectory, config, @@ -411,7 +448,7 @@ public static void indexMappedFbin(Path fbinPath, Config config, FieldCallback c null); return; } - try (MappedFbinDataset mapped = MappedFbinDataset.map(fbinPath)) { + try (MappedFbinDataset mapped = MappedFbinDataset.map(sourcePath)) { validateMappedShape(mapped.rows(), mapped.dimensions(), config); buildBorrowedSegment( config.targetDirectory(), @@ -423,8 +460,9 @@ public static void indexMappedFbin(Path fbinPath, Config config, FieldCallback c /** * Builds a non-self-contained Lucene index whose exact-vector scoring reads from the immutable - * FBIN represented by {@code source}. The caller must retain {@code source} through every reader - * lifetime and replicate the referenced FBIN together with the Lucene directory. + * FBIN represented by {@code source}. The supplied lease must remain open for this build. It may + * then be closed, but at least one registration for the same content ID must span every reader + * lifetime. Deployments must replicate the referenced FBIN together with the Lucene directory. */ public static void indexImmutableFbin( ExternalFbinFileRegistry.Registration source, Config config, ExternalFbinOptions options) @@ -443,16 +481,17 @@ public static void indexImmutableFbin( Objects.requireNonNull(source, "source"); Objects.requireNonNull(config, "config"); Objects.requireNonNull(options, "options"); - requireMappedBuildTarget(config); - FbinFileMetadata metadata = FbinFileMetadata.read(source.path()); + requireOneShotTarget(config); + Path sourcePath = resolveSourcePath(source.path(), config); + FbinFileMetadata metadata = FbinFileMetadata.read(sourcePath); int rows = metadata.rows(); int dimensions = metadata.dimensions(); validateMappedShape(rows, dimensions, config); + List slices = sliceEvenly(rows, config.numSegments()); if (config.numSegments() > 1) { - List slices = sliceEvenly(rows, config.numSegments()); ExternalFbinOptions trustedOptions = new ExternalFbinOptions(ExternalFbinBuildValidation.TRUSTED_IMMUTABLE, 0L); - buildBorrowedSegments( + buildPartitionedSegments( config, slices, (segmentDirectory, firstRow, rowCount, gpuPermit) -> { @@ -479,12 +518,26 @@ public static void indexImmutableFbin( } } - private static void requireMappedBuildTarget(Config config) { + private static void requireOneShotTarget(Config config) { if (config.targetDirectory() == null) { throw new IllegalArgumentException("targetDirectory is required for a one-shot bulk build"); } } + private static Path resolveSourcePath(Path source, Config config) throws IOException { + Path realSource = source.toRealPath(); + Path target = config.targetDirectory(); + boolean samePath = realSource.equals(target); + if (!samePath && Files.exists(target)) { + samePath = Files.isSameFile(realSource, target); + } + if (samePath) { + throw new IllegalArgumentException( + "Bulk index target directory must not be the source FBIN: " + realSource); + } + return realSource; + } + private static void validateMappedShape(int rows, int dimensions, Config config) { if (rows <= 0 || dimensions != config.dimensions()) { throw new IllegalArgumentException( @@ -589,7 +642,7 @@ void build(Path segmentDirectory, int firstRow, int rowCount, Semaphore gpuPermi throws Exception; } - private static void buildBorrowedSegments( + private static void buildPartitionedSegments( Config config, List slices, BorrowedSegmentBuilder segmentBuilder, @@ -609,7 +662,7 @@ private static void buildBorrowedSegments( try { if (validationReference == null || validationOptions.validation() == ExternalFbinBuildValidation.TRUSTED_IMMUTABLE) { - executeBorrowedSegmentBuilds(config, slices, segmentDirectories, segmentBuilder); + executeSegmentBuilds(config, slices, segmentDirectories, segmentBuilder); } else { long overlapStartedAt = CagraHnswBuildMetrics.start(); try { @@ -620,8 +673,7 @@ private static void buildBorrowedSegments( validationOptions.scanHeadStartBytes(), config.metrics()); coordinator.runAfterHeadStart( - () -> - executeBorrowedSegmentBuilds(config, slices, segmentDirectories, segmentBuilder)); + () -> executeSegmentBuilds(config, slices, segmentDirectories, segmentBuilder)); } finally { config.metrics().stop("external reference overlap wall [GPU+DISK]", overlapStartedAt); } @@ -638,7 +690,7 @@ private static void buildBorrowedSegments( } } - private static void executeBorrowedSegmentBuilds( + private static void executeSegmentBuilds( Config config, List slices, List segmentDirectories, @@ -717,71 +769,6 @@ private static void buildSequential( } } - /** - * 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 @@ -886,16 +873,31 @@ private static void combineByHardlink(Path targetDir, List segDirs) throws } /** 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<>(); + static List sliceEvenly(int total, int k) { + if (total <= 0) { + throw new IllegalArgumentException("Bulk vector count must be > 0, got " + total); + } + if (total > IndexWriter.MAX_DOCS) { + throw new IllegalArgumentException( + "Bulk vector count " + + total + + " exceeds Lucene's index-wide document limit " + + IndexWriter.MAX_DOCS + + "; reduce the input size"); + } + if (k <= 0) { + throw new IllegalArgumentException("numSegments must be > 0, got " + k); + } + if (k > total) { + throw new IllegalArgumentException( + "numSegments " + k + " exceeds bulk vector count " + total); + } + List slices = new ArrayList<>(k); 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; } @@ -921,7 +923,10 @@ private static void deleteRecursivelyQuietly(Path path) { } } - /** Immutable configuration for {@link CagraHnswBulkIndexWriter}. */ + /** + * Configuration for {@link CagraHnswBulkIndexWriter}. Scalar settings are immutable; the + * referenced {@link CagraHnswBuildMetrics} is a caller-owned cumulative accumulator. + */ public static final class Config { private final String fieldName; private final int dimensions; @@ -967,7 +972,10 @@ public AcceleratedHNSWParams graphBuildParams() { return graphBuildParams; } - /** Only consulted by the one-shot build methods. */ + /** + * Absolute normalized path exclusively owned and recreated by the one-shot build methods; not + * consulted by the direct single-segment constructor. + */ public Path targetDirectory() { return targetDirectory; } @@ -987,7 +995,10 @@ public int pipelineDepth() { return pipelineDepth; } - /** Per-build metrics accumulator shared by all segments in this configuration. */ + /** + * Cumulative metrics accumulator shared by all segments and every invocation that reuses this + * configuration. Use a fresh configuration or accumulator for per-build measurements. + */ public CagraHnswBuildMetrics metrics() { return metrics; } @@ -1034,20 +1045,26 @@ public Builder graphBuild(AcceleratedHNSWParams graphBuildParams) { return this; } - /** Sets the directory the final index is written to; required for {@link #indexFbin}/{@link #build}. */ + /** + * Sets the exclusively owned directory that a one-shot method recreates as the final index. + * The stored path is absolute and normalized. This is path normalization, not a cross-process + * ownership lock; callers must prevent concurrent use of the same target. + */ public Builder targetDirectory(Path targetDirectory) { - this.targetDirectory = Objects.requireNonNull(targetDirectory, "targetDirectory"); + this.targetDirectory = + Objects.requireNonNull(targetDirectory, "targetDirectory").toAbsolutePath().normalize(); 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). + * segment. The requested count is exact and must not exceed the source row count. 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}, builds up to {@link #pipelineDepth} slices - * concurrently instead of strictly sequentially. Supported by the FBIN entry points and - * ignored by {@link #build(VectorSource, Config)}, whose caller-owned source is + * concurrently instead of strictly sequentially. Supported by the FBIN entry points; + * {@link #build(VectorSource, Config)} rejects overlap because its caller-owned source is * single-consumer. */ public Builder segments(int numSegments, boolean overlapped) { @@ -1068,7 +1085,10 @@ public Builder pipelineDepth(int pipelineDepth) { return this; } - /** Supplies a caller-owned metrics accumulator for this build. */ + /** + * Supplies a caller-owned cumulative metrics accumulator. Reusing the resulting + * configuration aggregates subsequent invocations into the same accumulator. + */ public Builder metrics(CagraHnswBuildMetrics metrics) { this.metrics = Objects.requireNonNull(metrics, "metrics"); return this; diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java index 3917879970..906a310ede 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java @@ -49,7 +49,6 @@ import org.apache.lucene.store.ReadAdvice; import org.apache.lucene.util.Bits; import org.apache.lucene.util.IOUtils; -import org.apache.lucene.util.VectorUtil; import org.apache.lucene.util.hnsw.IntToIntFunction; /** @@ -447,19 +446,8 @@ private interface FloatToFloatFunction { * @return an instance of the FloatToFloatFunction */ private static FloatToFloatFunction getScoreNormalizationFunc(VectorSimilarityFunction sim) { - return cuvsValue -> toLuceneScore(sim, cuvsValue); - } - - /** Convert a public cuVS search value into the score contract used by Lucene's collectors. */ - static float toLuceneScore(VectorSimilarityFunction sim, float cuvsValue) { - return switch (sim) { - case EUCLIDEAN -> 1.0f / (1.0f + cuvsValue); - // Public cuVS InnerProduct results contain the natural dot product. - case DOT_PRODUCT -> Math.max((1.0f + cuvsValue) / 2.0f, 0.0f); - // cuVS CosineExpanded reports 1-cos(a,b), while Lucene scales cosine to [0, 1]. - case COSINE -> Math.max(1.0f - cuvsValue / 2.0f, 0.0f); - case MAXIMUM_INNER_PRODUCT -> VectorUtil.scaleMaxInnerProductScore(cuvsValue); - }; + // TODO: check for different similarities + return score -> (1f / (1f + score)); } /** diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinFileRegistry.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinFileRegistry.java index 0d99f8451f..80b031c4d7 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinFileRegistry.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinFileRegistry.java @@ -13,12 +13,13 @@ import java.util.Objects; /** - * Process-local allowlist that resolves content-addressed external FBIN references. + * JVM/classloader-local allowlist that resolves content-addressed external FBIN references. * *

    Lucene recreates codecs through no-argument SPI constructors, so writer-time configuration is - * not available when a process later opens an index. Every indexing or search process must register - * the immutable local file for each referenced content ID before opening the index. Registrations - * are reference counted and should remain alive for at least the complete reader lifetime. + * not available when a process later opens an index. Every indexing or search application must + * register the immutable local file for each referenced content ID through the same defining + * cuvs-lucene classloader before opening the index. Registrations are reference counted; at least + * one matching registration must remain alive for every dependent reader's complete lifetime. * *

    Registration validates the path and content-ID syntax but deliberately does not hash the file. * Use {@link ExternalFbinBuildValidation#VERIFY_SHA256} while building, or run Lucene integrity @@ -31,7 +32,10 @@ public final class ExternalFbinFileRegistry { private ExternalFbinFileRegistry() {} - /** Registers an allowlisted local file for a complete-file SHA-256 content identity. */ + /** + * Registers an allowlisted local file for a complete-file SHA-256 content identity. The path is + * resolved once to its absolute real path, which is returned by {@link Registration#path()}. + */ public static Registration register(Path path, String sha256Hex) throws IOException { Objects.requireNonNull(path, "path"); byte[] digest = ExternalFbinReference.parseSha256(sha256Hex); @@ -65,9 +69,9 @@ static Path resolve(ExternalFbinReference reference) throws IOException { Entry entry = ENTRIES.get(reference.contentId()); if (entry == null) { throw new IOException( - "No allowlisted external FBIN is registered for " + "No allowlisted external FBIN is registered in this cuvs-lucene classloader for " + reference.contentId() - + ". Register it with ExternalFbinFileRegistry before opening the index."); + + ". Register it through the same defining classloader before opening the index."); } return entry.path; } @@ -89,7 +93,11 @@ private Entry(Path path, int references) { } } - /** A scoped registry lease. Closing the final lease removes the path from the allowlist. */ + /** + * A scoped registry lease. Closing the final lease removes the path from this classloader's + * allowlist. The registry does not own the file, and every reader that depends on its content ID + * must be closed before the final matching lease is closed. + */ public static final class Registration implements Closeable { private final Path path; @@ -102,8 +110,7 @@ private Registration(Path path, String contentId) { } /** Creates a path-independent reference to a contiguous row range in the registered FBIN. */ - public synchronized ExternalFbinReference reference(int firstRow, int rowCount) - throws IOException { + synchronized ExternalFbinReference reference(int firstRow, int rowCount) throws IOException { ensureOpen(); return ExternalFbinReference.fromFile(path, contentId, firstRow, rowCount); } @@ -111,8 +118,7 @@ public synchronized ExternalFbinReference reference(int firstRow, int rowCount) /** * Maps a registered row range and couples its native build dataset to the persisted reference. */ - public synchronized ImmutableExternalFbinDataset map(int firstRow, int rowCount) - throws IOException { + synchronized ImmutableExternalFbinDataset map(int firstRow, int rowCount) throws IOException { ensureOpen(); return ImmutableExternalFbinDataset.map( ExternalFbinReference.fromFile(path, contentId, firstRow, rowCount)); @@ -128,6 +134,7 @@ public String contentId() { return contentId; } + /** Returns the absolute real path captured when this registration was created. */ public Path path() { return path; } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinFlatVectorsReader.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinFlatVectorsReader.java index a9c267a79c..5761220a90 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinFlatVectorsReader.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinFlatVectorsReader.java @@ -7,7 +7,6 @@ import java.io.IOException; import java.nio.ByteBuffer; import java.nio.ByteOrder; -import java.nio.file.Path; import org.apache.lucene.codecs.CodecUtil; import org.apache.lucene.codecs.hnsw.DefaultFlatVectorScorer; import org.apache.lucene.codecs.hnsw.FlatVectorsReader; @@ -21,7 +20,6 @@ import org.apache.lucene.index.VectorEncoding; import org.apache.lucene.index.VectorSimilarityFunction; import org.apache.lucene.store.ChecksumIndexInput; -import org.apache.lucene.store.Directory; import org.apache.lucene.store.IOContext; import org.apache.lucene.store.IndexInput; import org.apache.lucene.store.MMapDirectory; @@ -38,8 +36,7 @@ final class ExternalFbinFlatVectorsReader extends FlatVectorsReader { private final FieldInfo fieldInfo; private final ExternalFbinReference reference; - private final Directory externalDirectory; - private final IndexInput sourceInput; + private final ExternalFbinMappedInputCache.Lease sourceLease; private final IndexInput payloadInput; ExternalFbinFlatVectorsReader(SegmentReadState state) throws IOException { @@ -52,19 +49,15 @@ final class ExternalFbinFlatVectorsReader extends FlatVectorsReader { fieldInfo = descriptor.fieldInfo(); reference = descriptor.reference(); - Path sourcePath = ExternalFbinIO.validateAndResolve(reference); - Directory newDirectory = null; - IndexInput newSource = null; + var sourcePath = ExternalFbinIO.validateAndResolve(reference); + ExternalFbinMappedInputCache.Lease newSourceLease = null; IndexInput newPayload = null; boolean success = false; try { - Path parent = sourcePath.getParent(); - if (parent == null) { - throw new IOException("External FBIN path has no parent directory: " + sourcePath); - } - newDirectory = new MMapDirectory(parent, maxMmapChunkSize); IOContext context = state.context.withReadAdvice(ReadAdvice.RANDOM); - newSource = newDirectory.openInput(sourcePath.getFileName().toString(), context); + newSourceLease = + ExternalFbinMappedInputCache.acquire(sourcePath, reference, maxMmapChunkSize, context); + IndexInput newSource = newSourceLease.sourceInput(); validateOpenedSource(newSource, reference); newPayload = newSource.slice( @@ -75,11 +68,10 @@ final class ExternalFbinFlatVectorsReader extends FlatVectorsReader { success = true; } finally { if (!success) { - IOUtils.closeWhileHandlingException(newPayload, newSource, newDirectory); + IOUtils.closeWhileHandlingException(newPayload, newSourceLease); } } - externalDirectory = newDirectory; - sourceInput = newSource; + sourceLease = newSourceLease; payloadInput = newPayload; } @@ -276,6 +268,7 @@ public RandomVectorScorer getRandomVectorScorer(String field, byte[] target) { @Override public void checkIntegrity() throws IOException { + IndexInput sourceInput = sourceLease.sourceInput(); IndexInput clone = sourceInput.clone(); // Lucene IndexInput clones are non-owning. Multi-chunk MemorySegmentIndexInput clones share // the owner's segment array, so closing a clone would invalidate the live reader. @@ -305,7 +298,7 @@ public long ramBytesUsed() { @Override public void close() throws IOException { - IOUtils.close(payloadInput, sourceInput, externalDirectory); + IOUtils.close(payloadInput, sourceLease); } private record Descriptor(FieldInfo fieldInfo, ExternalFbinReference reference) {} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinIO.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinIO.java index ea15c224af..506187b19a 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinIO.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinIO.java @@ -24,9 +24,19 @@ private ExternalFbinIO() {} static Path validateAndResolve(ExternalFbinReference reference) throws IOException { Path path = ExternalFbinFileRegistry.resolve(reference); - ExternalFbinReference actual = - ExternalFbinReference.fromFile( - path, reference.sha256Hex(), Math.toIntExact(reference.firstRow()), reference.rows()); + final ExternalFbinReference actual; + try { + actual = + ExternalFbinReference.fromFile( + path, reference.sha256Hex(), Math.toIntExact(reference.firstRow()), reference.rows()); + } catch (IllegalArgumentException | ArithmeticException incompatibleFile) { + throw new IOException( + "Registered external FBIN " + + path + + " is incompatible with persisted reference " + + reference.contentId(), + incompatibleFile); + } if (!actual.equals(reference)) { throw new IOException( "Registered FBIN metadata does not match persisted reference for " diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinMappedInputCache.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinMappedInputCache.java new file mode 100644 index 0000000000..2500f56a38 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinMappedInputCache.java @@ -0,0 +1,146 @@ +/* + * 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; +import java.nio.file.Path; +import java.util.HashMap; +import java.util.Map; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.store.MMapDirectory; +import org.apache.lucene.util.IOUtils; + +/** Shares the one owning full-file FBIN mapping among every live reader for the same artifact. */ +final class ExternalFbinMappedInputCache { + + private static final Object LOCK = new Object(); + private static final Map OWNERS = new HashMap<>(); + + private ExternalFbinMappedInputCache() {} + + static Lease acquire( + Path realPath, ExternalFbinReference reference, long maxMmapChunkSize, IOContext context) + throws IOException { + Key key = new Key(realPath, reference.contentId(), reference.fileLength(), maxMmapChunkSize); + synchronized (LOCK) { + Owner owner = OWNERS.get(key); + if (owner == null) { + owner = openOwner(realPath, maxMmapChunkSize, context); + OWNERS.put(key, owner); + } + IndexInput sourceView = owner.rootInput.clone(); + owner.leases++; + return new Lease(key, owner, sourceView); + } + } + + private static Owner openOwner(Path realPath, long maxMmapChunkSize, IOContext context) + throws IOException { + Path parent = realPath.getParent(); + if (parent == null) { + throw new IOException("External FBIN path has no parent directory: " + realPath); + } + MMapDirectory directory = null; + IndexInput rootInput = null; + boolean success = false; + try { + directory = new MMapDirectory(parent, maxMmapChunkSize); + rootInput = directory.openInput(realPath.getFileName().toString(), context); + Owner owner = new Owner(directory, rootInput); + success = true; + return owner; + } finally { + if (!success) { + IOUtils.closeWhileHandlingException(rootInput, directory); + } + } + } + + static int activeOwnerCountForTests( + Path realPath, ExternalFbinReference reference, long maxMmapChunkSize) throws IOException { + synchronized (LOCK) { + return OWNERS.containsKey(key(realPath, reference, maxMmapChunkSize)) ? 1 : 0; + } + } + + static int activeLeaseCountForTests( + Path realPath, ExternalFbinReference reference, long maxMmapChunkSize) throws IOException { + synchronized (LOCK) { + Owner owner = OWNERS.get(key(realPath, reference, maxMmapChunkSize)); + return owner == null ? 0 : owner.leases; + } + } + + private static Key key(Path path, ExternalFbinReference reference, long maxMmapChunkSize) + throws IOException { + return new Key( + path.toRealPath(), reference.contentId(), reference.fileLength(), maxMmapChunkSize); + } + + private static final class Owner { + private final MMapDirectory directory; + private final IndexInput rootInput; + private int leases; + + private Owner(MMapDirectory directory, IndexInput rootInput) { + this.directory = directory; + this.rootInput = rootInput; + } + } + + /** + * A reader lease over a non-owning full-file clone. + * + *

    Lucene's multi-chunk mmap clone shares its segment array with the owning input, so the clone + * must never be closed directly. Closing this lease releases only the shared owner reference; + * the final lease closes the sole owning input and its directory. + */ + static final class Lease implements Closeable { + private final Key key; + private final Owner owner; + private IndexInput sourceView; + private boolean closed; + + private Lease(Key key, Owner owner, IndexInput sourceView) { + this.key = key; + this.owner = owner; + this.sourceView = sourceView; + } + + IndexInput sourceInput() { + synchronized (LOCK) { + if (closed) { + throw new IllegalStateException("External FBIN mapping lease is closed"); + } + return sourceView; + } + } + + @Override + public void close() throws IOException { + synchronized (LOCK) { + if (closed) { + return; + } + closed = true; + sourceView = null; + Owner current = OWNERS.get(key); + if (current != owner || owner.leases <= 0) { + throw new IllegalStateException("External FBIN mapping cache reference count is corrupt"); + } + owner.leases--; + if (owner.leases == 0) { + OWNERS.remove(key); + IOUtils.close(owner.rootInput, owner.directory); + } + } + } + } + + private record Key( + Path realPath, String contentId, long expectedFileLength, long maxMmapChunkSize) {} +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinReference.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinReference.java index 9821698227..568303645c 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinReference.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFbinReference.java @@ -23,7 +23,7 @@ * process must register an allowlisted local path for this content identity through {@link * ExternalFbinFileRegistry} before opening an index that references it. */ -public final class ExternalFbinReference { +final class ExternalFbinReference { static final long HEADER_BYTES = 2L * Integer.BYTES; static final int SHA256_BYTES = 32; @@ -79,8 +79,8 @@ private ExternalFbinReference( *

    This method does not hash the file. {@code sha256Hex} must be a previously established digest * of the complete FBIN file. */ - public static ExternalFbinReference fromFile( - Path path, String sha256Hex, int firstRow, int rowCount) throws IOException { + static ExternalFbinReference fromFile(Path path, String sha256Hex, int firstRow, int rowCount) + throws IOException { Objects.requireNonNull(path, "path"); if (firstRow < 0 || rowCount <= 0) { throw new IllegalArgumentException("firstRow must be non-negative and rowCount positive"); @@ -187,39 +187,39 @@ private static void readFully(FileChannel channel, ByteBuffer target, long posit } } - public String contentId() { + String contentId() { return contentId(sha256); } - public String sha256Hex() { + String sha256Hex() { return HEX.formatHex(sha256); } - public byte[] sha256() { + byte[] sha256() { return sha256.clone(); } - public long fileLength() { + long fileLength() { return fileLength; } - public long payloadOffset() { + long payloadOffset() { return payloadOffset; } - public long payloadLength() { + long payloadLength() { return payloadLength; } - public int rows() { + int rows() { return rows; } - public int dimensions() { + int dimensions() { return dimensions; } - public long firstRow() { + long firstRow() { return (payloadOffset - HEADER_BYTES) / ((long) dimensions * Float.BYTES); } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFloat32Dataset.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFloat32Dataset.java index 039f4ed0f2..e7578fb50a 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFloat32Dataset.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ExternalFloat32Dataset.java @@ -24,7 +24,7 @@ * for ensuring that every raw vector component is finite; unlike ordinary Lucene field ingestion, * this path deliberately does not scan the payload one float at a time. */ -public final class ExternalFloat32Dataset { +final class ExternalFloat32Dataset { private final MemorySegment memorySegment; private final CuVSHostMatrix matrix; @@ -42,7 +42,7 @@ private ExternalFloat32Dataset( /** * Wraps a contiguous native-memory region containing row-major float32 values without copying. */ - public static ExternalFloat32Dataset fromMemorySegment( + static ExternalFloat32Dataset fromMemorySegment( MemorySegment memorySegment, int rows, int dimensions) { Objects.requireNonNull(memorySegment, "memorySegment"); if (rows <= 0) { @@ -115,19 +115,19 @@ public static ExternalFloat32Dataset fromMemorySegment( return new ExternalFloat32Dataset(memorySegment, hostMatrix, rows, dimensions); } - public MemorySegment memorySegment() { + MemorySegment memorySegment() { return memorySegment; } - public CuVSHostMatrix matrix() { + CuVSHostMatrix matrix() { return matrix; } - public int rows() { + int rows() { return rows; } - public int dimensions() { + int dimensions() { return dimensions; } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ImmutableExternalFbinDataset.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ImmutableExternalFbinDataset.java index 54c9ac3a7c..c9df249b83 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ImmutableExternalFbinDataset.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/ImmutableExternalFbinDataset.java @@ -20,7 +20,7 @@ * persisting a reference to another. Closing it releases both the cuVS matrix wrapper and mapping; * callers must keep it open until the index writer has finished. */ -public final class ImmutableExternalFbinDataset implements AutoCloseable { +final class ImmutableExternalFbinDataset implements AutoCloseable { private final ExternalFbinReference reference; private final Arena arena; @@ -70,16 +70,16 @@ ExternalFloat32Dataset dataset() { return dataset; } - public ExternalFbinReference reference() { + ExternalFbinReference reference() { ensureOpen(); return reference; } - public int rows() { + int rows() { return reference.rows(); } - public int dimensions() { + int dimensions() { return reference.dimensions(); } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/MappedSelfContainedOutput.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/MappedSelfContainedOutput.java index 89b3f7822e..5985d0eceb 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/MappedSelfContainedOutput.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/MappedSelfContainedOutput.java @@ -6,11 +6,11 @@ import java.io.IOException; import java.io.InterruptedIOException; -import java.util.concurrent.CancellationException; -import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; import org.apache.lucene.index.DocsWithFieldSet; import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.SegmentWriteState; @@ -36,64 +36,129 @@ public void writeField( AcceleratedHnswGraphOutput graphOutput) throws IOException { long overlapStartedAt = CagraHnswBuildMetrics.start(); - try (ExecutorService executor = + try { + runOverlapped( + () -> { + long startedAt = CagraHnswBuildMetrics.start(); + flatOutput.writeField( + field, dataset.matrix(), dataset.memorySegment(), maxDoc, docsWithField); + metrics.stop( + "mapped flat-write worker [DISK]", startedAt, dataset.memorySegment().byteSize()); + }, + () -> graphOutput.writeBorrowedField(field, dataset.matrix())); + } finally { + metrics.stop("mapped flush overlap wall [GPU+DISK]", overlapStartedAt); + } + } + + /** + * Runs the flat-vector write on its worker while the caller builds the graph. Package visibility + * keeps the failure and cancellation protocol independently testable without native resources. + */ + static void runOverlapped(OverlapTask flatWrite, OverlapTask graphWrite) throws IOException { + ExecutorService executor = Executors.newSingleThreadExecutor( - task -> Thread.ofPlatform().name("cuvs-mapped-flat-writer").unstarted(task))) { - Future flatWrite = + task -> Thread.ofPlatform().name("cuvs-mapped-flat-writer").unstarted(task)); + AtomicReference flatFailure = new AtomicReference<>(); + Future flatFuture = null; + Throwable failure = null; + boolean interrupted = false; + + try { + flatFuture = executor.submit( () -> { - long startedAt = CagraHnswBuildMetrics.start(); - flatOutput.writeField( - field, dataset.matrix(), dataset.memorySegment(), maxDoc, docsWithField); - metrics.stop( - "mapped flat-write worker [DISK]", - startedAt, - dataset.memorySegment().byteSize()); - return null; + try { + flatWrite.run(); + } catch (Throwable thrown) { + // A cancelled Future can report completion before its running task has stopped + // and no longer exposes an exception thrown after cancellation. Retain the + // worker's actual failure until executor termination establishes visibility. + flatFailure.compareAndSet(null, thrown); + } }); - Throwable failure = null; - try { - graphOutput.writeBorrowedField(field, dataset.matrix()); - } catch (Throwable graphFailure) { - failure = graphFailure; + if (Thread.interrupted()) { + interrupted = true; + failure = interruption("Interrupted before mapped graph output", null); + } else { + try { + graphWrite.run(); + } catch (InterruptedException graphInterruption) { + interrupted = true; + failure = + interruption("Interrupted while writing mapped graph output", graphInterruption); + } catch (Throwable graphFailure) { + failure = graphFailure; + } } - failure = combine(failure, await(flatWrite)); - if (failure != null) { - throw Utils.handleThrowable(failure); + + if (Thread.interrupted()) { + interrupted = true; + failure = + combine(failure, interruption("Interrupted while writing mapped graph output", null)); } + } catch (Throwable orchestrationFailure) { + failure = combine(failure, orchestrationFailure); } finally { - metrics.stop("mapped flush overlap wall [GPU+DISK]", overlapStartedAt); - } - } + if (failure != null || interrupted) { + if (flatFuture != null) { + flatFuture.cancel(true); + } + executor.shutdownNow(); + } else { + executor.shutdown(); + } - private static Throwable await(Future future) { - boolean interrupted = false; - Throwable failure = null; - while (true) { - try { - future.get(); - break; - } catch (InterruptedException e) { + // Future.cancel(true) only changes Future state; it does not prove the task honored the + // interrupt. Keep all borrowed mappings and outputs alive until the worker has really exited. + while (executor.isTerminated() == false) { + try { + executor.awaitTermination(1L, TimeUnit.DAYS); + } catch (InterruptedException awaitInterruption) { + interrupted = true; + failure = + combine( + failure, + interruption( + "Interrupted while awaiting mapped flat-vector output", awaitInterruption)); + if (flatFuture != null) { + flatFuture.cancel(true); + } + executor.shutdownNow(); + } + } + + // Cover an interrupt arriving after the final await returned without clearing it early. + if (Thread.interrupted()) { interrupted = true; - } catch (ExecutionException e) { - failure = e.getCause(); - break; - } catch (CancellationException e) { - failure = e; - break; + failure = + combine( + failure, + interruption("Interrupted while awaiting mapped flat-vector output", null)); } - } - if (interrupted) { - Thread.currentThread().interrupt(); - InterruptedIOException interruption = - new InterruptedIOException("Interrupted while awaiting mapped flat-vector output"); - if (failure != null) { - interruption.addSuppressed(failure); + failure = combine(failure, flatFailure.get()); + if (interrupted) { + Thread.currentThread().interrupt(); } - return interruption; } - return failure; + + if (failure != null) { + throw Utils.handleThrowable(failure); + } + } + + private static InterruptedIOException interruption(String message, Throwable cause) { + InterruptedIOException interruption = new InterruptedIOException(message); + if (cause != null) { + interruption.initCause(cause); + } + return interruption; + } + + @FunctionalInterface + interface OverlapTask { + void run() throws Throwable; } private static Throwable combine(Throwable primary, Throwable additional) { diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatBufferedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatBufferedHNSWVectorsWriter.java index 0a92f8600e..960ff22dd5 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatBufferedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatBufferedHNSWVectorsWriter.java @@ -5,7 +5,6 @@ package com.nvidia.cuvs.lucene; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.printInfoStream; -import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.closeCuVSResourcesInstance; import static org.apache.lucene.index.VectorEncoding.FLOAT32; import static org.apache.lucene.util.RamUsageEstimator.shallowSizeOfInstance; @@ -183,10 +182,8 @@ public void close() throws IOException { } static void closeOutputsAndResources(Closeable... outputs) throws IOException { - try { + try (Closeable resourceCloser = ThreadLocalCuVSResourcesProvider::closeCuVSResourcesInstance) { IOUtils.close(outputs); - } finally { - closeCuVSResourcesInstance(); } } 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 index 3a23c657d4..c54337a6ed 100644 --- 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 @@ -8,6 +8,7 @@ import com.nvidia.cuvs.CuVSHostMatrix; import java.io.Closeable; import java.io.IOException; +import java.io.InterruptedIOException; import java.lang.foreign.MemorySegment; import java.lang.foreign.ValueLayout; import java.nio.ByteOrder; @@ -174,16 +175,37 @@ void writeField( static void copyRawFloat32Vectors( MemorySegment source, long byteCount, DataOutput output, CagraHnswBuildMetrics metrics) throws IOException { + copyRawFloat32Vectors(source, byteCount, output, metrics, RAW_CHUNK_BYTES); + } + + // The explicit chunk size makes cooperative cancellation deterministic in CPU-only tests. + static void copyRawFloat32Vectors( + MemorySegment source, + long byteCount, + DataOutput output, + CagraHnswBuildMetrics metrics, + int chunkBytes) + throws IOException { if (byteCount < 0L || byteCount > source.byteSize()) { throw new IllegalArgumentException("Invalid raw float32 byte count: " + byteCount); } - byte[] chunk = new byte[(int) Math.min(RAW_CHUNK_BYTES, Math.max(1L, byteCount))]; + if (chunkBytes <= 0) { + throw new IllegalArgumentException("Raw float32 chunk size must be positive"); + } + byte[] chunk = new byte[(int) Math.min(chunkBytes, Math.max(1L, byteCount))]; MemorySegment chunkSegment = MemorySegment.ofArray(chunk); long offset = 0L; long stagingCopyNanos = 0L; long outputWriteNanos = 0L; long chunks = 0L; while (offset < byteCount) { + if (Thread.currentThread().isInterrupted()) { + throw new InterruptedIOException( + "Interrupted while copying mapped flat-vector data at byte " + + offset + + " of " + + byteCount); + } int length = (int) Math.min(chunk.length, byteCount - offset); long startedAt = CagraHnswBuildMetrics.start(); MemorySegment.copy(source, offset, chunkSegment, 0L, length); diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/CuVSTestSupport.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/CuVSTestSupport.java new file mode 100644 index 0000000000..b68aa3afe8 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/CuVSTestSupport.java @@ -0,0 +1,37 @@ +/* + * 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 com.nvidia.cuvs.spi.CuVSProvider; +import org.junit.Assume; + +/** Keeps local no-GPU tests skippable while making the required GPU CI job fail closed. */ +final class CuVSTestSupport { + + private static final String REQUIRE_GPU_ENV = "CUVS_LUCENE_REQUIRE_GPU_TESTS"; + + private CuVSTestSupport() {} + + static void enableRmmOrSkip() { + try { + CuVSProvider.provider().enableRMMAsyncMemory(); + } catch (UnsupportedOperationException unsupported) { + requireOrSkip("cuVS RMM initialization failed: " + unsupported.getMessage(), false); + } + } + + static void requireCuvsOrSkip() { + requireOrSkip("cuVS Java/native support is unavailable", isSupported()); + } + + private static void requireOrSkip(String message, boolean condition) { + if (!condition && "1".equals(System.getenv(REQUIRE_GPU_ENV))) { + throw new AssertionError(message + " (required because " + REQUIRE_GPU_ENV + "=1)"); + } + Assume.assumeTrue(message, condition); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedThreeLayerRoundTrip.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedThreeLayerRoundTrip.java new file mode 100644 index 0000000000..91e2eb8c92 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedThreeLayerRoundTrip.java @@ -0,0 +1,216 @@ +/* + * 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.CuVSTestSupport.enableRmmOrSkip; +import static com.nvidia.cuvs.lucene.CuVSTestSupport.requireCuvsOrSkip; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; + +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; +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.StandardOpenOption; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.hnsw.HnswGraphProvider; +import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; +import org.apache.lucene.index.CodecReader; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.LeafReader; +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.SuppressFileSystems; +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.Before; +import org.junit.BeforeClass; +import org.junit.Test; + +/** Required GPU proof for the accelerated upper-layer build and persisted graph. */ +@SuppressFileSystems("*") +@SuppressSysoutChecks(bugUrl = "") +public class TestAcceleratedThreeLayerRoundTrip extends LuceneTestCase { + + private static final int ROWS = 512; + private static final int DIMENSIONS = 32; + private static final int GRAPH_DEGREE = 8; + private static final long LAYER_SEED = 0x2594L; + private static final String VECTOR_FIELD = "vector"; + + @BeforeClass + public static void beforeClass() { + enableRmmOrSkip(); + } + + @Before + public void requireCuvs() { + requireCuvsOrSkip(); + } + + @Test + public void testAcceleratedThreeLayerGraphPersistsAndSearches() throws Exception { + Path root = createTempDir("accelerated-three-layer"); + float[][] vectors = createVectors(); + Path fbin = writeFbin(root.resolve("vectors.fbin"), vectors); + Path index = root.resolve("index"); + CagraHnswBuildMetrics metrics = new CagraHnswBuildMetrics(); + AcceleratedHNSWParams graphBuild = + new AcceleratedHNSWParams.Builder() + .withStrategy(AcceleratedHNSWParams.Strategy.CUSTOM) + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(GRAPH_DEGREE) + .withIntermediateGraphDegree(16) + .withHNSWLayer(3) + .withHnswLayerSeed(LAYER_SEED) + .build(); + CagraHnswBulkIndexWriter.Config config = + CagraHnswBulkIndexWriter.Config.builder() + .field(VECTOR_FIELD, DIMENSIONS, EUCLIDEAN) + .graphBuild(graphBuild) + .segments(1, false) + .targetDirectory(index) + .metrics(metrics) + .build(); + + CagraHnswBulkIndexWriter.indexMappedFbin(fbin, config); + Files.delete(fbin); + + try (Directory directory = FSDirectory.open(index); + DirectoryReader reader = DirectoryReader.open(directory)) { + assertEquals(1, reader.leaves().size()); + LeafReader leaf = reader.leaves().getFirst().reader(); + HnswGraph graph = graphOf(leaf); + assertEquals(3, graph.numLevels()); + List expectedLayers = + AcceleratedHNSWUtils.selectUpperLayerNodes( + ROWS, 3, Math.ceilDiv(GRAPH_DEGREE, 2), LAYER_SEED); + assertEquals(2, expectedLayers.size()); + for (int level = 1; level < graph.numLevels(); level++) { + int[] nodes = NodesIterator.getSortedNodes(graph.getNodesOnLevel(level)); + assertArrayEquals(expectedLayers.get(level - 1), nodes); + assertNeighborsRemainOnLevel(graph, level, nodes); + } + int[] topLevel = NodesIterator.getSortedNodes(graph.getNodesOnLevel(2)); + assertTrue(Arrays.binarySearch(topLevel, graph.entryNode()) >= 0); + + FloatVectorValues values = leaf.getFloatVectorValues(VECTOR_FIELD); + assertArrayEquals(vectors[0], values.vectorValue(0), 0.0f); + assertArrayEquals(vectors[ROWS - 1], values.copy().vectorValue(ROWS - 1), 0.0f); + TopDocs hits = + new IndexSearcher(reader).search(new KnnFloatVectorQuery(VECTOR_FIELD, vectors[0], 5), 5); + assertEquals(5, hits.scoreDocs.length); + } + + Map snapshot = metrics.snapshot(); + assertEquals(1L, snapshot.get("stage/base cagra-build [GPU]/count").longValue()); + assertEquals(1L, snapshot.get("stage/hnsw-convert [GPU+PCIe+CPU]/count").longValue()); + assertEquals(GRAPH_DEGREE, snapshot.get("gauge/effective graph degree").intValue()); + } + + @Test + public void testMetricsSeparateSelectedAndNativeTruncatedGraphDegree() throws Exception { + int rows = 20; + int dimensions = 8; + int selectedDegree = 32; + Path root = createTempDir("actual-graph-degree"); + Path fbin = writeFbin(root.resolve("vectors.fbin"), createVectors(rows, dimensions)); + CagraHnswBuildMetrics metrics = new CagraHnswBuildMetrics(); + AcceleratedHNSWParams graphBuild = + new AcceleratedHNSWParams.Builder() + .withStrategy(AcceleratedHNSWParams.Strategy.CUSTOM) + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(selectedDegree) + .withIntermediateGraphDegree(64) + .build(); + CagraHnswBulkIndexWriter.Config config = + CagraHnswBulkIndexWriter.Config.builder() + .field(VECTOR_FIELD, dimensions, EUCLIDEAN) + .graphBuild(graphBuild) + .targetDirectory(root.resolve("index")) + .metrics(metrics) + .build(); + + CagraHnswBulkIndexWriter.indexMappedFbin(fbin, config); + + Map snapshot = metrics.snapshot(); + int effectiveDegree = snapshot.get("gauge/effective graph degree").intValue(); + assertEquals(selectedDegree, snapshot.get("gauge/selected graph degree").intValue()); + assertEquals(rows - 1, effectiveDegree); + assertEquals( + (long) rows * effectiveDegree * Integer.BYTES, + snapshot.get("counter/actual cagra adjacency bytes").longValue()); + } + + private static void assertNeighborsRemainOnLevel(HnswGraph graph, int level, int[] nodes) + throws Exception { + for (int node : nodes) { + graph.seek(level, node); + for (int neighbor = graph.nextNeighbor(); + neighbor != NO_MORE_DOCS; + neighbor = graph.nextNeighbor()) { + assertTrue( + "upper-layer neighbor " + neighbor + " is absent from level " + level, + Arrays.binarySearch(nodes, neighbor) >= 0); + } + } + } + + private static HnswGraph graphOf(LeafReader leaf) throws Exception { + KnnVectorsReader knnReader = ((CodecReader) leaf).getVectorReader(); + if (knnReader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) { + knnReader = fieldsReader.getFieldReader(VECTOR_FIELD); + } + return ((HnswGraphProvider) knnReader).getGraph(VECTOR_FIELD); + } + + private static float[][] createVectors() { + return createVectors(ROWS, DIMENSIONS); + } + + private static float[][] createVectors(int rows, int dimensions) { + float[][] vectors = new float[rows][dimensions]; + for (int row = 0; row < rows; row++) { + for (int dimension = 0; dimension < dimensions; dimension++) { + vectors[row][dimension] = (float) Math.sin((row + 1L) * (dimension + 3L) * 0.0009765625d); + } + } + return vectors; + } + + private static Path writeFbin(Path file, float[][] vectors) throws Exception { + int rows = vectors.length; + int dimensions = vectors[0].length; + ByteBuffer data = + ByteBuffer.allocate(2 * Integer.BYTES + rows * dimensions * Float.BYTES) + .order(ByteOrder.LITTLE_ENDIAN); + data.putInt(rows); + data.putInt(dimensions); + for (float[] vector : vectors) { + for (float value : vector) { + data.putFloat(value); + } + } + data.flip(); + try (FileChannel channel = + FileChannel.open(file, StandardOpenOption.CREATE_NEW, StandardOpenOption.WRITE)) { + while (data.hasRemaining()) { + channel.write(data); + } + } + return file; + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBulkExternalFbinApiSurface.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBulkExternalFbinApiSurface.java index f16670709c..849d23df87 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBulkExternalFbinApiSurface.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBulkExternalFbinApiSurface.java @@ -52,6 +52,21 @@ public void testExternalBuildIsOnlyExposedByBulkWriter() throws Exception { || type == ExternalFloat32Dataset.class)); } + @Test + public void testBorrowedDatasetImplementationTypesAreNotPublicApi() { + assertFalse(isPublic(ExternalFloat32Dataset.class.getModifiers())); + assertFalse(isPublic(ImmutableExternalFbinDataset.class.getModifiers())); + assertFalse(isPublic(ExternalFbinReference.class.getModifiers())); + assertThrows( + NoSuchMethodException.class, + () -> + ExternalFbinFileRegistry.Registration.class.getMethod( + "reference", int.class, int.class)); + assertThrows( + NoSuchMethodException.class, + () -> ExternalFbinFileRegistry.Registration.class.getMethod("map", int.class, int.class)); + } + @Test public void testExternalOptionsValidateHeadStart() { assertThrows( diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBulkFbinIndexStorage.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBulkFbinIndexStorage.java index 1220c244eb..bc294f317e 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBulkFbinIndexStorage.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBulkFbinIndexStorage.java @@ -4,11 +4,11 @@ */ package com.nvidia.cuvs.lucene; -import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static com.nvidia.cuvs.lucene.CuVSTestSupport.enableRmmOrSkip; +import static com.nvidia.cuvs.lucene.CuVSTestSupport.requireCuvsOrSkip; import static org.apache.lucene.index.VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT; import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; -import com.nvidia.cuvs.spi.CuVSProvider; import java.io.IOException; import java.nio.ByteBuffer; import java.nio.channels.FileChannel; @@ -46,16 +46,12 @@ public class TestBulkFbinIndexStorage extends LuceneTestCase { @BeforeClass public static void beforeClass() { - try { - CuVSProvider.provider().enableRMMAsyncMemory(); - } catch (UnsupportedOperationException unsupported) { - assumeTrue("cuVS not supported: " + unsupported.getMessage(), false); - } + enableRmmOrSkip(); } @Before public void requireCuvs() { - assumeTrue("cuVS not supported", isSupported()); + requireCuvsOrSkip(); } @After diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBulkFbinMultiSegmentStorage.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBulkFbinMultiSegmentStorage.java index c268575bbd..7a1a7aacbe 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBulkFbinMultiSegmentStorage.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestBulkFbinMultiSegmentStorage.java @@ -4,11 +4,11 @@ */ package com.nvidia.cuvs.lucene; -import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static com.nvidia.cuvs.lucene.CuVSTestSupport.enableRmmOrSkip; +import static com.nvidia.cuvs.lucene.CuVSTestSupport.requireCuvsOrSkip; import static org.apache.lucene.index.VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT; import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; -import com.nvidia.cuvs.spi.CuVSProvider; import java.io.IOException; import java.nio.ByteBuffer; import java.nio.channels.FileChannel; @@ -57,16 +57,12 @@ public class TestBulkFbinMultiSegmentStorage extends LuceneTestCase { @BeforeClass public static void beforeClass() { - try { - CuVSProvider.provider().enableRMMAsyncMemory(); - } catch (UnsupportedOperationException unsupported) { - assumeTrue("cuVS not supported: " + unsupported.getMessage(), false); - } + enableRmmOrSkip(); } @Before public void requireCuvs() { - assumeTrue("cuVS not supported", isSupported()); + requireCuvsOrSkip(); } @After @@ -140,7 +136,36 @@ public void testExternalBuildVerifiesOnceAndPreservesSlices() throws Exception { assertEquals(1L, snapshot.get("stage/external fbin SHA-256 [DISK+CPU]/count").longValue()); assertEquals( Files.size(fbin), snapshot.get("stage/external fbin SHA-256 [DISK+CPU]/bytes").longValue()); - assertEquals(SEGMENTS, snapshot.get("stage/cagra-build [GPU]/count").longValue()); + assertEquals(SEGMENTS, snapshot.get("stage/base cagra-build [GPU]/count").longValue()); + } + + @Test + public void testExternalIndexAndFbinCanRelocateAfterBuildLeaseCloses() throws Exception { + Path root = createTempDir("external-relocation"); + float[][] vectors = createVectors(); + Path originalFbin = writeFbin(root, vectors); + Path originalIndex = root.resolve("index-a"); + String digest = sha256(originalFbin); + + try (ExternalFbinFileRegistry.Registration registration = + ExternalFbinFileRegistry.register(originalFbin, digest)) { + CagraHnswBulkIndexWriter.indexImmutableFbin( + registration, + config(originalIndex, new CagraHnswBuildMetrics()), + ExternalFbinOptions.verifySha256(), + (document, id) -> document.add(new StoredField(CALLBACK_ID_FIELD, id))); + } + + Path relocatedRoot = Files.createDirectory(root.resolve("relocated")); + Path relocatedFbin = Files.move(originalFbin, relocatedRoot.resolve("vectors.fbin")); + Path relocatedIndex = Files.move(originalIndex, relocatedRoot.resolve("index-b")); + assertFalse(Files.exists(originalFbin)); + assertFalse(Files.exists(originalIndex)); + + try (ExternalFbinFileRegistry.Registration relocatedLease = + ExternalFbinFileRegistry.register(relocatedFbin, digest)) { + assertIndex(relocatedIndex, vectors); + } } @Test @@ -173,7 +198,7 @@ public void testExternalOverlappedBuildPrefetchesOnceAndPublishesCompleteIndex() assertEquals( Files.size(fbin) - ExternalFbinReference.HEADER_BYTES, snapshot.get("stage/external fbin prefetch [DISK]/bytes").longValue()); - assertEquals(SEGMENTS, snapshot.get("stage/cagra-build [GPU]/count").longValue()); + assertEquals(SEGMENTS, snapshot.get("stage/base cagra-build [GPU]/count").longValue()); } @Test diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBuildMetricsConcurrency.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBuildMetricsConcurrency.java index 79f6c6ba9c..ac7d0364aa 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBuildMetricsConcurrency.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBuildMetricsConcurrency.java @@ -4,6 +4,7 @@ */ package com.nvidia.cuvs.lucene; +import java.util.Map; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; @@ -14,6 +15,18 @@ public class TestCagraHnswBuildMetricsConcurrency extends LuceneTestCase { + @Test + public void testAccumulatorReuseIsExplicitlyCumulative() { + CagraHnswBuildMetrics metrics = new CagraHnswBuildMetrics(); + metrics.record("stage", 10L, 20L); + Map firstBuild = metrics.snapshot(); + metrics.record("stage", 30L, 40L); + + assertEquals(1L, firstBuild.get("stage/stage/count").longValue()); + assertEquals(2L, metrics.snapshot().get("stage/stage/count").longValue()); + assertEquals(60L, metrics.snapshot().get("stage/stage/bytes").longValue()); + } + @Test public void testGaugeAllowsRepeatedIdenticalValues() { CagraHnswBuildMetrics metrics = new CagraHnswBuildMetrics(); @@ -39,11 +52,11 @@ public void testGaugeReportsRangeForDifferentSegmentValues() { public void testCounterRemainsAdditive() { CagraHnswBuildMetrics metrics = new CagraHnswBuildMetrics(); - metrics.addCounter("logical cagra adjacency bytes", 4_096L); - metrics.addCounter("logical cagra adjacency bytes", 8_192L); + metrics.addCounter("actual cagra adjacency bytes", 4_096L); + metrics.addCounter("actual cagra adjacency bytes", 8_192L); assertEquals( - 12_288L, metrics.snapshot().get("counter/logical cagra adjacency bytes").longValue()); + 12_288L, metrics.snapshot().get("counter/actual cagra adjacency bytes").longValue()); } @Test 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 index 204039ff7d..618771b0a7 100644 --- 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 @@ -4,14 +4,15 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.CuVSTestSupport.enableRmmOrSkip; +import static com.nvidia.cuvs.lucene.CuVSTestSupport.requireCuvsOrSkip; 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.CuvsDistanceType; -import com.nvidia.cuvs.spi.CuVSProvider; import java.io.File; import java.io.IOException; +import java.nio.file.Files; import java.nio.file.Path; import java.nio.file.Paths; import java.util.Map; @@ -61,18 +62,12 @@ public class TestCagraHnswBulkIndexWriter extends LuceneTestCase { @BeforeClass public static void beforeClass() { - // It is recommended to enable RMM allocation mode at application start, before constructing - // any CagraHnswBulkIndexWriter, to avoid device-wide sync from the default allocator. - try { - CuVSProvider.provider().enableRMMAsyncMemory(); - } catch (UnsupportedOperationException unsupported) { - assumeTrue("cuVS not supported: " + unsupported.getMessage(), false); - } + enableRmmOrSkip(); } @Before public void beforeTest() { - assumeTrue("cuVS not supported", isSupported()); + requireCuvsOrSkip(); random = new Random(222); indexDirPath = Paths.get(UUID.randomUUID().toString()); fbinPath = Paths.get(UUID.randomUUID() + ".fbin"); @@ -127,10 +122,20 @@ public void testOverlappedBuildProducesKSegments() throws Exception { int k = 4; float[][] dataset = generateDataset(random, numDocs, dimension); TestUtils.writeFbin(fbinPath, dataset); + Path preExistingSibling = + indexDirPath.resolveSibling(indexDirPath.getFileName().toString() + "_p0"); + Files.createDirectories(preExistingSibling); + Path sentinel = Files.writeString(preExistingSibling.resolve("owned-by-caller"), "preserve"); - CagraHnswBulkIndexWriter.indexFbin(fbinPath, configFor(dimension, k, true)); + try { + CagraHnswBulkIndexWriter.indexFbin(fbinPath, configFor(dimension, k, true)); - assertSearchable(numDocs, dataset, /* expectedSegments= */ k); + assertSearchable(numDocs, dataset, /* expectedSegments= */ k); + assertEquals("preserve", Files.readString(sentinel)); + } finally { + Files.deleteIfExists(sentinel); + Files.deleteIfExists(preExistingSibling); + } } @Test diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBulkIndexWriterValidation.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBulkIndexWriterValidation.java new file mode 100644 index 0000000000..4e58b39f34 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBulkIndexWriterValidation.java @@ -0,0 +1,216 @@ +/* + * 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.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +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.store.Directory; +import org.apache.lucene.store.FilterDirectory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.junit.Test; + +public class TestCagraHnswBulkIndexWriterValidation extends LuceneTestCase { + + @Test + public void testSliceCountValidationIsBoundedAndExact() { + for (int total : new int[] {-1, 0}) { + expectThrows( + IllegalArgumentException.class, () -> CagraHnswBulkIndexWriter.sliceEvenly(total, 1)); + } + expectThrows(IllegalArgumentException.class, () -> CagraHnswBulkIndexWriter.sliceEvenly(1, 0)); + expectThrows(IllegalArgumentException.class, () -> CagraHnswBulkIndexWriter.sliceEvenly(1, 2)); + expectThrows( + IllegalArgumentException.class, + () -> CagraHnswBulkIndexWriter.sliceEvenly(1, Integer.MAX_VALUE)); + expectThrows( + IllegalArgumentException.class, + () -> CagraHnswBulkIndexWriter.sliceEvenly(IndexWriter.MAX_DOCS + 1, 2)); + + List slices = CagraHnswBulkIndexWriter.sliceEvenly(5, 3); + assertEquals(3, slices.size()); + assertArrayEquals(new int[] {0, 2}, slices.get(0)); + assertArrayEquals(new int[] {2, 2}, slices.get(1)); + assertArrayEquals(new int[] {4, 1}, slices.get(2)); + } + + @Test + public void testPublicBuildRejectsEmptySourceWithoutCreatingTarget() throws Exception { + Path target = createTempDir().resolve("empty-index"); + IllegalArgumentException failure = + expectThrows( + IllegalArgumentException.class, + () -> CagraHnswBulkIndexWriter.build(new FixedShapeSource(0), config(target, 1))); + assertTrue(failure.getMessage(), failure.getMessage().contains("must be > 0")); + assertFalse(Files.exists(target)); + } + + @Test + public void testPublicBuildRejectsExtremeSegmentsWithoutCreatingTarget() throws Exception { + Path target = createTempDir().resolve("too-many-segments-index"); + IllegalArgumentException failure = + expectThrows( + IllegalArgumentException.class, + () -> + CagraHnswBulkIndexWriter.build( + new FixedShapeSource(1), config(target, Integer.MAX_VALUE))); + assertTrue(failure.getMessage(), failure.getMessage().contains("exceeds bulk vector count")); + assertFalse(Files.exists(target)); + } + + @Test + public void testPublicBuildRejectsTotalAboveLuceneLimitWithoutCreatingTarget() throws Exception { + Path target = createTempDir().resolve("too-many-documents-index"); + IllegalArgumentException failure = + expectThrows( + IllegalArgumentException.class, + () -> + CagraHnswBulkIndexWriter.build( + new FixedShapeSource(IndexWriter.MAX_DOCS + 1), config(target, 2))); + assertTrue(failure.getMessage(), failure.getMessage().contains("index-wide document limit")); + assertFalse(Files.exists(target)); + } + + @Test + public void testOneShotBuildRequiresTarget() { + IllegalArgumentException failure = + expectThrows( + IllegalArgumentException.class, + () -> CagraHnswBulkIndexWriter.build(new FixedShapeSource(1), config())); + assertTrue(failure.getMessage(), failure.getMessage().contains("targetDirectory")); + } + + @Test + public void testTargetIsStoredAbsoluteAndNormalized() { + Path supplied = Path.of("build-area", "..", "index"); + assertEquals(supplied.toAbsolutePath().normalize(), config(supplied, 1).targetDirectory()); + } + + @Test + public void testFbinCannotAlsoBeTheTargetDirectory() throws Exception { + Path source = createTempDir().resolve("source.fbin"); + Files.write(source, new byte[8]); + IllegalArgumentException failure = + expectThrows( + IllegalArgumentException.class, + () -> CagraHnswBulkIndexWriter.indexFbin(source, config(source, 1))); + assertTrue(failure.getMessage(), failure.getMessage().contains("must not be the source FBIN")); + } + + @Test + public void testExactCountAcceptsLuceneSegmentLimitWithoutOverflow() throws Exception { + try (Directory directory = newDirectory()) { + CagraHnswBulkIndexWriter writer = + new CagraHnswBulkIndexWriter( + directory, new IndexWriterConfig(), config(), IndexWriter.MAX_DOCS); + writer.abort(); + } + } + + @Test + public void testExactCountRejectsValuesAboveLuceneIndexLimitBeforeWriterCreation() + throws Exception { + for (int count : new int[] {IndexWriter.MAX_DOCS + 1, Integer.MAX_VALUE}) { + try (Directory directory = newDirectory()) { + IllegalArgumentException failure = + expectThrows( + IllegalArgumentException.class, + () -> + new CagraHnswBulkIndexWriter( + directory, new IndexWriterConfig(), config(), count)); + assertTrue( + failure.getMessage(), failure.getMessage().contains("index-wide document limit")); + try (var lock = directory.obtainLock(IndexWriter.WRITE_LOCK_NAME)) { + assertNotNull(lock); + } + } + } + } + + @Test + public void testFailedExplicitCommitIsNotRetriedByClose() throws Exception { + try (FailFirstCommitRenameDirectory directory = + new FailFirstCommitRenameDirectory(newDirectory())) { + IndexWriterConfig owned = + CagraHnswBulkIndexWriter.ownedIndexWriterConfig( + new IndexWriterConfig(), Codec.getDefault(), 1); + IOException failure = + expectThrows( + IOException.class, + () -> { + try (IndexWriter writer = new IndexWriter(directory, owned)) { + Document document = new Document(); + document.add(new StringField("id", "0", Field.Store.NO)); + writer.addDocument(document); + writer.commit(); + } + }); + + assertEquals("first publication rename failed", failure.getMessage()); + assertEquals(1, directory.publicationRenameAttempts); + assertFalse(DirectoryReader.indexExists(directory)); + } + } + + private static CagraHnswBulkIndexWriter.Config config() { + return CagraHnswBulkIndexWriter.Config.builder() + .field("vector", 2, EUCLIDEAN) + .graphBuild(new AcceleratedHNSWParams.Builder().build()) + .build(); + } + + private static CagraHnswBulkIndexWriter.Config config(Path target, int segments) { + return CagraHnswBulkIndexWriter.Config.builder() + .field("vector", 2, EUCLIDEAN) + .graphBuild(new AcceleratedHNSWParams.Builder().build()) + .segments(segments, false) + .targetDirectory(target) + .build(); + } + + private record FixedShapeSource(int size) implements VectorSource { + @Override + public int dimensions() { + return 2; + } + + @Override + public void get(int index, float[] dst) { + throw new AssertionError("source must not be consumed during validation"); + } + + @Override + public void close() {} + } + + private static final class FailFirstCommitRenameDirectory extends FilterDirectory { + private int publicationRenameAttempts; + + private FailFirstCommitRenameDirectory(Directory in) { + super(in); + } + + @Override + public void rename(String source, String dest) throws IOException { + if (source.startsWith("pending_segments") && dest.startsWith("segments")) { + publicationRenameAttempts++; + if (publicationRenameAttempts == 1) { + throw new IOException("first publication rename failed"); + } + } + super.rename(source, dest); + } + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java index 996c1d7575..24182dc75e 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java @@ -154,9 +154,6 @@ public void testCagraToHnswSerializationAndSearchWithFallbackWriter() throws Exc "Id: " + id + " expected but not found", expectedIds.contains(Integer.valueOf(id))); } assertTrue("TopK results not returned", results.scoreDocs.length == topK); - - } catch (Exception e) { - e.printStackTrace(); } } } diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVS2510GPUVectorsReaderScoreNormalization.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVS2510GPUVectorsReaderScoreNormalization.java deleted file mode 100644 index fafc58d1b2..0000000000 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVS2510GPUVectorsReaderScoreNormalization.java +++ /dev/null @@ -1,41 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -package com.nvidia.cuvs.lucene; - -import org.apache.lucene.index.VectorSimilarityFunction; -import org.apache.lucene.tests.util.LuceneTestCase; -import org.junit.Test; - -public class TestCuVS2510GPUVectorsReaderScoreNormalization extends LuceneTestCase { - - @Test - public void testMaximumInnerProductValueUsesLuceneScoreContract() { - assertEquals( - 13.0f, - CuVS2510GPUVectorsReader.toLuceneScore( - VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT, 12.0f), - 0.0f); - assertEquals( - 0.2f, - CuVS2510GPUVectorsReader.toLuceneScore( - VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT, -4.0f), - 0.0f); - } - - @Test - public void testOtherCuVSValuesUseTheirLuceneScoreContracts() { - assertEquals( - 0.2f, - CuVS2510GPUVectorsReader.toLuceneScore(VectorSimilarityFunction.EUCLIDEAN, 4.0f), - 0.0f); - assertEquals( - 0.75f, - CuVS2510GPUVectorsReader.toLuceneScore(VectorSimilarityFunction.DOT_PRODUCT, 0.5f), - 0.0f); - assertEquals( - 0.75f, CuVS2510GPUVectorsReader.toLuceneScore(VectorSimilarityFunction.COSINE, 0.5f), 0.0f); - } -} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestExternalFbinFlatVectorsReader.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestExternalFbinFlatVectorsReader.java index fec0704bbf..284acc5ffe 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestExternalFbinFlatVectorsReader.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestExternalFbinFlatVectorsReader.java @@ -15,7 +15,9 @@ import java.nio.file.StandardOpenOption; import java.security.MessageDigest; import java.util.HexFormat; +import java.util.List; import java.util.Map; +import java.util.function.Consumer; import org.apache.lucene.codecs.Codec; import org.apache.lucene.codecs.CodecUtil; import org.apache.lucene.index.CorruptIndexException; @@ -25,6 +27,7 @@ import org.apache.lucene.index.FieldInfos; import org.apache.lucene.index.FloatVectorValues; import org.apache.lucene.index.IndexFileNames; +import org.apache.lucene.index.IndexFormatTooNewException; import org.apache.lucene.index.IndexOptions; import org.apache.lucene.index.SegmentInfo; import org.apache.lucene.index.SegmentReadState; @@ -35,6 +38,7 @@ import org.apache.lucene.store.IndexInput; import org.apache.lucene.store.IndexOutput; import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.util.IOUtils; import org.apache.lucene.util.InfoStream; import org.apache.lucene.util.StringHelper; import org.apache.lucene.util.Version; @@ -104,6 +108,140 @@ public void testCorruptDescriptorChecksumIsRejected() throws Exception { } } + public void testIndependentlyEncodedVersionZeroDescriptorIsReadable() throws Exception { + float[][] vectors = {{1.0f, 2.0f}, {3.0f, 4.0f}, {5.0f, 6.0f}}; + Path fbin = writeFbin(createTempDir().resolve("literal-v0.fbin"), vectors); + + try (Directory indexDirectory = newDirectory(); + ExternalFbinFileRegistry.Registration registration = + ExternalFbinFileRegistry.register(fbin, sha256(fbin))) { + ExternalFbinReference reference = registration.reference(1, 2); + SegmentFixture fixture = newFixture(indexDirectory, reference.rows(), reference.dimensions()); + writeLiteralDescriptor(fixture, new LiteralDescriptor(reference)); + + try (ExternalFbinFlatVectorsReader reader = + new ExternalFbinFlatVectorsReader(fixture.readState())) { + FloatVectorValues values = reader.getFloatVectorValues("vector"); + assertEquals(2, values.size()); + assertArrayEquals(vectors[1], values.vectorValue(0), 0.0f); + assertArrayEquals(vectors[2], values.vectorValue(1), 0.0f); + reader.checkIntegrity(); + } + } + } + + public void testWrongRegisteredShapeFailsWithContextualCheckedException() throws Exception { + float[][] expectedVectors = {{1.0f, 2.0f}, {3.0f, 4.0f}, {5.0f, 6.0f}}; + Path root = createTempDir(); + Path expectedFbin = writeFbin(root.resolve("expected.fbin"), expectedVectors); + String expectedDigest = sha256(expectedFbin); + ExternalFbinReference reference; + try (ExternalFbinFileRegistry.Registration registration = + ExternalFbinFileRegistry.register(expectedFbin, expectedDigest)) { + reference = registration.reference(1, 2); + } + Path wrongFbin = writeFbin(root.resolve("wrong.fbin"), new float[][] {{7.0f, 8.0f}}); + + try (Directory indexDirectory = newDirectory(); + ExternalFbinFileRegistry.Registration ignored = + ExternalFbinFileRegistry.register(wrongFbin, expectedDigest)) { + SegmentFixture fixture = newFixture(indexDirectory, reference.rows(), reference.dimensions()); + writeLiteralDescriptor(fixture, new LiteralDescriptor(reference)); + + IOException failure = + expectThrows( + IOException.class, () -> new ExternalFbinFlatVectorsReader(fixture.readState())); + assertTrue( + failure.getMessage(), failure.getMessage().contains(wrongFbin.toRealPath().toString())); + assertTrue(failure.getMessage(), failure.getMessage().contains(reference.contentId())); + assertTrue(failure.getCause() instanceof IllegalArgumentException); + } + } + + public void testWrongRegisteredOverflowingHeaderFailsWithContextualCheckedException() + throws Exception { + float[][] expectedVectors = {{1.0f, 2.0f}, {3.0f, 4.0f}, {5.0f, 6.0f}}; + Path root = createTempDir(); + Path expectedFbin = writeFbin(root.resolve("expected.fbin"), expectedVectors); + String expectedDigest = sha256(expectedFbin); + ExternalFbinReference reference; + try (ExternalFbinFileRegistry.Registration registration = + ExternalFbinFileRegistry.register(expectedFbin, expectedDigest)) { + reference = registration.reference(1, 2); + } + Path wrongFbin = root.resolve("overflow.fbin"); + writeFbinHeader(wrongFbin, Integer.MAX_VALUE, Integer.MAX_VALUE); + + try (Directory indexDirectory = newDirectory(); + ExternalFbinFileRegistry.Registration ignored = + ExternalFbinFileRegistry.register(wrongFbin, expectedDigest)) { + SegmentFixture fixture = newFixture(indexDirectory, reference.rows(), reference.dimensions()); + writeLiteralDescriptor(fixture, new LiteralDescriptor(reference)); + + IOException failure = + expectThrows( + IOException.class, () -> new ExternalFbinFlatVectorsReader(fixture.readState())); + assertTrue( + failure.getMessage(), failure.getMessage().contains(wrongFbin.toRealPath().toString())); + assertTrue(failure.getMessage(), failure.getMessage().contains(reference.contentId())); + assertTrue(failure.getCause() instanceof ArithmeticException); + } + } + + public void testIndependentlyEncodedSemanticDescriptorFailures() throws Exception { + float[][] vectors = {{1.0f, 2.0f}, {3.0f, 4.0f}, {5.0f, 6.0f}}; + Path fbin = writeFbin(createTempDir().resolve("semantic-errors.fbin"), vectors); + + try (ExternalFbinFileRegistry.Registration registration = + ExternalFbinFileRegistry.register(fbin, sha256(fbin))) { + ExternalFbinReference reference = registration.reference(1, 2); + List mutations = + List.of( + new DescriptorMutation( + "invalid encoding", descriptor -> descriptor.encodingName = "NOT_AN_ENCODING"), + new DescriptorMutation("unknown field", descriptor -> descriptor.fieldNumber = 7), + new DescriptorMutation( + "invalid range", descriptor -> descriptor.payloadLength = Long.BYTES), + new DescriptorMutation("inconsistent rows", descriptor -> descriptor.fileRows += 1), + new DescriptorMutation("extra field", descriptor -> descriptor.sentinel = 0), + new DescriptorMutation( + "trailing payload", descriptor -> descriptor.trailingByte = true)); + + for (DescriptorMutation mutation : mutations) { + try (Directory indexDirectory = newDirectory()) { + SegmentFixture fixture = + newFixture(indexDirectory, reference.rows(), reference.dimensions()); + LiteralDescriptor descriptor = new LiteralDescriptor(reference); + mutation.mutation().accept(descriptor); + writeLiteralDescriptor(fixture, descriptor); + expectThrows( + CorruptIndexException.class, + mutation.name(), + () -> new ExternalFbinFlatVectorsReader(fixture.readState())); + } + } + } + } + + public void testIndependentlyEncodedFutureDescriptorVersionIsRejected() throws Exception { + float[][] vectors = {{1.0f, 2.0f}, {3.0f, 4.0f}, {5.0f, 6.0f}}; + Path fbin = writeFbin(createTempDir().resolve("future-version.fbin"), vectors); + + try (Directory indexDirectory = newDirectory(); + ExternalFbinFileRegistry.Registration registration = + ExternalFbinFileRegistry.register(fbin, sha256(fbin))) { + ExternalFbinReference reference = registration.reference(1, 2); + SegmentFixture fixture = newFixture(indexDirectory, reference.rows(), reference.dimensions()); + LiteralDescriptor descriptor = new LiteralDescriptor(reference); + descriptor.version = 1; + writeLiteralDescriptor(fixture, descriptor); + + expectThrows( + IndexFormatTooNewException.class, + () -> new ExternalFbinFlatVectorsReader(fixture.readState())); + } + } + public void testDimensionBoundariesWithSlicedRowsAcrossMmapChunks() throws Exception { for (int dimensions : GENERALIZATION_DIMENSIONS) { float[][] vectors = createVectors(5, dimensions); @@ -132,6 +270,66 @@ public void testDimensionBoundariesWithSlicedRowsAcrossMmapChunks() throws Excep } } + public void testSharedMappingOwnerSurvivesStaggeredReaderClose() throws Exception { + int dimensions = 64; + float[][] vectors = createVectors(6, dimensions); + Path fbin = writeFbin(createTempDir().resolve("shared-mapping.fbin"), vectors); + + try (Directory firstIndex = newDirectory(); + Directory secondIndex = newDirectory(); + ExternalFbinFileRegistry.Registration registration = + ExternalFbinFileRegistry.register(fbin, sha256(fbin))) { + ExternalFbinReference firstReference = registration.reference(0, 3); + ExternalFbinReference secondReference = registration.reference(3, 3); + SegmentFixture firstFixture = writeDescriptor(firstIndex, firstReference); + SegmentFixture secondFixture = writeDescriptor(secondIndex, secondReference); + ExternalFbinFlatVectorsReader first = null; + ExternalFbinFlatVectorsReader second = null; + try { + first = new ExternalFbinFlatVectorsReader(firstFixture.readState(), FORCED_MMAP_CHUNK_SIZE); + second = + new ExternalFbinFlatVectorsReader(secondFixture.readState(), FORCED_MMAP_CHUNK_SIZE); + assertEquals( + 1, + ExternalFbinMappedInputCache.activeOwnerCountForTests( + fbin, firstReference, FORCED_MMAP_CHUNK_SIZE)); + assertEquals( + 2, + ExternalFbinMappedInputCache.activeLeaseCountForTests( + fbin, firstReference, FORCED_MMAP_CHUNK_SIZE)); + assertArrayEquals(vectors[0], first.getFloatVectorValues("vector").vectorValue(0), 0.0f); + assertArrayEquals(vectors[5], second.getFloatVectorValues("vector").vectorValue(2), 0.0f); + first.checkIntegrity(); + + first.close(); + first = null; + assertEquals( + 1, + ExternalFbinMappedInputCache.activeOwnerCountForTests( + fbin, secondReference, FORCED_MMAP_CHUNK_SIZE)); + assertEquals( + 1, + ExternalFbinMappedInputCache.activeLeaseCountForTests( + fbin, secondReference, FORCED_MMAP_CHUNK_SIZE)); + assertArrayEquals(vectors[5], second.getFloatVectorValues("vector").vectorValue(2), 0.0f); + second.checkIntegrity(); + + second.close(); + second = null; + assertEquals( + 0, + ExternalFbinMappedInputCache.activeOwnerCountForTests( + fbin, secondReference, FORCED_MMAP_CHUNK_SIZE)); + assertEquals( + 0, + ExternalFbinMappedInputCache.activeLeaseCountForTests( + fbin, secondReference, FORCED_MMAP_CHUNK_SIZE)); + } finally { + IOUtils.closeWhileHandlingException(first, second); + } + } + } + public void testSparseFileReadsFinalRowAcrossTwoGiBBoundary() throws Exception { int dimensions = 64; long rowBytes = (long) dimensions * Float.BYTES; @@ -334,6 +532,16 @@ public void testOrphanDescriptorDoesNotSelectExternalReader() throws Exception { private static SegmentFixture writeDescriptor( Directory directory, ExternalFbinReference reference) throws Exception { + SegmentFixture fixture = newFixture(directory, reference.rows(), reference.dimensions()); + try (ExternalFbinReferenceWriter writer = + new ExternalFbinReferenceWriter(fixture.writeState())) { + writer.writeField(fixture.writeState().fieldInfos.fieldInfo("vector"), reference); + writer.finish(); + } + return fixture; + } + + private static SegmentFixture newFixture(Directory directory, int rows, int dimensions) { FieldInfo fieldInfo = new FieldInfo( "vector", @@ -349,7 +557,7 @@ private static SegmentFixture writeDescriptor( 0, 0, 0, - reference.dimensions(), + dimensions, VectorEncoding.FLOAT32, EUCLIDEAN, false, @@ -361,7 +569,7 @@ private static SegmentFixture writeDescriptor( Version.LATEST, Version.LATEST, "_0", - reference.rows(), + rows, false, false, Codec.getDefault(), @@ -369,18 +577,45 @@ private static SegmentFixture writeDescriptor( StringHelper.randomId(), Map.of(), null); - segmentInfo.putAttribute( - ExternalFbinReferenceWriter.segmentAttribute(""), - ExternalFbinReferenceWriter.SEGMENT_ATTRIBUTE_VALUE); + segmentInfo.putAttribute("cuvs.external_fbin.", "1"); SegmentWriteState writeState = new SegmentWriteState( InfoStream.NO_OUTPUT, directory, segmentInfo, fieldInfos, null, IOContext.DEFAULT); - try (ExternalFbinReferenceWriter writer = new ExternalFbinReferenceWriter(writeState)) { - writer.writeField(fieldInfo, reference); - writer.finish(); - } return new SegmentFixture( - new SegmentReadState(directory, segmentInfo, fieldInfos, IOContext.DEFAULT)); + writeState, new SegmentReadState(directory, segmentInfo, fieldInfos, IOContext.DEFAULT)); + } + + /** Writes the version-zero wire layout without calling the production descriptor writer. */ + private static void writeLiteralDescriptor(SegmentFixture fixture, LiteralDescriptor descriptor) + throws IOException { + SegmentWriteState state = fixture.writeState(); + String fileName = + IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, "vefr"); + try (IndexOutput output = state.directory.createOutput(fileName, state.context)) { + CodecUtil.writeIndexHeader( + output, + "CuVSExternalFbinReference", + descriptor.version, + state.segmentInfo.getId(), + state.segmentSuffix); + output.writeInt(descriptor.fieldNumber); + output.writeString(descriptor.fieldName); + output.writeString(descriptor.encodingName); + output.writeString(descriptor.similarityName); + output.writeVInt(descriptor.dimensions); + output.writeInt(descriptor.rows); + output.writeInt(descriptor.fileRows); + output.writeLong(descriptor.firstRow); + output.writeLong(descriptor.fileLength); + output.writeLong(descriptor.payloadOffset); + output.writeLong(descriptor.payloadLength); + output.writeBytes(descriptor.sha256, descriptor.sha256.length); + output.writeInt(descriptor.sentinel); + if (descriptor.trailingByte) { + output.writeByte((byte) 0x5a); + } + CodecUtil.writeFooter(output); + } } private static void corruptDescriptorPayload(Directory directory, SegmentReadState state) @@ -429,6 +664,19 @@ private static Path writeFbin(Path file, float[][] vectors) throws IOException { return file; } + private static void writeFbinHeader(Path file, int rows, int dimensions) throws IOException { + ByteBuffer header = + ByteBuffer.allocate((int) ExternalFbinReference.HEADER_BYTES) + .order(ByteOrder.LITTLE_ENDIAN) + .putInt(rows) + .putInt(dimensions); + header.flip(); + try (FileChannel channel = + FileChannel.open(file, StandardOpenOption.CREATE_NEW, StandardOpenOption.WRITE)) { + writeFully(channel, header, 0L); + } + } + private static void writeSparseFbin( Path file, int rows, int dimensions, long finalRowOffset, float[] finalRow) throws IOException { @@ -510,7 +758,41 @@ private static String sha256(Path file) throws Exception { return HexFormat.of().formatHex(digest.digest()); } - private record SegmentFixture(SegmentReadState readState) {} + private static final class LiteralDescriptor { + private int version = 0; + private int fieldNumber = 0; + private String fieldName = "vector"; + private String encodingName = "FLOAT32"; + private String similarityName = "EUCLIDEAN"; + private int dimensions; + private int rows; + private int fileRows; + private long firstRow; + private long fileLength; + private long payloadOffset; + private long payloadLength; + private byte[] sha256; + private int sentinel = -1; + private boolean trailingByte; + + private LiteralDescriptor(ExternalFbinReference reference) { + dimensions = reference.dimensions(); + rows = reference.rows(); + fileRows = + Math.toIntExact( + (reference.fileLength() - ExternalFbinReference.HEADER_BYTES) + / ((long) dimensions * Float.BYTES)); + firstRow = reference.firstRow(); + fileLength = reference.fileLength(); + payloadOffset = reference.payloadOffset(); + payloadLength = reference.payloadLength(); + sha256 = reference.sha256(); + } + } + + private record DescriptorMutation(String name, Consumer mutation) {} + + private record SegmentFixture(SegmentWriteState writeState, SegmentReadState readState) {} private record SparseRow(long offset, float[] values) {} } diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMappedSelfContainedOutputCancellation.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMappedSelfContainedOutputCancellation.java new file mode 100644 index 0000000000..4f29302531 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMappedSelfContainedOutputCancellation.java @@ -0,0 +1,190 @@ +/* + * 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.io.InterruptedIOException; +import java.lang.foreign.MemorySegment; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import org.apache.lucene.store.DataOutput; +import org.apache.lucene.tests.util.LuceneTestCase; + +public class TestMappedSelfContainedOutputCancellation extends LuceneTestCase { + + public void testGraphFailureCancelsAndWaitsForFlatWorker() throws Exception { + CountDownLatch flatStarted = new CountDownLatch(1); + CountDownLatch flatInterrupted = new CountDownLatch(1); + CountDownLatch releaseFlat = new CountDownLatch(1); + AtomicBoolean flatExited = new AtomicBoolean(); + AtomicBoolean returnedAfterFlatExit = new AtomicBoolean(); + AtomicReference observed = new AtomicReference<>(); + IOException graphFailure = new IOException("graph failed"); + IOException flatFailure = new IOException("flat worker failed after cancellation"); + + Thread caller = + Thread.ofPlatform() + .name("mapped-overlap-graph-failure-test") + .unstarted( + () -> { + try { + MappedSelfContainedOutput.runOverlapped( + () -> { + flatStarted.countDown(); + try { + new CountDownLatch(1).await(); + throw new AssertionError("flat worker unexpectedly resumed"); + } catch (InterruptedException expected) { + flatInterrupted.countDown(); + awaitUninterruptibly(releaseFlat); + throw flatFailure; + } finally { + flatExited.set(true); + } + }, + () -> { + assertTrue(flatStarted.await(10L, TimeUnit.SECONDS)); + throw graphFailure; + }); + observed.set(new AssertionError("expected graph failure")); + } catch (Throwable thrown) { + observed.set(thrown); + } finally { + returnedAfterFlatExit.set(flatExited.get()); + } + }); + + caller.start(); + try { + assertTrue(flatInterrupted.await(10L, TimeUnit.SECONDS)); + assertFalse(flatExited.get()); + } finally { + releaseFlat.countDown(); + caller.join(10_000L); + if (caller.isAlive()) { + caller.interrupt(); + } + } + + assertFalse("overlap caller did not terminate", caller.isAlive()); + assertSame(graphFailure, observed.get()); + assertTrue(flatExited.get()); + assertTrue("overlap returned before the flat worker exited", returnedAfterFlatExit.get()); + assertArrayEquals(new Throwable[] {flatFailure}, graphFailure.getSuppressed()); + } + + public void testCallerInterruptionCancelsWorkerAndRestoresInterrupt() throws Exception { + CountDownLatch flatStarted = new CountDownLatch(1); + CountDownLatch graphFinished = new CountDownLatch(1); + CountDownLatch releaseFlat = new CountDownLatch(1); + AtomicBoolean flatInterrupted = new AtomicBoolean(); + AtomicBoolean flatExited = new AtomicBoolean(); + AtomicBoolean returnedAfterFlatExit = new AtomicBoolean(); + AtomicBoolean interruptRestored = new AtomicBoolean(); + AtomicReference observed = new AtomicReference<>(); + + Thread caller = + Thread.ofPlatform() + .name("mapped-overlap-caller-interruption-test") + .unstarted( + () -> { + try { + MappedSelfContainedOutput.runOverlapped( + () -> { + flatStarted.countDown(); + try { + releaseFlat.await(); + } catch (InterruptedException expected) { + flatInterrupted.set(true); + } finally { + flatExited.set(true); + } + }, + () -> { + assertTrue(flatStarted.await(10L, TimeUnit.SECONDS)); + graphFinished.countDown(); + }); + observed.set(new AssertionError("expected caller interruption")); + } catch (Throwable thrown) { + observed.set(thrown); + interruptRestored.set(Thread.currentThread().isInterrupted()); + } finally { + returnedAfterFlatExit.set(flatExited.get()); + } + }); + + caller.start(); + try { + assertTrue(graphFinished.await(10L, TimeUnit.SECONDS)); + caller.interrupt(); + caller.join(10_000L); + } finally { + releaseFlat.countDown(); + if (caller.isAlive()) { + caller.interrupt(); + } + } + + assertFalse("interrupted overlap caller did not terminate", caller.isAlive()); + assertTrue(observed.get() instanceof InterruptedIOException); + assertTrue(flatInterrupted.get()); + assertTrue(flatExited.get()); + assertTrue("overlap returned before the flat worker exited", returnedAfterFlatExit.get()); + assertTrue("caller interrupt status was not restored", interruptRestored.get()); + } + + public void testRawCopyChecksInterruptionBetweenChunks() throws Exception { + AtomicInteger writes = new AtomicInteger(); + DataOutput interruptingOutput = + new DataOutput() { + @Override + public void writeByte(byte value) { + throw new AssertionError("unexpected single-byte write"); + } + + @Override + public void writeBytes(byte[] bytes, int offset, int length) { + if (writes.incrementAndGet() == 1) { + Thread.currentThread().interrupt(); + } + } + }; + + try { + InterruptedIOException failure = + expectThrows( + InterruptedIOException.class, + () -> + NativeFlatVectorsWriter.copyRawFloat32Vectors( + MemorySegment.ofArray(new byte[8]), + 8L, + interruptingOutput, + new CagraHnswBuildMetrics(), + 4)); + assertTrue(failure.getMessage(), failure.getMessage().contains("at byte 4 of 8")); + assertEquals(1, writes.get()); + assertTrue(Thread.currentThread().isInterrupted()); + } finally { + Thread.interrupted(); + } + } + + private static void awaitUninterruptibly(CountDownLatch latch) { + boolean interrupted = false; + while (latch.getCount() != 0L) { + try { + latch.await(); + } catch (InterruptedException expected) { + interrupted = true; + } + } + if (interrupted) { + Thread.currentThread().interrupt(); + } + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferedCleanup.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferedCleanup.java index 4ee42dc4b4..ecd8d2de96 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferedCleanup.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferedCleanup.java @@ -4,12 +4,26 @@ */ package com.nvidia.cuvs.lucene; +import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.CuVSResources; import java.io.Closeable; import java.io.IOException; import java.lang.reflect.Proxy; +import java.util.Map; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FieldInfos; +import org.apache.lucene.index.SegmentInfo; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.tests.store.MockDirectoryWrapper; import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.util.InfoStream; +import org.apache.lucene.util.StringHelper; +import org.apache.lucene.util.Version; import org.junit.Test; public class TestNativeFlatBufferedCleanup extends LuceneTestCase { @@ -33,6 +47,92 @@ public void testThreadResourcesCloseWhenOutputCloseFails() { assertTrue("cuVS resources were not closed", resourcesClosed.get()); } + @Test + public void testBorrowedWriterConstructorFailureClosesThreadResource() throws Exception { + AtomicInteger closeCount = new AtomicInteger(); + AtomicReference constructorFailure = new AtomicReference<>(); + try (MockDirectoryWrapper directory = newMockDirectory()) { + SegmentInfo segmentInfo = + new SegmentInfo( + directory, + Version.LATEST, + Version.LATEST, + "_0", + 1, + false, + false, + Codec.getDefault(), + Map.of(), + StringHelper.randomId(), + Map.of(), + null); + SegmentWriteState state = + new SegmentWriteState( + InfoStream.NO_OUTPUT, + directory, + segmentInfo, + new FieldInfos(new FieldInfo[0]), + null, + IOContext.DEFAULT); + + Thread constructorThread = + new Thread( + () -> { + ThreadLocalCuVSResourcesProvider.setCuVSResourcesInstance( + resourcesThatCountsClose(closeCount)); + try { + new BorrowedDatasetHnswVectorsWriter( + state, + new AcceleratedHNSWParams.Builder().build(), + BulkIndexingContext.nativeBuffered(1, new CagraHnswBuildMetrics())); + constructorFailure.set(new AssertionError("constructor unexpectedly succeeded")); + } catch (Throwable failure) { + constructorFailure.set(failure); + } + }, + "borrowed-writer-constructor-failure"); + constructorThread.start(); + constructorThread.join(); + + assertNotNull(constructorFailure.get()); + assertTrue(constructorFailure.get() instanceof IllegalArgumentException); + assertTrue( + constructorFailure.get().getMessage(), + constructorFailure.get().getMessage().contains("requires a mapped dataset")); + assertEquals(1, closeCount.get()); + assertEquals(0, directory.getFileHandleCount()); + } + } + + @Test + public void testOwnedDatasetCloseFailureIsSuppressedOntoWriteFailure() { + CuVSMatrix dataset = + (CuVSMatrix) + Proxy.newProxyInstance( + CuVSMatrix.class.getClassLoader(), + new Class[] {CuVSMatrix.class}, + (proxy, method, arguments) -> { + if (method.getName().equals("close")) { + throw new IllegalStateException("dataset close failed"); + } + throw new UnsupportedOperationException(method.getName()); + }); + + IOException failure = + expectThrows( + IOException.class, + () -> + AcceleratedHnswGraphOutput.withOwnedDataset( + dataset, + () -> { + throw new IOException("graph write failed"); + })); + + assertEquals("graph write failed", failure.getMessage()); + assertEquals(1, failure.getSuppressed().length); + assertEquals("dataset close failed", failure.getSuppressed()[0].getMessage()); + } + private static CuVSResources resourcesThatRecordClose(AtomicBoolean closed) { return (CuVSResources) Proxy.newProxyInstance( @@ -46,4 +146,18 @@ private static CuVSResources resourcesThatRecordClose(AtomicBoolean closed) { throw new UnsupportedOperationException(method.getName()); }); } + + private static CuVSResources resourcesThatCountsClose(AtomicInteger closeCount) { + return (CuVSResources) + Proxy.newProxyInstance( + CuVSResources.class.getClassLoader(), + new Class[] {CuVSResources.class}, + (proxy, method, arguments) -> { + if (method.getName().equals("close")) { + closeCount.incrementAndGet(); + return null; + } + throw new UnsupportedOperationException(method.getName()); + }); + } } diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestSingleVectorHnswGraphCleanup.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestSingleVectorHnswGraphCleanup.java new file mode 100644 index 0000000000..ddf2952fa0 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestSingleVectorHnswGraphCleanup.java @@ -0,0 +1,77 @@ +/* + * 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.CuVSMatrix; +import java.lang.reflect.Proxy; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.junit.Test; + +public class TestSingleVectorHnswGraphCleanup extends LuceneTestCase { + + @Test + public void testOwnedAdjacencyMatrixIsClosedAfterMaterialization() throws Throwable { + AtomicBoolean rowRead = new AtomicBoolean(); + AtomicInteger closeCount = new AtomicInteger(); + CuVSMatrix matrix = matrix(rowRead, closeCount, null, null); + + GPUBuiltHnswGraph graph = AcceleratedHNSWUtils.createSingleVectorHnswGraph(1, 4, matrix); + + assertTrue("the adjacency row was not materialized", rowRead.get()); + assertEquals("the owned matrix must be closed exactly once", 1, closeCount.get()); + assertEquals(1, graph.size()); + assertEquals(1, graph.numLevels()); + } + + @Test + public void testCloseFailureIsSuppressedOntoMaterializationFailure() { + AtomicBoolean rowRead = new AtomicBoolean(); + AtomicInteger closeCount = new AtomicInteger(); + IllegalStateException materializationFailure = + new IllegalStateException("materialization failed"); + IllegalStateException closeFailure = new IllegalStateException("matrix close failed"); + CuVSMatrix matrix = matrix(rowRead, closeCount, materializationFailure, closeFailure); + + IllegalStateException failure = + expectThrows( + IllegalStateException.class, + () -> AcceleratedHNSWUtils.createSingleVectorHnswGraph(1, 4, matrix)); + + assertSame(materializationFailure, failure); + assertTrue("the failing adjacency row was not read", rowRead.get()); + assertEquals("the failing owned matrix must be closed exactly once", 1, closeCount.get()); + assertArrayEquals(new Throwable[] {closeFailure}, failure.getSuppressed()); + } + + private static CuVSMatrix matrix( + AtomicBoolean rowRead, + AtomicInteger closeCount, + RuntimeException materializationFailure, + RuntimeException closeFailure) { + return (CuVSMatrix) + Proxy.newProxyInstance( + CuVSMatrix.class.getClassLoader(), + new Class[] {CuVSMatrix.class}, + (proxy, method, arguments) -> { + if (method.getName().equals("getRow")) { + rowRead.set(true); + if (materializationFailure != null) { + throw materializationFailure; + } + return null; + } + if (method.getName().equals("close")) { + closeCount.incrementAndGet(); + if (closeFailure != null) { + throw closeFailure; + } + return null; + } + throw new UnsupportedOperationException(method.getName()); + }); + } +}