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()); + }); + } +}