diff --git a/java/cuvs-lucene/examples/README.md b/java/cuvs-lucene/examples/README.md index 19013675b6..0d61d0a7c0 100644 --- a/java/cuvs-lucene/examples/README.md +++ b/java/cuvs-lucene/examples/README.md @@ -32,3 +32,17 @@ To run the Index and Search on GPU example do: ```sh mvn clean install && java -Djava.util.logging.config.file=src/main/resources/logging.properties -cp target/examples-26.10.0-jar-with-merged-services.jar com.nvidia.cuvs.lucene.examples.IndexAndSearchonGPUExample ``` + +To run the optimized CAGRA-HNSW build example (reference pattern for efficiently building an +accelerated HNSW index from a large `.fbin` with every ingest-side knob on — open the file once and +stream sequential prefetched chunks that overlap the disk read with indexing, hold at most two chunks +in memory, reuse a single vector array, size a native flat buffer per segment, auto-select the CAGRA +graph-build algorithm, and optionally partition into K segments built sequentially or overlapped) do: + +```sh +mvn clean install && java -Djava.util.logging.config.file=src/main/resources/logging.properties -cp target/examples-26.10.0-jar-with-merged-services.jar com.nvidia.cuvs.lucene.examples.OptimizedCagraHnswBuildExample +``` + +With no arguments it generates and indexes a small demo `.fbin` as a single segment; pass a real file, +chunk size, segment count, and overlap flag as +`... OptimizedCagraHnswBuildExample `. diff --git a/java/cuvs-lucene/examples/pom.xml b/java/cuvs-lucene/examples/pom.xml index 4432883f64..e8b1833c8a 100644 --- a/java/cuvs-lucene/examples/pom.xml +++ b/java/cuvs-lucene/examples/pom.xml @@ -92,6 +92,11 @@ lucene-backward-codecs 10.2.0 + + org.apache.lucene + lucene-misc + 10.2.0 + commons-io commons-io diff --git a/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java b/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java new file mode 100644 index 0000000000..f76b899243 --- /dev/null +++ b/java/cuvs-lucene/examples/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java @@ -0,0 +1,657 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene.examples; + +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; +import com.nvidia.cuvs.lucene.AcceleratedHNSWParams; +import com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWCodec; +import com.nvidia.cuvs.spi.CuVSProvider; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.channels.FileChannel; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.nio.file.StandardOpenOption; +import java.util.ArrayList; +import java.util.List; +import java.util.Random; +import java.util.UUID; +import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.Semaphore; +import java.util.logging.Logger; +import org.apache.commons.io.FileUtils; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.NoMergePolicy; +import org.apache.lucene.misc.store.HardlinkCopyDirectoryWrapper; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; + +/** + * Reference pattern for building an accelerated HNSW index (whose graph is built on the GPU with + * CAGRA) from a LARGE {@code .fbin} vector file with every optimization turned on — spanning + * both how vectors are ingested and how the GPU build is configured and partitioned. cuvs-lucene is a + * Lucene codec that sits below {@code addDocument}, so how you read your source data and how you shape + * segments is application code — this example shows the full recipe. + * + *

The knobs, and why each one matters

+ * + *
    + *
  1. Streaming, prefetched, bounded-memory reads. {@link PrefetchingFbinReader} opens the + * file ONCE, reads it front-to-back in large sequential chunks (no per-vector seek/close "fd + * churn"), holds at most two chunks (no whole-file heap copy), fills the NEXT chunk on a + * background thread while the ingest thread drains the current one (disk read hidden behind + * indexing), and unpacks into a caller-reused {@code float[]} (no per-vector allocation — safe + * because Lucene copies the value eagerly inside {@code addDocument}). + *
  2. Native flat buffering. {@link AcceleratedHNSWParams.Builder#withNumInputVectors} sizes + * a native host matrix to exactly the segment's vector count, so vectors stream straight into + * the buffer the GPU build consumes instead of piling up as a {@code List} on the JVM + * heap that must then be assembled. This removes the assembly copy and roughly halves peak host + * memory. It requires a single-segment build (no flush before commit, no merge). + *
  3. Automatic graph-build algorithm. {@code HEURISTIC} strategy defaults to + * {@code AUTO_SELECT}, letting cuVS pick the CAGRA build algorithm by dataset size (NN_DESCENT + * below ~5M vectors, IVF_PQ at or above) and auto-tune its parameters. NN_DESCENT generally + * reaches higher recall but takes longer to build; IVF_PQ is faster but slightly lower recall. + * Force one explicitly via {@code withCagraGraphBuildAlgo} only under expert guidance. + *
  4. Partitioned multi-segment build. Because native flat buffering is single-segment, "K + * segments" means K independent single-segment builds over contiguous slices, combined at the + * end. This is a deliberate memory/throughput lever — see the assumptions below. + *
+ * + *

Assumptions for the user-specified number of segments

+ * + *

{@code numSegments} is your choice and it is a trade-off, not a free speedup: + * + *

    + *
  • Each segment is one native-flat build over a {@code 1/K} slice, so peak host memory scales + * as {@code 1/K} (sequential mode) — this is the point of partitioning for large or + * memory-bounded datasets. + *
  • The GPU is serialized: only one CAGRA build runs on the device at a time. Extra + * segments do NOT parallelize the graph build; they only reduce host memory and (in overlap + * mode) let one segment's host-side ingest run during a prior segment's GPU commit. + *
  • Search fans out across all K segment graphs, so more segments trade a little query throughput + * (and can shift recall) for lower build-time memory. Pick K to fit your memory budget, not + * higher. + *
  • Each slice must stay a single segment: {@code maxBufferedDocs > sliceSize} and {@link + * NoMergePolicy} (no flush, no merge) so native flat buffering stays valid. + *
+ * + *

Sequential vs overlapped

+ * + *

With {@code overlap=false} the K slices are built as K sequential passes appended to one + * directory; peak host memory is a single slice's buffer ({@code N/K}). With {@code overlap=true} a + * bounded pool builds up to {@code PIPELINE_DEPTH} segments at once into their own directories, so a + * segment's ingest overlaps a prior segment's (serialized) GPU commit; the finished per-segment + * indexes are then combined by hardlinking their files into the final directory ({@link + * HardlinkCopyDirectoryWrapper} + {@code addIndexes}, no bulk copy of the vector data). Overlap costs + * up to {@code PIPELINE_DEPTH * (N/K)} peak host memory and can hide most of the ingest time behind + * the GPU build, bounded by disk contention between concurrent reads and the committing segment's + * writes. + * + *

Adapt {@link PrefetchingFbinReader} to your own source (a DB, object store, or stream). The + * properties above are what matter, not the {@code .fbin} specifics. + * + *

Usage: {@code OptimizedCagraHnswBuildExample [] [] [] + * []}. With no arguments a small demo {@code .fbin} is generated and indexed as a + * single segment. + */ +public class OptimizedCagraHnswBuildExample { + + private static final Logger log = + Logger.getLogger(OptimizedCagraHnswBuildExample.class.getName()); + private static final String ID_FIELD = "id"; + private static final String VECTOR_FIELD = "vector_field"; + + /** Max segments built concurrently in overlap mode; peak host memory is this many slice buffers. */ + private static final int PIPELINE_DEPTH = 2; + + public static void main(String[] args) throws Exception { + int chunkSizeMB = args.length >= 2 ? Integer.parseInt(args[1]) : 32; + int numSegments = args.length >= 3 ? Math.max(1, Integer.parseInt(args[2])) : 1; + boolean overlap = args.length >= 4 && Boolean.parseBoolean(args[3]); + Path indexDirPath = Paths.get(UUID.randomUUID().toString()); + + Path fbinPath; + boolean generated = false; + if (args.length >= 1) { + fbinPath = Paths.get(args[0]); + } else { + fbinPath = Paths.get("demo-" + UUID.randomUUID() + ".fbin"); + writeDemoFbin(fbinPath, 5000, 32, new Random(222)); + generated = true; + log.info("No .fbin provided; generated a demo file at " + fbinPath); + } + + // Must not be called when using a CPU-only Lucene codec: those paths never load the cuVS + // native library, so CuVSProvider resolves to UnsupportedProvider and throws. + CuVSProvider.provider().enableRMMAsyncMemory(); + try { + buildIndex(fbinPath, indexDirPath, chunkSizeMB, numSegments, overlap); + runSampleSearch(indexDirPath, fbinPath, 5); + } finally { + FileUtils.deleteDirectory(indexDirPath.toFile()); + if (generated) { + Files.deleteIfExists(fbinPath); + } + } + } + + /** + * Builds an accelerated HNSW index from {@code fbinPath} into {@code numSegments} contiguous + * slices, dispatching to the sequential or overlapped partitioned build. {@code numSegments == 1} + * is the plain single-segment native-flat build. + */ + private static void buildIndex( + Path fbinPath, Path indexDirPath, int chunkSizeMB, int numSegments, boolean overlap) + throws Exception { + int total; + int dim; + try (PrefetchingFbinReader reader = new PrefetchingFbinReader(fbinPath, 1)) { + total = reader.size(); + dim = reader.dimension(); + } + List slices = sliceEvenly(total, numSegments); // [start, size] per segment + + log.info( + "Indexing " + + total + + " vectors (" + + dim + + "-dim) from " + + fbinPath + + " into " + + slices.size() + + " segment(s), " + + (overlap && slices.size() > 1 ? "overlapped" : "sequential") + + " build, " + + chunkSizeMB + + " MB prefetched chunks"); + + if (overlap && slices.size() > 1) { + buildOverlapped(fbinPath, indexDirPath, chunkSizeMB, slices, dim); + } else { + buildSequential(fbinPath, indexDirPath, chunkSizeMB, slices, dim); + } + log.info("Index build complete: " + indexDirPath); + } + + /** + * Sequential partitioned build: one forward-only prefetch reader streamed front-to-back across all + * K slices, each slice built as a single native-flat segment appended to the same directory (first + * pass {@code CREATE}, later passes {@code APPEND}). Peak host memory is one slice's native buffer. + */ + private static void buildSequential( + Path fbinPath, Path indexDirPath, int chunkSizeMB, List slices, int dim) + throws Exception { + float[] scratch = new float[dim]; // one reusable array for the whole build + try (PrefetchingFbinReader reader = new PrefetchingFbinReader(fbinPath, chunkSizeMB); + Directory dir = FSDirectory.open(indexDirPath)) { + for (int p = 0; p < slices.size(); p++) { + int[] slice = slices.get(p); + log.info( + "Building segment " + + (p + 1) + + "/" + + slices.size() + + ": docs [" + + slice[0] + + ", " + + (slice[0] + slice[1]) + + ")"); + buildSegment(dir, reader, scratch, slice[0], slice[1], p == 0, null); + } + } + } + + /** + * Overlapped partitioned build: a bounded pool builds up to {@link #PIPELINE_DEPTH} segments at + * once, each into its OWN directory with its OWN prefetch reader over just its slice, so a + * segment's ingest overlaps a prior segment's GPU commit. The GPU build is serialized on a single + * permit. The finished per-segment indexes are then hardlinked into the final directory (no bulk + * copy of the vector data). + */ + private static void buildOverlapped( + Path fbinPath, Path indexDirPath, int chunkSizeMB, List slices, int dim) + throws Exception { + int depth = Math.min(slices.size(), PIPELINE_DEPTH); + int maxSlice = slices.stream().mapToInt(s -> s[1]).max().orElse(0); + double peakHostGb = (double) depth * maxSlice * dim * Float.BYTES / 1e9; + log.info( + "Overlapped build: " + + slices.size() + + " segment(s), pipeline depth " + + depth + + " (up to " + + depth + + " co-resident native host buffers, ~" + + String.format("%.2f", peakHostGb) + + " GB peak host)"); + + List segDirs = new ArrayList<>(); + for (int p = 0; p < slices.size(); p++) { + segDirs.add(Paths.get(indexDirPath + "_p" + p)); + } + // Start from fresh per-segment temp dirs, and always remove them afterwards (even on failure) + // so + // a crashed build does not leave orphaned per-segment indexes behind. + for (Path segDir : segDirs) { + FileUtils.deleteQuietly(segDir.toFile()); + } + try { + Semaphore gpuPermit = new Semaphore(1); // serialize the GPU CAGRA build across segments + ExecutorService pool = Executors.newFixedThreadPool(depth); + List> futures = new ArrayList<>(); + for (int p = 0; p < slices.size(); p++) { + int[] slice = slices.get(p); + Path segDir = segDirs.get(p); + futures.add( + pool.submit( + () -> { + float[] scratch = new float[dim]; + try (PrefetchingFbinReader reader = + new PrefetchingFbinReader(fbinPath, slice[0], slice[1], chunkSizeMB); + Directory d = FSDirectory.open(segDir)) { + // createNew=true: each segment is a fresh single-segment index in its own dir. + buildSegment(d, reader, scratch, slice[0], slice[1], true, gpuPermit); + } + return null; + })); + } + pool.shutdown(); + try { + for (Future f : futures) { + f.get(); // propagate any build failure + } + } finally { + pool.shutdownNow(); + } + combineByHardlink(indexDirPath, segDirs); + } finally { + for (Path segDir : segDirs) { + FileUtils.deleteQuietly(segDir.toFile()); + } + } + } + + /** + * Builds one segment from the contiguous slice {@code [start, start + size)}: a single-threaded + * {@link IndexWriter} whose codec's {@code numInputVectors} is sized to the slice (native flat + * buffering, single segment). When {@code gpuPermit} is non-null the commit — which runs the GPU + * CAGRA build — is serialized on it while other segments' host-side ingest may proceed. + */ + private static void buildSegment( + Directory dir, + PrefetchingFbinReader reader, + float[] scratch, + int start, + int size, + boolean createNew, + Semaphore gpuPermit) + throws Exception { + Codec codec = codecFor(size); + + // Keep this slice in ONE segment: raise the doc-count flush threshold above the slice size and + // disable RAM-based flushing so nothing flushes before commit, and forbid merges. This is what + // makes native flat buffering valid. Order matters: enable the doc-count trigger BEFORE + // disabling + // the RAM trigger, since Lucene rejects a config where both are disabled at once. + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(codec) + .setUseCompoundFile(false) + .setMaxBufferedDocs(Math.max(2, size + 1)) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) + .setMergePolicy(NoMergePolicy.INSTANCE) + .setOpenMode( + createNew ? IndexWriterConfig.OpenMode.CREATE : IndexWriterConfig.OpenMode.APPEND); + + try (IndexWriter writer = new IndexWriter(dir, config)) { + for (int i = 0; i < size; i++) { + int id = start + i; + reader.get(id, scratch); // sequential access -> served from the prefetched chunk, no alloc + Document doc = new Document(); + doc.add(new StringField(ID_FIELD, Integer.toString(id), Field.Store.YES)); + doc.add(new KnnFloatVectorField(VECTOR_FIELD, scratch, EUCLIDEAN)); + writer.addDocument(doc); // copies the vector -> 'scratch' is safe to reuse next iteration + } + // The single flush at commit is where the GPU CAGRA build runs; serialize it if asked. + if (gpuPermit != null) { + gpuPermit.acquire(); + try { + writer.commit(); + } finally { + gpuPermit.release(); + } + } else { + writer.commit(); + } + } + } + + /** + * Combines the per-segment indexes into {@code indexDirPath} by hardlinking their files (same + * filesystem) rather than copying the vector data. {@link HardlinkCopyDirectoryWrapper} falls back + * to a byte copy automatically if the segment dirs and the final dir are on different filesystems. + */ + private static void combineByHardlink(Path indexDirPath, List segDirs) throws IOException { + Directory[] sources = new Directory[segDirs.size()]; + try { + for (int i = 0; i < segDirs.size(); i++) { + sources[i] = FSDirectory.open(segDirs.get(i)); + } + IndexWriterConfig iwc = + new IndexWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE); // keep segments separate + try (Directory target = new HardlinkCopyDirectoryWrapper(FSDirectory.open(indexDirPath)); + IndexWriter combiner = new IndexWriter(target, iwc)) { + combiner.addIndexes(sources); + } + } finally { + for (Directory s : sources) { + if (s != null) { + s.close(); + } + } + } + } + + /** Builds a codec with all knobs on, sizing the native flat buffer to {@code numInputVectors}. */ + private static Codec codecFor(int numInputVectors) throws Exception { + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder() + // HEURISTIC lets cuVS pick the build algorithm and auto-tune its parameters based on + // maxConn and beamWidth below. + .withStrategy(AcceleratedHNSWParams.Strategy.HEURISTIC) + // Primary recall/graph-size knobs. Higher values improve recall at the cost of a + // larger graph and longer build. Match to your dataset and recall target. + .withMaxConn(32) + .withBeamWidth(32) + // Must match the distance metric used when querying the index. + .withCuvsDistanceType(CuvsDistanceType.L2Expanded) + // Starting point: one thread per logical CPU. Profile and tune for your hardware. + .withWriterThreads(Runtime.getRuntime().availableProcessors()) + // Native flat buffering: the value MUST equal the number of vectors actually ingested + // into this segment; the writer fails fast if they differ. If some input vectors are + // excluded (e.g. filtered during ingest), either pre-count the survivors in a separate + // pass or omit withNumInputVectors (pass 0) to fall back to the heap-buffered path, + // which buffers all vectors in a List on the JVM heap before building and + // therefore uses more peak host memory. Index-sorted segments (IndexWriterConfig + // .setIndexSort) are also unsupported, and binary/scalar quantized fields are not yet + // supported. + .withNumInputVectors(numInputVectors) + .build(); + return new Lucene101AcceleratedHNSWCodec(params); + } + + /** Splits {@code total} into {@code k} contiguous [start, size] slices, spreading the remainder. */ + private static List sliceEvenly(int total, int k) { + List slices = new ArrayList<>(); + int base = total / k; + int rem = total % k; + int start = 0; + for (int p = 0; p < k; p++) { + int size = base + (p < rem ? 1 : 0); // spread the remainder over the first slices + if (size <= 0) { + continue; + } + slices.add(new int[] {start, size}); + start += size; + } + return slices; + } + + /** Runs one k-NN query using the first vector in the file to show the index is searchable. */ + private static void runSampleSearch(Path indexDirPath, Path fbinPath, int topK) throws Exception { + float[] queryVector; + try (PrefetchingFbinReader reader = new PrefetchingFbinReader(fbinPath, 1)) { + queryVector = reader.get(0); + } + try (Directory dir = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(dir)) { + IndexSearcher searcher = new IndexSearcher(reader); + TopDocs results = + searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK), topK); + log.info("Sample search returned " + results.scoreDocs.length + " hits:"); + for (int i = 0; i < results.scoreDocs.length; i++) { + ScoreDoc sd = results.scoreDocs[i]; + String id = searcher.storedFields().document(sd.doc).get(ID_FIELD); + log.info(" rank " + (i + 1) + ": id=" + id + " score=" + sd.score); + } + } + } + + /** Writes a small random {@code .fbin} so the example is runnable without external data. */ + private static void writeDemoFbin(Path path, int numVectors, int dim, Random random) + throws IOException { + ByteBuffer buf = + ByteBuffer.allocate(8 + numVectors * dim * Float.BYTES).order(ByteOrder.LITTLE_ENDIAN); + buf.putInt(numVectors); // .fbin header: [num_vectors int32][dimension int32] + buf.putInt(dim); + for (int i = 0; i < numVectors; i++) { + for (int j = 0; j < dim; j++) { + buf.putFloat(random.nextFloat() * 100); + } + } + buf.flip(); + try (FileChannel ch = + FileChannel.open( + path, + StandardOpenOption.CREATE, + StandardOpenOption.WRITE, + StandardOpenOption.TRUNCATE_EXISTING)) { + while (buf.hasRemaining()) { + ch.write(buf); + } + } + } + + /** + * Prefetching, double-buffered sequential reader for uncompressed {@code .fbin} files ({@code + * [num_vectors int32][dim int32]} header, then contiguous little-endian float32 rows). + * + *

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

Forward-only, single-consumer. {@code get} must be called with non-decreasing indices + * from a single thread — the intended pattern for streaming ingestion. To read only a slice (e.g. + * one segment of a partitioned build), construct it with a {@code [firstVector, count)} range so + * each segment streams just its own portion of the file. + */ + static final class PrefetchingFbinReader implements AutoCloseable { + + private static final long HEADER_BYTES = 8; + + private static final class Chunk { + final ByteBuffer buf; + final long start; + final int len; + + Chunk(ByteBuffer buf, long start, int len) { + this.buf = buf; + this.start = start; + this.len = len; + } + } + + /** Sentinel placed on the ready queue once the reader has produced the final chunk. */ + private static final Chunk POISON = new Chunk(null, -1, 0); + + private final FileChannel channel; + private final int dimension; + private final int firstVector; // absolute index of the first vector this reader serves + private final int endVector; // absolute index just past the last vector this reader serves + private final int vectorBytes; + private final int chunkVectors; + + private final BlockingQueue free = new ArrayBlockingQueue<>(2); + private final BlockingQueue ready = new ArrayBlockingQueue<>(2); + private final Thread reader; + private volatile IOException readerError; + + private Chunk current; // consumer-owned; the chunk currently being served + + /** Reads the whole file from index 0. */ + PrefetchingFbinReader(Path path, int chunkSizeMB) throws IOException { + this(path, 0, -1, chunkSizeMB); + } + + /** + * Reads the contiguous range {@code [firstVector, firstVector + count)}, or to end of file if + * {@code count <= 0}. {@link #get} then serves absolute file indices within that range. + */ + PrefetchingFbinReader(Path path, int firstVector, int count, int chunkSizeMB) + throws IOException { + this.channel = FileChannel.open(path, StandardOpenOption.READ); + ByteBuffer header = ByteBuffer.allocate((int) HEADER_BYTES).order(ByteOrder.LITTLE_ENDIAN); + readFully(header, 0); + header.flip(); + int numVectors = header.getInt(); + this.dimension = header.getInt(); + this.vectorBytes = dimension * Float.BYTES; + this.firstVector = firstVector; + this.endVector = + count > 0 ? (int) Math.min((long) firstVector + count, numVectors) : numVectors; + + long chunkBytes = (long) Math.max(1, chunkSizeMB) * 1024 * 1024; + int cap = (Integer.MAX_VALUE - 16) / vectorBytes; // keep chunkVectors * vectorBytes in an int + this.chunkVectors = (int) Math.max(1, Math.min(chunkBytes / vectorBytes, cap)); + + // Two reusable direct buffers: the reader fills one while the consumer drains the other. + for (int i = 0; i < 2; i++) { + free.add( + ByteBuffer.allocateDirect(chunkVectors * vectorBytes).order(ByteOrder.LITTLE_ENDIAN)); + } + + this.reader = new Thread(this::readLoop, "fbin-prefetch-reader"); + this.reader.setDaemon(true); + this.reader.start(); + } + + int size() { + return endVector - firstVector; + } + + int dimension() { + return dimension; + } + + /** Reader thread: fill chunks front-to-back, blocking on a free buffer between chunks. */ + private void readLoop() { + long next = firstVector; + try { + while (next < endVector) { + ByteBuffer buf = free.take(); + int toRead = (int) Math.min(chunkVectors, endVector - next); + buf.clear(); + buf.limit(toRead * vectorBytes); + readFully(buf, HEADER_BYTES + next * (long) vectorBytes); + ready.put(new Chunk(buf, next, toRead)); + next += toRead; + } + ready.put(POISON); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); // close() requested; stop quietly + } catch (IOException e) { + readerError = e; + try { + ready.put(POISON); // unblock the consumer so it can observe the error + } catch (InterruptedException ignored) { + Thread.currentThread().interrupt(); + } + } + } + + private void advance() throws IOException { + if (current != null) { + try { + free.put(current.buf); // hand the drained buffer back to the reader + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException("Interrupted returning chunk buffer", e); + } + current = null; + } + Chunk next; + try { + next = ready.take(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException("Interrupted awaiting next chunk", e); + } + if (next == POISON) { + if (readerError != null) { + throw new IOException("Prefetch reader failed", readerError); + } + throw new IOException("No more chunks available (unexpected EOF in prefetch)"); + } + current = next; + } + + /** Fills {@code dst} with the vector at absolute {@code index} (no allocation). */ + void get(int index, float[] dst) throws IOException { + if (index < firstVector || index >= endVector) { + throw new IndexOutOfBoundsException( + "Index " + index + " out of bounds [" + firstVector + ", " + endVector + ")"); + } + while (current == null || index >= current.start + current.len) { + advance(); + } + if (index < current.start) { + throw new UnsupportedOperationException( + "PrefetchingFbinReader requires forward-only sequential access; got index " + + index + + " before current chunk start " + + current.start); + } + int base = (int) (index - current.start) * vectorBytes; + for (int i = 0; i < dimension; i++) { + dst[i] = current.buf.getFloat(base + i * Float.BYTES); + } + } + + /** Convenience allocating variant (e.g. for a one-off query vector). */ + float[] get(int index) throws IOException { + float[] dst = new float[dimension]; + get(index, dst); + return dst; + } + + private void readFully(ByteBuffer buf, long position) throws IOException { + long pos = position; + while (buf.hasRemaining()) { + int n = channel.read(buf, pos); + if (n < 0) { + throw new IOException("Unexpected EOF reading at position " + pos); + } + pos += n; + } + } + + @Override + public void close() throws IOException { + reader.interrupt(); + channel.close(); + } + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java index a5f164b70b..0cdead8ecf 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java @@ -59,13 +59,14 @@ public static enum Strategy { public static final int DEFAULT_MAX_CONN = 32; public static final int DEFAULT_BEAM_WIDTH = 32; public static final CagraGraphBuildAlgo DEFAULT_CAGRA_GRAPH_BUILD_ALGO = - CagraGraphBuildAlgo.NN_DESCENT; + CagraGraphBuildAlgo.AUTO_SELECT; public static final int DEFAULT_NUM_MERGE_WORKERS = 1; public static final Strategy DEFAULT_STRATEGY = Strategy.HEURISTIC; public static final CuvsDistanceType DEFAULT_CUVS_DISTANCE_TYPE = CuvsDistanceType.L2Expanded; public static final int DEFAULT_NN_DESCENT_NUM_ITERATIONS = 20; public static final HnswHeuristicType DEFAULT_HNSW_HEURISTIC_TYPE = HnswHeuristicType.SAME_GRAPH_FOOTPRINT; + public static final int DEFAULT_NUM_INPUT_VECTORS = 0; public static final Supplier DEFAULT_IVF_PQ_PARAMS = () -> { @@ -85,12 +86,14 @@ public static enum Strategy { private final int beamWidth; private final CagraGraphBuildAlgo cagraGraphBuildAlgo; private final CuVSIvfPqParams cuVSIvfPqParams; + private final boolean cuVSIvfPqParamsExplicit; private final int numMergeWorkers; private final ExecutorService mergeExec; private final Strategy strategy; private final CuvsDistanceType cuvsDistanceType; private final int nnDescentNumIterations; private final HnswHeuristicType hnswHeuristicType; + private final int numInputVectors; /** * Constructs an instance of {@link AcceleratedHNSWParams} with specific parameter values. @@ -103,12 +106,16 @@ public static enum Strategy { * @param beamWidth The beam width parameter used when building HNSW index with the fallback mechanism. * @param cagraGraphBuildAlgo The CAGRA graph build algorithm to use [NN_DESCENT, IVF_PQ]. * @param cuVSIvfPqParams An instance of CuVSIvfPqParams containing IVF_PQ specific parameters. + * @param cuVSIvfPqParamsExplicit whether cuVSIvfPqParams was set explicitly by the caller, as + * opposed to defaulted; consulted under HEURISTIC with an explicit IVF_PQ override so a + * caller-supplied value is honored instead of silently replaced by the auto-tuned one. * @param numMergeWorkers The number of merge workers to use with the fallback mechanism. * @param mergeExec The instance of {@link ExecutorService} to use with the fallback mechanism. * @param strategy either HEURISTIC [Default] that delegates the CAGRA build parameters to cuVS (derived from the HNSW-equivalent maxConn and beamWidth) or CUSTOM that uses the parameters passed through this class. * @param cuvsDistanceType the cuvsDistanceType. The default option is L2Expanded. * @param nnDescentNumIterations the number of Iterations to run if building with NN_DESCENT. * @param hnswHeuristicType the heuristic cuVS applies when deriving the CAGRA build parameters from maxConn and beamWidth under the HEURISTIC strategy. + * @param numInputVectors exact number of vectors to be indexed, used to pre-size the native flat buffer (0 = disabled). */ private AcceleratedHNSWParams( int writerThreads, @@ -119,12 +126,14 @@ private AcceleratedHNSWParams( int beamWidth, CagraGraphBuildAlgo cagraGraphBuildAlgo, CuVSIvfPqParams cuVSIvfPqParams, + boolean cuVSIvfPqParamsExplicit, int numMergeWorkers, ExecutorService mergeExec, Strategy strategy, CuvsDistanceType cuvsDistanceType, int nnDescentNumIterations, - HnswHeuristicType hnswHeuristicType) { + HnswHeuristicType hnswHeuristicType, + int numInputVectors) { super(); this.writerThreads = writerThreads; this.intermediateGraphDegree = intermediateGraphDegree; @@ -134,12 +143,14 @@ private AcceleratedHNSWParams( this.beamWidth = beamWidth; this.cagraGraphBuildAlgo = cagraGraphBuildAlgo; this.cuVSIvfPqParams = cuVSIvfPqParams; + this.cuVSIvfPqParamsExplicit = cuVSIvfPqParamsExplicit; this.numMergeWorkers = numMergeWorkers; this.mergeExec = mergeExec; this.strategy = strategy; this.cuvsDistanceType = cuvsDistanceType; this.nnDescentNumIterations = nnDescentNumIterations; this.hnswHeuristicType = hnswHeuristicType; + this.numInputVectors = numInputVectors; } /** @@ -214,6 +225,16 @@ public CuVSIvfPqParams getCuVSIvfPqParams() { return cuVSIvfPqParams; } + /** + * Whether {@link #getCuVSIvfPqParams()} was set explicitly via {@link + * Builder#withCuVSIvfPqParams(CuVSIvfPqParams)}, as opposed to defaulted. + * + * @return true if the caller explicitly set cuVSIvfPqParams + */ + public boolean isCuVSIvfPqParamsExplicit() { + return cuVSIvfPqParamsExplicit; + } + /** * Get the number of merge workers set to be used in the fallback mechanism * @@ -272,6 +293,17 @@ public HnswHeuristicType getHnswHeuristicType() { return hnswHeuristicType; } + /** + * Get the number of input vectors used to pre-size the native flat buffer. A value of + * {@value DEFAULT_NUM_INPUT_VECTORS} means unset (the writer uses the default heap-buffered + * flat path). + * + * @return the number of vectors to be indexed, or 0 if unset + */ + public int getNumInputVectors() { + return numInputVectors; + } + @Override public String toString() { return "AcceleratedHNSWParams [writerThreads=" @@ -302,6 +334,8 @@ public String toString() { + nnDescentNumIterations + ", hnswHeuristicType=" + hnswHeuristicType + + ", numInputVectors=" + + numInputVectors + "]"; } @@ -319,11 +353,13 @@ public static class Builder { private CagraGraphBuildAlgo cagraGraphBuildAlgo = DEFAULT_CAGRA_GRAPH_BUILD_ALGO; private int numMergeWorkers = DEFAULT_NUM_MERGE_WORKERS; private CuVSIvfPqParams cuVSIvfPqParams = null; + private boolean cuVSIvfPqParamsExplicit = false; private ExecutorService mergeExec = null; private Strategy strategy = DEFAULT_STRATEGY; private CuvsDistanceType cuvsDistanceType = DEFAULT_CUVS_DISTANCE_TYPE; private int nnDescentNumIterations = DEFAULT_NN_DESCENT_NUM_ITERATIONS; private HnswHeuristicType hnswHeuristicType = DEFAULT_HNSW_HEURISTIC_TYPE; + private int numInputVectors = DEFAULT_NUM_INPUT_VECTORS; /** * Set the number of cuVS writer threads while building the index @@ -423,6 +459,7 @@ public Builder withCagraGraphBuildAlgo(CagraGraphBuildAlgo cagraGraphBuildAlgo) */ public Builder withCuVSIvfPqParams(CuVSIvfPqParams cuVSIvfPqParams) { this.cuVSIvfPqParams = cuVSIvfPqParams; + this.cuVSIvfPqParamsExplicit = true; return this; } @@ -507,6 +544,23 @@ public Builder withHnswHeuristicType(HnswHeuristicType hnswHeuristicType) { return this; } + /** + * Set the exact number of vectors to be indexed, used to pre-allocate a single contiguous + * native flat buffer (avoiding the on-heap {@code List} and the extra host-matrix + * copy). The native buffer is sized for exactly this many rows, so the value MUST equal the + * number of vectors actually added; the writer fails fast otherwise. Only supported for the + * unsorted single-segment CAGRA_HNSW build (no merges). Not yet supported for the + * binary/scalar quantized writers. A value of {@value DEFAULT_NUM_INPUT_VECTORS} (the + * default) disables it and uses the default heap-buffered flat path. + * + * @param numInputVectors the exact number of vectors to be indexed, or 0 to disable + * @return instance of {@link Builder} + */ + public Builder withNumInputVectors(int numInputVectors) { + this.numInputVectors = numInputVectors; + return this; + } + /** * Validates the input parameters. * @@ -591,6 +645,9 @@ private void validate() throws IllegalArgumentException { + MAX_NN_DESCENT_NUM_ITERATIONS + "]"); } + if (numInputVectors < 0) { + throw new IllegalArgumentException("numInputVectors cannot be negative."); + } } /** @@ -615,12 +672,14 @@ public AcceleratedHNSWParams build() { beamWidth, cagraGraphBuildAlgo, cuVSIvfPqParams, + cuVSIvfPqParamsExplicit, numMergeWorkers, mergeExec, strategy, cuvsDistanceType, nnDescentNumIterations, - hnswHeuristicType); + hnswHeuristicType, + numInputVectors); } } } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index 9c49c07fe0..730940ca61 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -19,8 +19,14 @@ import java.util.Random; import java.util.SortedSet; import java.util.TreeSet; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.store.ByteBuffersDataOutput; +import org.apache.lucene.store.DataOutput; import org.apache.lucene.store.IndexOutput; import org.apache.lucene.util.InfoStream; import org.apache.lucene.util.hnsw.HnswGraph; @@ -69,7 +75,7 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens layerAdjacencies.add(adjacencyMatrix); // Create the single-layer graph - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies, 1); } /** @@ -78,18 +84,25 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens * (its column count). Ceil is used to accommodate odd graph degrees. * Each layer contains 1/M nodes from the previous layer * Creates layers until the highest layer has ≤ M nodes + *

+ * Vectors for higher-layer subsets are read directly from the native matrix + * via {@link CuVSMatrix#getRow(long)} and {@link RowView#toArray(float[])}, + * avoiding any additional heap allocation of the full dataset. Used by both + * the flush and merge paths; the caller provides the vectors as a + * {@link CuVSMatrix}. */ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( FieldInfo fieldInfo, - int size, int dimensions, CuVSMatrix adjacencyListMatrix, - List vectors, + CuVSMatrix vectorDataset, int hnswLayers, CagraIndexParams params, - QuantizationType quantization) + QuantizationType quantization, + int numThreads) throws Throwable { + int size = (int) vectorDataset.size(); int M = Math.ceilDiv((int) adjacencyListMatrix.columns(), 2); // Store all layers data @@ -119,8 +132,7 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( // Select from previous layer nodes int[] prevLayerNodes = layerNodes.get(layerNodes.size() - 1); while (selectedNodesSet.size() < nextLayerSize) { - int idx = random.nextInt(prevLayerNodes.length); - selectedNodesSet.add(prevLayerNodes[idx]); + selectedNodesSet.add(prevLayerNodes[random.nextInt(prevLayerNodes.length)]); } } @@ -131,24 +143,22 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( layerNodes.add(selectedNodes); if (quantization == QuantizationType.NONE) { - // Extract vectors for selected nodes - float[][] selectedVectors = new float[nextLayerSize][]; + // Read only the sampled rows from the native matrix — no full-dataset heap copy + float[][] selectedVectors = new float[nextLayerSize][dimensions]; for (int i = 0; i < nextLayerSize; i++) { - selectedVectors[i] = (float[]) vectors.get(selectedNodes[i]); + vectorDataset.getRow(selectedNodes[i]).toArray(selectedVectors[i]); } // Build CAGRA graph for this layer layerAdjacencies.add( buildCagraGraphForSubset( selectedVectors, selectedNodes, 0, params, dimensions, quantization)); - } else { - - // Extract vectors for selected nodes - int bytesPerVector = (dimensions + 7) / 8; - byte[][] selectedVectors = new byte[nextLayerSize][]; + // Byte width comes from the matrix itself: binary packs 8 dims/byte, scalar is 1 byte/dim. + int bytesPerVector = (int) vectorDataset.columns(); + byte[][] selectedVectors = new byte[nextLayerSize][bytesPerVector]; for (int i = 0; i < nextLayerSize; i++) { - selectedVectors[i] = (byte[]) vectors.get(selectedNodes[i]); + vectorDataset.getRow(selectedNodes[i]).toArray(selectedVectors[i]); } // Build CAGRA graph for this layer @@ -166,7 +176,7 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( } // Create the multi-layer graph with all layers - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies, numThreads); } /** @@ -184,11 +194,9 @@ private static CuVSMatrix buildCagraGraphForSubset( CuVSMatrix subsetDataset; if (quantization == QuantizationType.BINARY) { - subsetDataset = - createByteMatrixFromArray((byte[][]) vectors, bytesPerVector, getCuVSResourcesInstance()); + subsetDataset = createByteMatrixFromArray((byte[][]) vectors, bytesPerVector); } else if (quantization == QuantizationType.SCALAR) { - subsetDataset = - createByteMatrixFromArray((byte[][]) vectors, dimensions, getCuVSResourcesInstance()); + subsetDataset = createByteMatrixFromArray((byte[][]) vectors, dimensions); } else { subsetDataset = CuVSMatrix.ofArray((float[][]) vectors); } @@ -235,56 +243,149 @@ private static CuVSMatrix buildCagraGraphForSubset( * @return a 2D array of offsets * @throws IOException I/O Exceptions */ - public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex) + public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex, int numThreads) throws IOException { // write vectors' neighbors on each level into the vectorIndex file int countOnLevel0 = graph.size(); - int[][] offsets = new int[graph.numLevels()][]; - int[] scratch = new int[graph.maxConn() * 2]; - for (int level = 0; level < graph.numLevels(); level++) { + int numLevels = graph.numLevels(); + int[][] offsets = new int[numLevels][]; + + // Level 0 holds all nodes and dominates serialization cost. Each node's delta/VInt block is + // independent, so encode level 0 in parallel and concatenate the per-thread buffers serially in + // node order, in memory-bounded waves. Higher levels are tiny and stay serial. The on-disk + // bytes + // are identical to the fully-serial path (blocks in node order, offsets = per-node byte + // lengths). + int[] level0Nodes = NodesIterator.getSortedNodes(graph.getNodesOnLevel(0)); + offsets[0] = new int[level0Nodes.length]; + if (numThreads > 1 && level0Nodes.length >= PARALLEL_MIN_NODES) { + writeLevel0Parallel(graph, vectorIndex, level0Nodes, offsets[0], countOnLevel0, numThreads); + } else { + writeLevelSerial(graph, vectorIndex, 0, level0Nodes, offsets[0], countOnLevel0); + } + + for (int level = 1; level < numLevels; level++) { int[] sortedNodes = NodesIterator.getSortedNodes(graph.getNodesOnLevel(level)); offsets[level] = new int[sortedNodes.length]; - int nodeOffsetId = 0; - - for (int node : sortedNodes) { - // Get node neighbors - NeighborArray neighbors = graph.getNeighbors(level, node); - // Get the size of the neighbor array - int size = neighbors.size(); - // Write size in VInt as the neighbors list is typically small - long offsetStart = vectorIndex.getFilePointer(); - // Get neighbors - int[] nnodes = neighbors.nodes(); - // Sort them - Arrays.sort(nnodes, 0, size); - // Now that we have sorted, do delta encoding to minimize the required bits to store the - // information - int actualSize = 0; - if (size > 0) { - scratch[0] = nnodes[0]; - actualSize = 1; - } - // De-duplication - for (int i = 1; i < size; i++) { - assert nnodes[i] < countOnLevel0 : "node too large: " + nnodes[i] + ">=" + countOnLevel0; - // Sorting step helps here - if (nnodes[i - 1] == nnodes[i]) { + writeLevelSerial(graph, vectorIndex, level, sortedNodes, offsets[level], countOnLevel0); + } + // Return offsets (information written while writing the meta info) + return offsets; + } + + /** Node count below which parallel level-0 serialization is not worth the overhead. */ + private static final int PARALLEL_MIN_NODES = 1 << 16; + + /** Nodes per wave — bounds the transient encode buffer regardless of dataset size. */ + private static final int WAVE_NODES = 1 << 20; + + /** Serially encodes a level's nodes into {@code out}, recording per-node byte lengths. */ + private static void writeLevelSerial( + GPUBuiltHnswGraph graph, + IndexOutput out, + int level, + int[] sortedNodes, + int[] offsets, + int countOnLevel0) + throws IOException { + int[] scratch = new int[graph.maxConn() * 2]; + int idx = 0; + for (int node : sortedNodes) { + long start = out.getFilePointer(); + encodeNode(graph.getNeighbors(level, node), scratch, out, countOnLevel0); + offsets[idx++] = Math.toIntExact(out.getFilePointer() - start); + } + } + + /** + * Encodes level 0 in parallel: within memory-bounded waves, threads encode contiguous node + * sub-ranges into per-thread buffers, which are then concatenated to {@code out} in node order + * (identical layout to the serial path). + */ + private static void writeLevel0Parallel( + GPUBuiltHnswGraph graph, + IndexOutput out, + int[] nodes, + int[] offsets, + int countOnLevel0, + int numThreads) + throws IOException { + ExecutorService pool = Executors.newFixedThreadPool(numThreads); + try { + int n = nodes.length; + for (int waveStart = 0; waveStart < n; waveStart += WAVE_NODES) { + int waveEnd = Math.min(waveStart + WAVE_NODES, n); + int perThread = (waveEnd - waveStart + numThreads - 1) / numThreads; + + ByteBuffersDataOutput[] buffers = new ByteBuffersDataOutput[numThreads]; + List> futures = new ArrayList<>(numThreads); + for (int t = 0; t < numThreads; t++) { + final int subStart = waveStart + t * perThread; + final int subEnd = Math.min(subStart + perThread, waveEnd); + final int slot = t; + if (subStart >= subEnd) { continue; } - scratch[actualSize++] = nnodes[i] - nnodes[i - 1]; + futures.add( + pool.submit( + () -> { + ByteBuffersDataOutput buffer = new ByteBuffersDataOutput(); + int[] scratch = new int[graph.maxConn() * 2]; + for (int i = subStart; i < subEnd; i++) { + long before = buffer.size(); + encodeNode(graph.getNeighbors(0, nodes[i]), scratch, buffer, countOnLevel0); + offsets[i] = Math.toIntExact(buffer.size() - before); + } + buffers[slot] = buffer; + return null; + })); } - // Write the size after duplicates are removed - vectorIndex.writeVInt(actualSize); - // Write de-duplicated neighbors - for (int i = 0; i < actualSize; i++) { - vectorIndex.writeVInt(scratch[i]); + for (Future f : futures) { + f.get(); + } + // Concatenate in thread order (== node order), preserving the serial byte layout. + for (ByteBuffersDataOutput buffer : buffers) { + if (buffer != null) { + buffer.copyTo(out); + } } - offsets[level][nodeOffsetId++] = - Math.toIntExact(vectorIndex.getFilePointer() - offsetStart); } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException("Interrupted during parallel writeGraph", e); + } catch (ExecutionException e) { + throw new IOException("Parallel writeGraph failed", e.getCause()); + } finally { + pool.shutdown(); + } + } + + /** + * Sorts, delta-encodes and de-duplicates a node's neighbors and writes the block (VInt size + VInt + * deltas) to {@code out}. Shared by the serial and parallel paths so encoding is identical. + */ + private static void encodeNode( + NeighborArray neighbors, int[] scratch, DataOutput out, int countOnLevel0) + throws IOException { + int size = neighbors == null ? 0 : neighbors.size(); + int actualSize = 0; + if (size > 0) { + int[] nnodes = neighbors.nodes(); + Arrays.sort(nnodes, 0, size); + scratch[0] = nnodes[0]; + actualSize = 1; + for (int i = 1; i < size; i++) { + assert nnodes[i] < countOnLevel0 : "node too large: " + nnodes[i] + ">=" + countOnLevel0; + if (nnodes[i - 1] == nnodes[i]) { + continue; + } + scratch[actualSize++] = nnodes[i] - nnodes[i - 1]; + } + } + out.writeVInt(actualSize); + for (int i = 0; i < actualSize; i++) { + out.writeVInt(scratch[i]); } - // Return offsets (information written while writing the meta info) - return offsets; } /** diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java index 8bb46faf61..c257f0be28 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java @@ -6,6 +6,13 @@ package com.nvidia.cuvs.lucene; import com.nvidia.cuvs.CagraIndexParams; +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; +import com.nvidia.cuvs.CagraIndexParams.CodebookGen; +import com.nvidia.cuvs.CagraIndexParams.CudaDataType; +import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; +import com.nvidia.cuvs.CuVSIvfPqIndexParams; +import com.nvidia.cuvs.CuVSIvfPqParams; +import com.nvidia.cuvs.CuVSIvfPqSearchParams; /** * A centralized place for producing {@link CagraIndexParams} from the cuvs-lucene input parameter @@ -18,6 +25,121 @@ public class CagraIndexParamsFactory { private CagraIndexParamsFactory() {} + /** + * Translation of the internal logic found in the constructor of {@code struct ivf_pq_params} in + * cuVS's {@code cpp/include/cuvs/neighbors/ivf_pq.hpp}. + * + * Ideally we should hook into the internal API but this is currently replicated to avoid complications + * in other parts of code base. + * + * TODO: cuvs-java has no standalone binding for that C++ constructor -- only fromHnswParams and + * fromDataset, which bundle algorithm selection together with parameter tuning and offer no way + * to pin the algorithm while still getting cuVS's own auto-tuned IVF-PQ parameters. If cuvs-java + * exposed the ivf_pq_params(dataset_extents, metric) constructor directly (or an equivalent + * narrower entry point that derives IVF-PQ parameters without also choosing the algorithm), this + * method could delegate to it instead of reimplementing the heuristic here. + */ + private static CuVSIvfPqParams getCuVSIvfPqParams(long rows, long dimension) { + int pqDim; + int pqBits; + if (dimension <= 32) { + pqDim = 16; + pqBits = 8; + } else { + pqBits = 4; + if (dimension <= 64) { + pqDim = 32; + } else if (dimension <= 128) { + pqDim = 64; + } else if (dimension <= 192) { + pqDim = 96; + } else { + pqDim = (int) roundUpSafe(dimension / 2, 128); + } + } + int nLists = (int) Math.max(1, rows / 2000); + final int kmeansNIters = 10; + final double kMinPointsPerCluster = 32; + double minKmeansTrainsetPoints = kMinPointsPerCluster * nLists; + final double maxKmeansTrainsetFraction = 1.0; + double minKmeansTrainsetFraction = + Math.min(maxKmeansTrainsetFraction, minKmeansTrainsetPoints / rows); + double kmeansTrainsetFraction = + Math.clamp( + 1.0 / Math.sqrt(rows * 1e-5), minKmeansTrainsetFraction, maxKmeansTrainsetFraction); + int nProbes = (int) Math.round(Math.sqrt(nLists) / 20 + 4); + CuVSIvfPqIndexParams cuVSIvfPqIndexParams = + new CuVSIvfPqIndexParams.Builder() + .withCodebookKind(CodebookGen.PER_SUBSPACE) + .withKmeansNIters(kmeansNIters) + .withKmeansTrainsetFraction(kmeansTrainsetFraction) + .withNLists(nLists) + .withPqBits(pqBits) + .withPqDim(pqDim) + .withAddDataOnBuild(true) + .withConservativeMemoryAllocation(true) + .build(); + CuVSIvfPqSearchParams cuVSIvfPqSearchParams = + new CuVSIvfPqSearchParams.Builder() + .withLutDtype(CudaDataType.CUDA_R_16F) + .withInternalDistanceDtype(CudaDataType.CUDA_R_16F) + .withNProbes(nProbes) + .build(); + return new CuVSIvfPqParams.Builder() + .withCuVSIvfPqIndexParams(cuVSIvfPqIndexParams) + .withCuVSIvfPqSearchParams(cuVSIvfPqSearchParams) + .withRefinementRate(1) + .build(); + } + + private static long roundUpSafe(long numberToRound, long modulus) { + long remainder = numberToRound % modulus; + if (remainder == 0) { + return numberToRound; + } + return numberToRound - remainder + modulus; + } + + private static CagraIndexParams getNNDescentParams( + int graphDegree, + int intGraphDegree, + int writerThreads, + long nnDescentNumIterations, + CuvsDistanceType cuvsDistanceType) { + return new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(graphDegree) + .withIntermediateGraphDegree(intGraphDegree) + .withNNDescentNumIterations(nnDescentNumIterations) + .withNumWriterThreads(writerThreads) + .withMetric(cuvsDistanceType) + .build(); + } + + /** + * @param explicitIvfPqParams the caller's own IVF-PQ params if they set one ({@link + * AcceleratedHNSWParams#isCuVSIvfPqParamsExplicit()}), honored as-is; otherwise {@code + * null}, in which case params are auto-tuned from {@code rows}/{@code dimension}. + */ + private static CagraIndexParams getIVFPQParams( + int graphDegree, + int intGraphDegree, + int writerThreads, + long rows, + long dimension, + CuvsDistanceType cuvsDistanceType, + CuVSIvfPqParams explicitIvfPqParams) { + return new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.IVF_PQ) + .withCuVSIvfPqParams( + explicitIvfPqParams != null ? explicitIvfPqParams : getCuVSIvfPqParams(rows, dimension)) + .withNumWriterThreads(writerThreads) + .withIntermediateGraphDegree(intGraphDegree) + .withGraphDegree(graphDegree) + .withMetric(cuvsDistanceType) + .build(); + } + /** * Creates an instance of {@link CagraIndexParams} for the GPU-native CAGRA index based on the * chosen strategy in the {@link GPUSearchParams}. @@ -74,29 +196,53 @@ public static CagraIndexParams create( public static CagraIndexParams create( AcceleratedHNSWParams acceleratedHNSWParams, long rows, long dimension) { if (acceleratedHNSWParams.getStrategy().equals(AcceleratedHNSWParams.Strategy.HEURISTIC)) { - // Delegate the derivation of the graph degrees, build algorithm and its parameters to cuVS, - // expressed in terms of the HNSW-equivalent maxConn/beamWidth. - CagraIndexParams derived = - CagraIndexParams.fromHnswParams( - rows, - dimension, - acceleratedHNSWParams.getMaxConn(), - acceleratedHNSWParams.getBeamWidth(), - acceleratedHNSWParams.getHnswHeuristicType(), - acceleratedHNSWParams.getCuvsDistanceType()); - // TODO: fromHnswParams has no writerThreads argument, so its result carries the cuVS default - // (not a heuristic value). We can rebuild the CagraIndexParams with the caller-supplied - // writerThreads for now but should fix this in cuVS in the future. - return new CagraIndexParams.Builder() - .withGraphDegree(derived.getGraphDegree()) - .withIntermediateGraphDegree(derived.getIntermediateGraphDegree()) - .withCagraGraphBuildAlgo(derived.getCagraGraphBuildAlgo()) - .withCuVSIvfPqParams(derived.getCuVSIvfPqParams()) - .withNNDescentNumIterations(derived.getNNDescentNumIterations()) - .withMetric(acceleratedHNSWParams.getCuvsDistanceType()) - .withNumWriterThreads(acceleratedHNSWParams.getWriterThreads()) - .build(); + // An explicit cagraGraphBuildAlgo of IVF_PQ or NN_DESCENT overrides the heuristic choice; + // AUTO_SELECT (the default) delegates all parameter derivation to cuVS via fromHnswParams. + CagraGraphBuildAlgo algo = acceleratedHNSWParams.getCagraGraphBuildAlgo(); + if (algo == CagraGraphBuildAlgo.IVF_PQ) { + return getIVFPQParams( + acceleratedHNSWParams.getGraphdegree(), + acceleratedHNSWParams.getIntermediateGraphDegree(), + acceleratedHNSWParams.getWriterThreads(), + rows, + dimension, + acceleratedHNSWParams.getCuvsDistanceType(), + acceleratedHNSWParams.isCuVSIvfPqParamsExplicit() + ? acceleratedHNSWParams.getCuVSIvfPqParams() + : null); + } else if (algo == CagraGraphBuildAlgo.NN_DESCENT) { + return getNNDescentParams( + acceleratedHNSWParams.getGraphdegree(), + acceleratedHNSWParams.getIntermediateGraphDegree(), + acceleratedHNSWParams.getWriterThreads(), + acceleratedHNSWParams.getNNDescentNumIterations(), + acceleratedHNSWParams.getCuvsDistanceType()); + } else { + // AUTO_SELECT: delegate the derivation of the graph degrees, build algorithm and its + // parameters to cuVS, expressed in terms of the HNSW-equivalent maxConn/beamWidth. + CagraIndexParams derived = + CagraIndexParams.fromHnswParams( + rows, + dimension, + acceleratedHNSWParams.getMaxConn(), + acceleratedHNSWParams.getBeamWidth(), + acceleratedHNSWParams.getHnswHeuristicType(), + acceleratedHNSWParams.getCuvsDistanceType()); + // TODO: fromHnswParams has no writerThreads argument, so its result carries the cuVS + // default (not a heuristic value). We can rebuild the CagraIndexParams with the + // caller-supplied writerThreads for now but should fix this in cuVS in the future. + return new CagraIndexParams.Builder() + .withGraphDegree(derived.getGraphDegree()) + .withIntermediateGraphDegree(derived.getIntermediateGraphDegree()) + .withCagraGraphBuildAlgo(derived.getCagraGraphBuildAlgo()) + .withCuVSIvfPqParams(derived.getCuVSIvfPqParams()) + .withNNDescentNumIterations(derived.getNNDescentNumIterations()) + .withMetric(acceleratedHNSWParams.getCuvsDistanceType()) + .withNumWriterThreads(acceleratedHNSWParams.getWriterThreads()) + .build(); + } } + // CUSTOM: forward the caller's algorithm and the parameters it consumes. return new CagraIndexParams.Builder() .withGraphDegree(acceleratedHNSWParams.getGraphdegree()) .withIntermediateGraphDegree(acceleratedHNSWParams.getIntermediateGraphDegree()) diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java index 97fedf924f..ce77f36dea 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java @@ -201,8 +201,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro var cagraIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); try { CuVSMatrix cagraDataset = - Utils.createFloatMatrix( - vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance()); + Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); writeCagraIndex(cagraIndexOutputStream, cagraDataset); } catch (Throwable t) { // Fallback to brute force in a few cases, for now. @@ -215,8 +214,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro if (indexType.isBruteForce()) { var bruteForceIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); CuVSMatrix bruteforceDataset = - Utils.createFloatMatrix( - vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance()); + Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); writeBruteForceIndex(bruteForceIndexOutputStream, bruteforceDataset); bruteForceIndexLength = cuvsIndex.getFilePointer() - bruteForceIndexOffset; diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java index 8e2d9a70e2..92df78224b 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java @@ -8,6 +8,8 @@ import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.quantizeFloatVectorsToBinary; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.quantizeFloatVectorsToScalar; +import com.nvidia.cuvs.CuVSHostMatrix; +import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.QuantizationType; import java.io.IOException; import java.util.List; @@ -23,18 +25,64 @@ public class FieldWriter extends KnnFieldVectorsWriter { RamUsageEstimator.shallowSizeOfInstance(FieldWriter.class); private final FieldInfo fieldInfo; + private final int dimension; private final FlatFieldVectorsWriter flatFieldVectorsWriter; private int lastDocID = -1; private QuantizationType quantizationType; + /** + * Native-buffering state. When {@code numInputVectors > 0} the writer streams incoming vectors + * directly into a native host matrix ({@link CuVSMatrix#hostBuilder}) instead of accumulating a + * {@code List} on the Java heap via {@link #flatFieldVectorsWriter}. That matrix is reused + * as the CAGRA build input, so the full dataset is never held twice (heap list + native copy). + * + *

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

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

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

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

Ported code — version-pinned. The dense float32 layout (format constants, header, meta + * field order, and footer) is transcribed from Lucene 10.2.0, which must stay equal to the + * {@code lucene-core} version in {@code pom.xml}. This is enforced by + * {@code TestNativeFlatVectorsWriterFormatConstants}, which fails the build if the resolved + * {@code lucene-core} classpath version drifts from the pin below — update + * {@code PINNED_LUCENE_VERSION} there together with this class on a verified upgrade. Sources (tag + * {@code releases/lucene/10.2.0}): + * + *

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

On a Lucene upgrade: re-verify this class against the new version's two files above. The + * constants are package-private in Lucene (hence re-declared here), and the file format is validated + * on read via {@code CodecUtil.checkIndexHeader} (codec name + version range) plus the fixed meta + * layout — so if the new version bumps {@code VERSION_CURRENT}, renames a codec, or changes the meta + * field order, files written here will be silently incompatible and the stock + * {@code Lucene99FlatVectorsReader} will reject or misread them. Concretely: + * + *

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

Copies vectors directly into a native host matrix via {@link CuVSMatrix#hostBuilder}, + * without creating an intermediate {@code float[][]} on the heap. * * @param data The float vectors - * @param dimensions The number float elements in each vector - * @param resources The CuVS resources for device matrix creation - * @return an instance of CuVSMatrix + * @param dimensions The number of float elements in each vector + * @return a host-memory CuVSMatrix */ - static CuVSMatrix createFloatMatrix(List data, int dimensions, CuVSResources resources) { - // Use Builder pattern to avoid intermediate float[][] allocation - // and copy directly from List to device memory + static CuVSMatrix createFloatMatrix(List data, int dimensions) { CuVSMatrix.Builder builder = - CuVSMatrix.deviceBuilder( - resources, - data.size(), // rows (number of vectors) - dimensions, // columns (vector dimension) - CuVSMatrix.DataType.FLOAT); - - // Add vectors one by one - builder copies directly to device memory + CuVSMatrix.hostBuilder(data.size(), dimensions, CuVSMatrix.DataType.FLOAT); for (float[] vector : data) { builder.addVector(vector); } - return builder.build(); } /** - * A method to build a CuVSMatrix from a list of byte vectors (for binary quantized vectors). - * - * Uses CuVSMatrix.Builder to copy vectors directly to device memory - * without creating intermediate heap arrays. + * Builds a host-memory CuVSMatrix from a list of byte vectors (e.g. quantized vectors). * * @param data The byte vectors (packed bits for binary quantization) * @param bytesPerVector The number of bytes in each vector - * @param resources The CuVS resources for device matrix creation - * @return an instance of CuVSMatrix with BYTE data type + * @return a host-memory CuVSMatrix with BYTE data type */ - static CuVSMatrix createByteMatrix( - List data, int bytesPerVector, CuVSResources resources) { - // Use Builder pattern to avoid intermediate byte[][] allocation - // and copy directly from List to device memory + static CuVSMatrix createByteMatrix(List data, int bytesPerVector) { CuVSMatrix.Builder builder = - CuVSMatrix.deviceBuilder( - resources, - data.size(), // rows (number of vectors) - bytesPerVector, // columns (bytes per vector) - CuVSMatrix.DataType.BYTE); - - // Add vectors one by one - builder copies directly to device memory + CuVSMatrix.hostBuilder(data.size(), bytesPerVector, CuVSMatrix.DataType.BYTE); for (byte[] vector : data) { builder.addVector(vector); } - return builder.build(); } /** - * A method to build a CuVSMatrix from a 2D byte array (for binary quantized vectors). + * Builds a host-memory CuVSMatrix from a 2D byte array (e.g. quantized vectors). * * @param data The 2D byte array (packed bits for binary quantization) * @param bytesPerVector The number of bytes in each vector - * @param resources The CuVS resources for device matrix creation - * @return an instance of CuVSMatrix with BYTE data type + * @return a host-memory CuVSMatrix with BYTE data type */ - static CuVSMatrix createByteMatrixFromArray( - byte[][] data, int bytesPerVector, CuVSResources resources) { + static CuVSMatrix createByteMatrixFromArray(byte[][] data, int bytesPerVector) { CuVSMatrix.Builder builder = - CuVSMatrix.deviceBuilder( - resources, - data.length, // rows (number of vectors) - bytesPerVector, // columns (bytes per vector) - CuVSMatrix.DataType.BYTE); - - // Add vectors one by one - builder copies directly to device memory + CuVSMatrix.hostBuilder(data.length, bytesPerVector, CuVSMatrix.DataType.BYTE); for (byte[] vector : data) { builder.addVector(vector); } diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMergedGraphOrdinalBounds.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMergedGraphOrdinalBounds.java new file mode 100644 index 0000000000..12b00c6e12 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMergedGraphOrdinalBounds.java @@ -0,0 +1,198 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; + +import java.util.Random; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.hnsw.HnswGraphProvider; +import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.CodecReader; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.NoMergePolicy; +import org.apache.lucene.index.Term; +import org.apache.lucene.index.TieredMergePolicy; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.tests.util.TestUtil; +import org.apache.lucene.util.hnsw.HnswGraph; +import org.junit.Test; + +/** + * Repro for the CI-observed {@code EOFException} in {@code OffHeapFloatVectorValues} during + * concurrent KNN search over an accelerated-HNSW index built via {@link + * TestAcceleratedHNSWDeletedDocuments}/{@link TestCuVSAcceleratedHNSWDeletedDocuments} (both: + * deletions + a real merge, heap-buffered path, no {@code numInputVectors}). + * + *

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

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

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

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

{@link NativeFlatVectorsWriter}'s class javadoc says its dense float32 layout (format + * constants, header, meta field order, footer) is transcribed from Lucene 10.2.0 and "must stay + * equal to the lucene-core version in pom.xml." This asserts that invariant against the resolved + * {@code lucene-core} jar actually on the classpath (its manifest {@code Specification-Version}, + * not a parse of {@code pom.xml}'s text), so a Lucene upgrade fails this test immediately rather + * than silently writing {@code .vec}/{@code .vemf} files the stock + * {@code Lucene99FlatVectorsReader} may no longer read correctly. + * + *

This intentionally fires on any version change, not just ones that actually alter + * the format — per the class javadoc, every upgrade needs a manual re-verification pass against + * the new version's {@code Lucene99FlatVectorsFormat}/{@code Lucene99FlatVectorsWriter} sources. + */ +public class TestNativeFlatVectorsWriterFormatConstants { + + private static final String PINNED_LUCENE_VERSION = "10.2.0"; + + @Test + public void lucenePinMatchesResolvedClasspathVersion() { + String resolved = Lucene99FlatVectorsFormat.class.getPackage().getSpecificationVersion(); + assertEquals( + "lucene-core on the classpath is " + + resolved + + ", but NativeFlatVectorsWriter is pinned to " + + PINNED_LUCENE_VERSION + + " (per its class javadoc). Re-verify NativeFlatVectorsWriter against the new" + + " version's Lucene99FlatVectorsFormat/Lucene99FlatVectorsWriter sources, update the" + + " mirrored constants and write sequence if needed, then bump PINNED_LUCENE_VERSION" + + " here.", + PINNED_LUCENE_VERSION, + resolved); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatVectorsWriterRoundTrip.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatVectorsWriterRoundTrip.java new file mode 100644 index 0000000000..de4bdf6659 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatVectorsWriterRoundTrip.java @@ -0,0 +1,135 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import java.io.File; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Random; +import java.util.UUID; +import org.apache.commons.io.FileUtils; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.NoMergePolicy; +import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +/** + * Positive round-trip check for {@link NativeFlatVectorsWriter}: builds a single-segment index + * with {@code AcceleratedHNSWParams.numInputVectors} set (native flat buffering) and reads the + * {@code .vec}/{@code .vemf} files back through the stock {@code Lucene99FlatVectorsReader}, + * asserting every vector round-trips byte-exact. + * + *

This is the check called for by {@link NativeFlatVectorsWriter}'s "on a Lucene upgrade" + * class javadoc: {@link TestNativeFlatVectorsWriterFormatConstants} catches a version-pin drift, + * but only this test actually confirms the hand-transcribed format is still readable by Lucene's + * real reader. + */ +@SuppressSysoutChecks(bugUrl = "") +public class TestNativeFlatVectorsWriterRoundTrip extends LuceneTestCase { + + private static final String ID_FIELD = "id"; + private static final String VECTOR_FIELD = "vector_field"; + + private Random random; + private Path indexDirPath; + + @Before + public void beforeTest() throws Exception { + assumeTrue("cuVS not supported", isSupported()); + random = new Random(222); + indexDirPath = Paths.get(UUID.randomUUID().toString()); + } + + @After + public void afterTest() throws Exception { + if (indexDirPath == null) { + return; + } + File indexDirPathFile = indexDirPath.toFile(); + if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { + FileUtils.deleteDirectory(indexDirPathFile); + } + } + + @Test + public void vectorsRoundTripThroughStockLucene99FlatVectorsReader() throws Exception { + int numDocs = 500; + int dimension = 32; + float[][] dataset = generateDataset(random, numDocs, dimension); + + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder().withNumInputVectors(numDocs).build(); + Codec codec = new Lucene101AcceleratedHNSWCodec(params); + + // Force everything into a single unsorted, unmerged flush: native flat buffering requires + // numInputVectors to equal the exact number of vectors landing in that one flush. + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(codec) + .setUseCompoundFile(false) + .setMaxBufferedDocs(numDocs + 1) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) + .setMergePolicy(NoMergePolicy.INSTANCE); + + try (Directory dir = FSDirectory.open(indexDirPath); + IndexWriter writer = new IndexWriter(dir, config)) { + for (int i = 0; i < numDocs; i++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + writer.addDocument(document); + } + writer.commit(); + } + + try (Directory dir = FSDirectory.open(indexDirPath); + DirectoryReader reader = DirectoryReader.open(dir)) { + assertEquals("expected a single native-flat-buffered segment", 1, reader.leaves().size()); + + LeafReader leafReader = reader.leaves().get(0).reader(); + FloatVectorValues values = leafReader.getFloatVectorValues(VECTOR_FIELD); + assertNotNull(values); + assertEquals(numDocs, values.size()); + assertEquals(dimension, values.dimension()); + + int seen = 0; + KnnVectorValues.DocIndexIterator it = values.iterator(); + for (int doc = it.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = it.nextDoc()) { + String id = leafReader.storedFields().document(doc).get(ID_FIELD); + float[] roundTripped = values.vectorValue(it.index()); + assertArrayEquals( + "vector for id=" + + id + + " did not round-trip byte-exact through the stock" + + " Lucene99FlatVectorsReader", + dataset[Integer.parseInt(id)], + roundTripped, + 0f); + seen++; + } + assertEquals("did not visit every vector", numDocs, seen); + } + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphEquivalence.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphEquivalence.java new file mode 100644 index 0000000000..fd4f9079f4 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphEquivalence.java @@ -0,0 +1,154 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; + +import com.nvidia.cuvs.CuVSMatrix; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Random; +import org.apache.lucene.store.ByteBuffersDirectory; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.util.hnsw.HnswGraph; +import org.apache.lucene.util.hnsw.HnswGraph.NodesIterator; +import org.junit.Test; + +/** + * Verifies that {@code writerThreads > 1} produces the same result as the serial path for the + * two parallelizations gated on {@code AcceleratedHNSWUtils}/{@code GPUBuiltHnswGraph}'s {@code + * PARALLEL_MIN_NODES} threshold: materializing the CAGRA adjacency into {@code NeighborArray}s + * (the {@code GPUBuiltHnswGraph} constructor), and encoding level 0 to disk ({@code + * AcceleratedHNSWUtils#writeGraph}). + * + *

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