diff --git a/README.md b/README.md index f37f0a878..a16da866a 100644 --- a/README.md +++ b/README.md @@ -46,6 +46,8 @@ Introductory tutorials for JVector are available in [docs/tutorials](./docs/tuto For file-backed input datasets, see [Batched input vector access](./docs/batched-vector-access.md), including runnable 100k-vector examples and optional integration with PQ and NVQ. +For writing completed graphs with FP or NVQ features, see [Deferred batch serialization](./docs/deferred-batch-writer.md). It works with application vector sources and preloaded data; dataset prefetching is optional. + The older step-by-step guide for JV can be found [here](./docs/legacy/jvector-step-by-step.md). New users should start with the tutorials mentioned earler, but the step-by-step guide contains useful commentary for advanced users. diff --git a/benchmarks-jmh/README.md b/benchmarks-jmh/README.md index 1da3d8d9f..88a0e1c07 100644 --- a/benchmarks-jmh/README.md +++ b/benchmarks-jmh/README.md @@ -133,7 +133,7 @@ The formatter will output a clean table showing: - **Visited_Count**: Average number of nodes visited during search - **Expanded_Count_BaseLayer**: Average number of nodes expanded in base layer -Example output: +Example output:[^machine] ``` k PQ_Subspaces Time_ms Recall ReRanked_Count Visited_Count Expanded_Count_BaseLayer 50 0 19.283 1.000 0.0 3290.8 253.7 @@ -143,3 +143,5 @@ Example output: + +[^machine]: Machine type and core count were not recorded for this formatting example; its timings are illustrative, not a hardware baseline. diff --git a/docs/benchmarking.md b/docs/benchmarking.md index 79ddba82e..5809b9152 100644 --- a/docs/benchmarking.md +++ b/docs/benchmarking.md @@ -115,3 +115,11 @@ To add a custom fvecs/ivecs dataset: ``` For remote datasets, use `base_url` to specify where files should be downloaded from. The `${VAR}` and `${VAR:-default}` syntax is supported for environment variable expansion. See the example config for details. + +## Final index serialization + +BenchYAML and AutoBenchYAML default to parallel complete-record serialization +(default batch size 64) after graph cleanup. `-Djvector.parallel_writer=false` +selects the sequential path, which encodes FP/NVQ features during insertion. +Dataset prefetching is a separate option and is not required. See +[Deferred batch serialization](deferred-batch-writer.md) for activation and caller examples. diff --git a/docs/deferred-batch-writer.md b/docs/deferred-batch-writer.md new file mode 100644 index 000000000..b36015074 --- /dev/null +++ b/docs/deferred-batch-writer.md @@ -0,0 +1,134 @@ +# Deferred batch serialization + +The parallel graph writer can serialize complete layer-zero records into bounded +worker buffers. Each record contains its ordinal, inline features, neighbor count, +and adjacency list. By default, a worker writes up to 64 complete records together. +The final partial batch is written without adding records or padding to the format. +This preserves the V6 layout and existing readers. + +## Production sources + +The library writer does not depend on Grid, BenchYAML, dataset preloading, or +#744's buffered-access API. It accepts the existing +`Map>`: callers supply feature states for the +requested original ordinals from their own source. That source can be resident +vectors, a database's flush/compaction source, or another source; it must support +concurrent supplier calls and remain available through final serialization. +The writer does not require the original dataset to be fully resident. + +Complete-record batching also applies when all input vectors are preloaded. Its +benefits come from deferring per-insertion feature writes and grouping final writes, +independently of how input vectors are loaded. Buffered range access is an optional +Grid integration, not a prerequisite for the core writer. + +## Caller migration: early writes versus deferred batches + +Use the parallel writer when deferring NVQ/FP encoding to final serialization. +A sequential final writer performs that encoding on one thread and can regress +construction time. Sequential callers should retain encoding during insertion +and `writeFeaturesInline`. + +Both examples assume the same trained `nvq`, graph `builder`, node `count`, +`writerWorkers`, and output path. The builder's construction scoring is unchanged; +its peer vectors are supplied separately. `vectorForOrdinal` is the application's +lookup of an original input vector, not a harness loader. + +### Old: write each inline feature during insertion + +```java +try (var writer = new OnDiskParallelGraphIndexWriter.Builder(builder.getGraph(), outputPath) + .withParallelWorkerThreads(writerWorkers) + .withMapper(new OrdinalMapper.IdentityMapper(count - 1)) + .with(new NVQ(nvq)) + .build()) { + for (int ordinal = 0; ordinal < count; ordinal++) { + var vector = vectorForOrdinal.apply(ordinal); + builder.addGraphNode(ordinal, vector); + writer.writeFeaturesInline(ordinal, Feature.singleState( + FeatureId.NVQ_VECTORS, new NVQ.State(nvq.encode(vector)))); + } + builder.cleanup(); + writer.write(Map.of()); // Preserve the earlier feature writes; add adjacency. +} +``` + +### New: supply the features when writing the completed graph + +```java +for (int ordinal = 0; ordinal < count; ordinal++) { + builder.addGraphNode(ordinal, vectorForOrdinal.apply(ordinal)); +} +builder.cleanup(); + +try (var writer = new OnDiskParallelGraphIndexWriter.Builder(builder.getGraph(), outputPath) + .withParallelWorkerThreads(writerWorkers) + .withMapper(new OrdinalMapper.IdentityMapper(count - 1)) + .with(new NVQ(nvq)) + .build()) { + writer.write(Feature.singleStateFactory(FeatureId.NVQ_VECTORS, + ordinal -> new NVQ.State(nvq.encode(vectorForOrdinal.apply(ordinal))))); +} +``` + +The second call supplies NVQ payloads during final serialization, together with +adjacency, in batches of up to 64 complete records. No `writeFeaturesInline` pass +is needed. The same pattern works for FP using `InlineVectors` and +`InlineVectors.State`, or another supported inline feature. Supply every configured +inline feature to obtain complete-record batching, including fused PQ if configured. + +Final suppliers run concurrently. `vectorForOrdinal` must remain available through +`write` and provide independent borrowed vectors per worker, or stable vectors. +For an existing `RandomAccessVectorValues`, `threadLocalSupplier()` supplies those +independent views; it does not require buffered hooks. A caller can instead return +already encoded feature states from its own storage, without retaining original FP +vectors for this write. + +Integrations that already call `write` with every inline feature supplier receive +batching when using the parallel writer. Integrations that write features early +must move those calls to final suppliers to use the deferred path. The incremental +API remains compatible: a one-pass source whose values cannot be obtained again +must retain encoded states for final writing or continue using incremental writes. + +## Batch size and memory + +The default is 64 records. To override it, call `withParallelBatchSize(records)` +before `with(...)` on either `OnDiskParallelGraphIndexWriter.Builder` or +`RandomAccessOnDiskGraphIndexWriter.Builder`. The value must be positive. Each active +worker holds one buffer of `recordSize * records` bytes; feature encoders and vector +sources have their own memory costs. No dataset-sized output buffer is allocated. + +The unified `RandomAccessOnDiskGraphIndexWriter.Builder` retains its existing +parallel/sequential selection. Setting the batch size does not enable parallel +writing, and is ignored by the sequential writer. A caller-supplied executor +remains caller-owned and must outlive the write. + +## Dataset prefetching in Grid + +When the parallel writer is selected, Grid defers its inline NVQ or FP feature serialization +until after cleanup. Its final serialization suppliers consume worker-local ranges through +`VectorAccess.openRange`. With the sequential writer, which calls suppliers on one +thread, Grid keeps encoding these features on the insertion threads and writing them inline. +Sources implementing the buffered-access capability use that implementation; +ordinary sources retain the ordinary-access fallback. Borrowed vectors are consumed +before advancing a cursor, and ranges are released after the workers finish. + +This integrates with #744's optional `useDatasetPrefetching` setting. It does not +change construction scoring, quantization parameters, or the query path. Production +integrations can opt into deferred serialization by supplying their features at the +final write; their existing vector-source and incremental-write APIs remain valid. + +## Running the benchmark writer paths + +BenchYAML and AutoBenchYAML select the parallel writer by default. To use the +sequential writer instead, add `-Djvector.parallel_writer=false` to the Java VM +options. In that mode, Grid encodes FP/NVQ features during insertion and preserves +the earlier inline-write path. No YAML setting is needed. + +Grid sizes its writer pool from `PhysicalCoreExecutor`; no machine-specific worker +count is embedded in this change. Dataset prefetching is a separate optional setting +in the index YAML. An ordinary preloaded dataset can use the same deferred writer. + +This default applies only to the benchmark entry points. The library's unified +writer builder retains its existing sequential default, and application callers +still explicitly choose parallel writing. Code calling Grid directly can select +its writer through `GraphIndexBuilderConfig.getInstance().setParallelBuild(...)`. diff --git a/docs/legacy/jvector-step-by-step.md b/docs/legacy/jvector-step-by-step.md index 176ff38f8..b7424ac2e 100644 --- a/docs/legacy/jvector-step-by-step.md +++ b/docs/legacy/jvector-step-by-step.md @@ -1,5 +1,9 @@ ## JVector step-by-step +> Historical guide: API names and early-write examples below reflect an older +> version. Use the [current tutorials](../tutorials/2-disk-tutorial.md) and +> [deferred serialization guide](../deferred-batch-writer.md) for current construction APIs. + All code samples are from [SiftSmall](https://github.com/jbellis/jvector/blob/main/jvector-examples/src/main/java/io/github/jbellis/jvector/example/SiftSmall.java) in the JVector source repo, which includes the siftsmall dataset as well. Just import the project in your IDE and click Run to try it out! #### Step 1: Build and query an index in memory diff --git a/docs/release notes/4.0.2/747.performance.md b/docs/release notes/4.0.2/747.performance.md new file mode 100644 index 000000000..3c36111ee --- /dev/null +++ b/docs/release notes/4.0.2/747.performance.md @@ -0,0 +1,36 @@ +### Deferred batch graph serialization + +Parallel graph serialization can now write complete layer-zero records in bounded +batches of 64, with an optional positive batch-size override. When the parallel writer +is selected, Grid supplies inline NVQ or full-precision features during final +serialization, after graph cleanup, and uses buffered vector ranges when the input +source supports them. With the sequential writer selected, Grid keeps encoding those +features during insertion. + +The core writer accepts application-supplied feature states independently of the +benchmark harness or dataset prefetching. Complete-record batching applies equally +to fully preloaded inputs; the source does not have to preload the whole dataset. + +Integrators should use the parallel writer when deferring NVQ/FP encoding to +final serialization. A sequential final writer performs that encoding on one +thread and can regress construction time. Sequential callers should retain +encoding during insertion and `writeFeaturesInline`. + +The existing V6 layout, readers, ordinary vector sources, and incremental +`writeFeaturesInline` API remain supported. Prewritten features are preserved. +Concurrent feature suppliers must keep their input source available until writing +finishes. BenchYAML and AutoBenchYAML default to parallel writing, with +`-Djvector.parallel_writer=false` selecting the sequential path. The library +writer default remains unchanged. + +Recorded construction studies with dataset prefetching reduced ample-memory final +write time from 136.10 to 50.69 seconds on MPNet 10M, 148.50 to 64.02 seconds on +Cohere 10M, and 162.64 to 51.76 seconds on CAP 6M.[^machine] These are separate fresh runs, +not a universal throughput or memory guarantee. See +[Deferred batch serialization](../../deferred-batch-writer.md) for usage. + +Construction tutorials and examples now demonstrate final feature suppliers for +deferred FP/NVQ writing, with per-worker source views and an explicit source-lifetime +requirement. The query algorithms and quantization settings are unchanged. + +[^machine]: GCE VM, 48-core Intel Xeon 6985P-C (96 threads), local SSD RAID10, JDK 23, 128 GiB heap; native Highway, 48 workers. diff --git a/docs/tutorials/2-disk-tutorial.md b/docs/tutorials/2-disk-tutorial.md index 4040ff7de..e9f87bb87 100644 --- a/docs/tutorials/2-disk-tutorial.md +++ b/docs/tutorials/2-disk-tutorial.md @@ -8,7 +8,7 @@ Unlike the previous example, this time we'll use a proper ANN dataset based on O // This is a preconfigured dataset that will be downloaded automatically. DataSet dataset = DataSets.loadDataSet("ada002-100k").orElseThrow(() -> new RuntimeException("Dataset doesn't exist or wasn't configured correctly") -); +).getDataSet(); ``` > [!TIP] @@ -45,7 +45,8 @@ Path graphPath = Files.createTempFile("jvector-example-graph", null); // or whe // Create a writer for the on-heap graph we just built. // Remember to close when done. -GraphIndexWriter writer = GraphIndexWriter.getBuilderFor(GraphIndexWriterTypes.RANDOM_ACCESS_PARALLEL, heapGraph, graphPath) +GraphIndexWriter writer = new OnDiskParallelGraphIndexWriter.Builder(heapGraph, graphPath) + .withParallelWorkerThreads(PhysicalCoreExecutor.getPhysicalCoreCount()) // Let the writer know that we'll also be passing in the actual vector data // to be saved "inline" with the data for each corresponding graph node. .with(new InlineVectors(dim)) @@ -60,13 +61,20 @@ Since we want to include a feature, we need to tell the writer how to obtain the // Supply one map entry for each feature. // The key is a FeatureId enum corresponding to the feature // and the value is a function which generates the feature state for each graph node. +var writerVectors = ravv.threadLocalSupplier(); writer.write(Map.of( FeatureId.INLINE_VECTORS, // we already have a RAVV, so we'll just use that to supply the writer. - nodeId -> new InlineVectors.State(ravv.getVector(nodeId)))); + nodeId -> new InlineVectors.State(writerVectors.get().getVector(nodeId)))); // writer.close() if not using try-with-resources ``` +Supplying every inline feature lets the parallel writer serialize complete records +in batches of up to 64. The source must remain readable until `write` finishes. +`threadLocalSupplier()` provides independent views when a source borrows a scratch +vector; it also works with ordinary preloaded sources. See +[Deferred batch serialization](../deferred-batch-writer.md) for NVQ and migration examples. + At this point the graph index has been written to disk and can be used by creating an instance of `OnDiskGraphIndex`. To do so, we'll need to specify a `RandomAccessReader` implementation which JVector will use to read parts of the file as needed. This interface isn't thread safe, so when creating an `OnDiskGraphIndex` we pass in a `ReaderSupplier` object which can be used by the `OnDiskGraphIndex` to create `RandomAccessReader`s as needed. diff --git a/docs/tutorials/3-larger-than-memory-tutorial.md b/docs/tutorials/3-larger-than-memory-tutorial.md index 3389f6e9a..6001c8e0a 100644 --- a/docs/tutorials/3-larger-than-memory-tutorial.md +++ b/docs/tutorials/3-larger-than-memory-tutorial.md @@ -1,7 +1,7 @@ # JVector Tutorial Part 3: Larger-Than-Memory Indexes with Product Quantization In the previous tutorials, we built indexes in-memory then wrote them to disk. However, this requires enough memory to hold all the vectors during construction. In this tutorial we'll build larger-than-memory indexes using two key techniques: -- Write vectors to disk as they are streamed in +- Write complete records after graph cleanup, reading FP payloads from a replayable source - Keep compressed versions of the vectors in the memory to allow JVector to compute similarities during construction. We'll use Product Quantization (PQ). > [!TIP] @@ -29,7 +29,7 @@ We'll change things up with a different dataset: // the base vectors in-memory. DataSet dataset = DataSets.loadDataSet("e5-small-v2-100k").orElseThrow(() -> new RuntimeException("Dataset doesn't exist or wasn't configured correctly") -); +).getDataSet(); // Remember that RAVVs need not be in-memory in the general case. // We will sample from this RAVV to compute the PQ codebooks. @@ -69,44 +69,35 @@ var pqVectors = new MutablePQVectors(pq); BuildScoreProvider bsp = BuildScoreProvider.pqBuildScoreProvider(vsf, pqVectors); ``` -Since we have a BuildScoreProvider, we're now ready to create and use a `GraphIndexBuilder`. We will also create the corresponding `GraphIndexWriter` at the same time. Unlike the last tutorial we created the writer only after building the complete graph, we'll write full-resolution vectors to disk in tandem with adding them to the graph. +The graph builder uses resident compressed vectors for neighbor comparisons. +Original input vectors are consumed for insertion and reread for final FP payloads; +the source must remain available through serialization. It need not preload the +whole dataset. A one-pass source must retain its final feature states or use the +compatible incremental-writing API instead. ```java -// Initialize the graph parameters M, ef etc -// ... - +// Initialize the graph parameters M, ef etc. Path graphPath = Files.createTempFile("jvector-ltm-graph", null); -// remember to close these manually or use try-with-resources +// Close the builder after final writing, or use try-with-resources. GraphIndexBuilder builder = new GraphIndexBuilder(bsp, dim, M, ef, overflow, alpha, addHierarchy); -OnDiskGraphIndexWriter writer = new OnDiskGraphIndexWriter.Builder(builder.getGraph(), graphPath) - .with(new InlineVectors(dim)) - // Since we start with an empty graph, the writer will, by default, - // assume an ordinal mapping of size 0 (which is obviously incorrect). - // This is easy to rectify if you know the number of vectors beforehand, - // if not you may need to implement OrdinalMapper yourself. - .withMapper(new OrdinalMapper.IdentityMapper(ravv.size() - 1)) - .build(); +var vectors = ravv.threadLocalSupplier(); ``` ## Incremental Index Construction We'll proceed to add vectors to the index one at a time. The general procedure is: - Encode the vector using the PQ codebooks created earlier, and add it to the collection of PQ vectors -- Write the full-resolution vector to disk - Add the new vector to the graph ```java for (int ordinal = 0; ordinal < ravv.size(); ordinal++) { - VectorFloat v = ravv.getVector(ordinal); + VectorFloat v = vectors.get().getVector(ordinal); // Encode and add the vector to the working set of PQ vectors, // which allows the graph builder to access it through the BuildScoreProvder. pqVectors.encodeAndSet(ordinal, v); - // Write the feature (full-resolution vector) for a single vector instead of all at once. - writer.writeFeaturesInline(ordinal, Feature.singleState(FeatureId.INLINE_VECTORS, new InlineVectors.State(v))); - builder.addGraphNode(ordinal, v); } ``` @@ -121,12 +112,25 @@ Once we've added all the vectors we can write the graph structure to disk: // Enforces maximum degree constraint among other things. builder.cleanup(); -// Write the graph structure (neighbor lists) to disk. -// No need to pass-in a feature supplier since we wrote the features incrementally -// using writer.writeInLine -writer.write(Map.of()); +// Serialize FP payloads and adjacency together after cleanup. +try (var writer = new OnDiskParallelGraphIndexWriter.Builder(builder.getGraph(), graphPath) + .withParallelWorkerThreads(PhysicalCoreExecutor.getPhysicalCoreCount()) + .with(new InlineVectors(dim)) + .withMapper(new OrdinalMapper.IdentityMapper(ravv.size() - 1)) + .build()) { + writer.write(Feature.singleStateFactory(FeatureId.INLINE_VECTORS, + ordinal -> new InlineVectors.State(vectors.get().getVector(ordinal)))); +} +builder.close(); ``` +The parallel writer batches up to 64 complete records without changing V6. +Callbacks run concurrently, so the per-worker source views above are important for +borrowed vectors. For NVQ reranking, supply `NVQ.State` instead of `InlineVectors.State`; +see [Deferred batch serialization](../deferred-batch-writer.md) and +[`NvqExample.java`](../../jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/NvqExample.java). +The construction scorer and two-stage query algorithm are unchanged. + We should also save the PQ vectors we created since we'll need them later on when it's time to run searches: ```java @@ -154,7 +158,7 @@ ReaderSupplier graphSupplier = ReaderSupplierFactory.open(graphPath); OnDiskGraphIndex graph = OnDiskGraphIndex.load(graphSupplier); // except that we also need a reader for the PQ vectors ReaderSupplier pqSupplier = ReaderSupplierFactory.open(pqPath); -RandomAccessReader pqReader = pqSupplier.get() +RandomAccessReader pqReader = pqSupplier.get(); // don't forget to close all of the above! (or just use try-with-resources) // we need to have the PQ vectors in memory @@ -167,7 +171,7 @@ int rerankK = (int) (topK * overqueryFactor); // a friendly reminder that searchers need closing too GraphSearcher searcher = new GraphSearcher(graph); var graphRavv = (RandomAccessVectorValues) searcher.getView(); -VectorFloat query : dataset.getQueryVectors().get(0); +VectorFloat query = dataset.getQueryVectors().get(0); // Two-phase search: // 1. ApproximateScoreFunction (ASF) uses compressed vectors for fast initial search diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/NodeRecordTask.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/NodeRecordTask.java index 7e756dab9..a825a8480 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/NodeRecordTask.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/NodeRecordTask.java @@ -28,6 +28,7 @@ import java.util.Map; import java.util.concurrent.Callable; import java.util.concurrent.ExecutionException; +import java.util.function.BooleanSupplier; import java.util.function.IntFunction; /** @@ -38,8 +39,8 @@ * Each task processes a contiguous range of ordinals to reduce task creation overhead. *

* This writes directly to the AsynchronousFileChannel using position-based writes with writeFully - * to ensure all bytes are written before returning. This eliminates race conditions where the OS - * buffer cache hasn't flushed data before subsequent reads occur. + * to ensure each write completes before its buffer is reused. Durability is handled by + * the caller's final channel.force(true). */ class NodeRecordTask implements Callable { private final int startOrdinal; // Inclusive @@ -53,6 +54,7 @@ class NodeRecordTask implements Callable { private final long baseOffset; // Base file offset for L0 (offsets calculated per-ordinal) private final AsynchronousFileChannel channel; private final ByteBuffer buffer; // Thread-local buffer for building record components + private final BooleanSupplier stopRequested; NodeRecordTask(int startOrdinal, int endOrdinal, @@ -65,6 +67,25 @@ class NodeRecordTask implements Callable { long baseOffset, AsynchronousFileChannel channel, ByteBuffer buffer) { + this(startOrdinal, endOrdinal, ordinalMapper, graph, view, inlineFeatures, featureStateSuppliers, + recordSize, baseOffset, channel, buffer, () -> false); + } + + /// Create a task that stops before its next record once `stopRequested` returns true, for + /// example after another task failed or the writing thread was interrupted. A stopped task + /// returns normally; the writer reports the failure that requested the stop. + NodeRecordTask(int startOrdinal, + int endOrdinal, + OrdinalMapper ordinalMapper, + ImmutableGraphIndex graph, + ImmutableGraphIndex.View view, + List inlineFeatures, + Map> featureStateSuppliers, + int recordSize, + long baseOffset, + AsynchronousFileChannel channel, + ByteBuffer buffer, + BooleanSupplier stopRequested) { this.startOrdinal = startOrdinal; this.endOrdinal = endOrdinal; this.ordinalMapper = ordinalMapper; @@ -76,6 +97,7 @@ class NodeRecordTask implements Callable { this.baseOffset = baseOffset; this.channel = channel; this.buffer = buffer; + this.stopRequested = stopRequested; } /** @@ -95,8 +117,8 @@ private static void writeFully(AsynchronousFileChannel channel, ByteBuffer buffe long currentPosition = position; while (buffer.hasRemaining()) { int written = channel.write(buffer, currentPosition).get(); - if (written < 0) { - throw new IOException("Channel closed while writing"); + if (written <= 0) { + throw new IOException("Channel made no progress while writing"); } currentPosition += written; } @@ -107,16 +129,61 @@ public Void call() throws Exception { // Reuse writer and buffer across all ordinals in this range var writer = new ByteBufferIndexWriter(buffer); + int payloadBytes = Integer.BYTES * (2 + graph.getDegree(0)) + + inlineFeatures.stream().mapToInt(Feature::featureSize).sum(); + if (recordSize != payloadBytes) { + throw new IllegalArgumentException("Record size must match its serialized payload"); + } + final boolean completeRecord = inlineFeatures.stream().allMatch(f -> featureStateSuppliers.get(f.id()) != null); + final int batch = buffer.capacity() / recordSize; for (int newOrdinal = startOrdinal; newOrdinal < endOrdinal; newOrdinal++) { + if (stopRequested.getAsBoolean()) return null; var originalOrdinal = ordinalMapper.newToOld(newOrdinal); long recordOffset = baseOffset + (long) newOrdinal * recordSize; long currentPosition = recordOffset; + // A complete record can share one bounded buffer and one positioned write. + if (completeRecord) { + if (buffer.position() == buffer.limit()) writer.reset(); + int recordBegin = buffer.position(); + writer.writeInt(originalOrdinal == OrdinalMapper.OMITTED ? -1 : newOrdinal); + if (originalOrdinal == OrdinalMapper.OMITTED) { + for (var feature : inlineFeatures) + for (int i = 0; i < feature.featureSize(); i++) writer.writeByte(0); + writer.writeInt(0); + for (int n = 0; n < graph.getDegree(0); n++) writer.writeInt(-1); + } else { + if (!graph.containsNode(originalOrdinal)) { + throw new IllegalStateException("Missing graph node " + originalOrdinal); + } + for (var feature : inlineFeatures) + feature.writeInline(writer, featureStateSuppliers.get(feature.id()).apply(originalOrdinal)); + var neighbors = view.getNeighborsIterator(0, originalOrdinal); + if (neighbors.size() > graph.getDegree(0)) throw new IllegalStateException("Too many neighbors for " + originalOrdinal); + writer.writeInt(neighbors.size()); + int n = 0; + while (neighbors.hasNext()) { + int mapped = ordinalMapper.oldToNew(neighbors.nextInt()); + if (mapped < 0 || mapped > ordinalMapper.maxOrdinal()) throw new IllegalStateException("Invalid mapped neighbor " + mapped); + writer.writeInt(mapped); n++; + } + for (; n < graph.getDegree(0); n++) writer.writeInt(-1); + } + if (writer.position() != recordBegin + recordSize) throw new IllegalStateException("Wrong serialized record length"); + if (buffer.position() < recordSize * batch && newOrdinal != endOrdinal - 1) continue; + long batchOffset = recordOffset - recordBegin; + buffer.flip(); + // writeFully waits for completion before this worker reuses its buffer. + writeFully(channel, buffer, batchOffset); + buffer.clear(); + continue; + } + // Reset buffer for this ordinal writer.reset(); // Write node ordinal - writer.writeInt(newOrdinal); + writer.writeInt(originalOrdinal == OrdinalMapper.OMITTED ? -1 : newOrdinal); ByteBuffer ordinalData = writer.cloneBuffer(); writeFully(channel, ordinalData, currentPosition); currentPosition += Integer.BYTES; @@ -132,8 +199,10 @@ public Void call() throws Exception { } } ByteBuffer featureData = writer.cloneBuffer(); + // writeFully consumes the buffer; retain its length before writing. + int featureBytes = featureData.remaining(); writeFully(channel, featureData, currentPosition); - currentPosition += featureData.remaining(); + currentPosition += featureBytes; // Write empty neighbor list writer.reset(); diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskParallelGraphIndexWriter.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskParallelGraphIndexWriter.java index 4e8e7ac2d..0207225f2 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskParallelGraphIndexWriter.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskParallelGraphIndexWriter.java @@ -70,6 +70,7 @@ public class OnDiskParallelGraphIndexWriter extends RandomAccessOnDiskGraphIndex private final int parallelWorkerThreads; private final boolean parallelUseDirectBuffers; private final ExecutorService parallelExecutor; + private final int parallelBatchSize; /** * Constructs an OnDiskParallelGraphIndexWriter with all parameters including file path @@ -118,6 +119,24 @@ public class OnDiskParallelGraphIndexWriter extends RandomAccessOnDiskGraphIndex int parallelWorkerThreads, boolean parallelUseDirectBuffers, ExecutorService parallelExecutor) + { + this(randomAccessWriter, version, startOffset, graph, oldToNewOrdinals, dimension, features, + filePath, parallelWorkerThreads, parallelUseDirectBuffers, parallelExecutor, + ParallelGraphWriter.DEFAULT_RECORDS_PER_BATCH); + } + + OnDiskParallelGraphIndexWriter(RandomAccessWriter randomAccessWriter, + int version, + long startOffset, + ImmutableGraphIndex graph, + OrdinalMapper oldToNewOrdinals, + int dimension, + EnumMap features, + Path filePath, + int parallelWorkerThreads, + boolean parallelUseDirectBuffers, + ExecutorService parallelExecutor, + int parallelBatchSize) { super(randomAccessWriter, version, startOffset, graph, oldToNewOrdinals, dimension, features); if (filePath == null) { @@ -127,6 +146,7 @@ public class OnDiskParallelGraphIndexWriter extends RandomAccessOnDiskGraphIndex this.parallelWorkerThreads = parallelWorkerThreads; this.parallelUseDirectBuffers = parallelUseDirectBuffers; this.parallelExecutor = parallelExecutor; + this.parallelBatchSize = parallelBatchSize; } /** @@ -153,7 +173,8 @@ protected void writeL0Records(ImmutableGraphIndex.View view, var config = new ParallelGraphWriter.Config( parallelWorkerThreads, parallelUseDirectBuffers, - 4 // Default task multiplier (4x cores) + 4, // Default task multiplier (4x cores) + parallelBatchSize ); try (var parallelWriter = new ParallelGraphWriter( @@ -200,6 +221,7 @@ public static class Builder extends AbstractGraphIndexWriter.Builder */ class ParallelGraphWriter implements AutoCloseable { + // Amortize positioned writes with bounded worker-local buffers. + static final int DEFAULT_RECORDS_PER_BATCH = 64; private final RandomAccessWriter writer; private final ImmutableGraphIndex graph; private final ExecutorService executor; private final boolean ownsExecutor; private final ThreadLocal viewPerThread; - private final ThreadLocal bufferPerThread; + private final ExplicitThreadLocal bufferPerThread; private final CopyOnWriteArrayList allViews = new CopyOnWriteArrayList<>(); private final int recordSize; private final Path filePath; @@ -81,6 +85,7 @@ static class Config { final int workerThreads; final boolean useDirectBuffers; final int taskMultiplier; + final int recordsPerBatch; /** * @param workerThreads number of worker threads for building records (0 = use available processors) @@ -89,6 +94,14 @@ static class Config { * (4x = good balance for most use cases, higher = more fine-grained parallelism) */ public Config(int workerThreads, boolean useDirectBuffers, int taskMultiplier) { + this(workerThreads, useDirectBuffers, taskMultiplier, DEFAULT_RECORDS_PER_BATCH); + } + + public Config(int workerThreads, boolean useDirectBuffers, int taskMultiplier, int recordsPerBatch) { + if (recordsPerBatch <= 0) { + throw new IllegalArgumentException("Record batch size must be positive"); + } + this.recordsPerBatch = recordsPerBatch; this.workerThreads = workerThreads <= 0 ? Runtime.getRuntime().availableProcessors() : workerThreads; this.useDirectBuffers = useDirectBuffers; this.taskMultiplier = taskMultiplier <= 0 ? 4 : taskMultiplier; @@ -143,6 +156,10 @@ public ParallelGraphWriter(RandomAccessWriter writer, this.graph = graph; this.filePath = Objects.requireNonNull(filePath); this.taskMultiplier = config.taskMultiplier; + // Validate buffer size before creating a pool that would otherwise need cleanup. + this.recordSize = Math.addExact(Integer.BYTES * (2 + graph.getDegree(0)), + inlineFeatures.stream().mapToInt(Feature::featureSize).sum()); + final int bufferSize = Math.multiplyExact(recordSize, config.recordsPerBatch); if (externalExecutor != null) { this.executor = externalExecutor; this.ownsExecutor = false; @@ -157,12 +174,6 @@ public ParallelGraphWriter(RandomAccessWriter writer, this.ownsExecutor = true; } - // Compute fixed record size for L0 - this.recordSize = Integer.BYTES // node ordinal - + inlineFeatures.stream().mapToInt(Feature::featureSize).sum() - + Integer.BYTES // neighbor count - + graph.getDegree(0) * Integer.BYTES; // neighbors + padding - // Thread-local views for safe neighbor iteration // CopyOnWriteArrayList handles concurrent additions safely this.viewPerThread = ThreadLocal.withInitial(() -> { @@ -173,9 +184,8 @@ public ParallelGraphWriter(RandomAccessWriter writer, // Thread-local buffers to avoid allocation overhead // Use BIG_ENDIAN to match Java DataOutput specification - final int bufferSize = recordSize; final boolean useDirect = config.useDirectBuffers; - this.bufferPerThread = ThreadLocal.withInitial(() -> { + this.bufferPerThread = ExplicitThreadLocal.withInitial(() -> { ByteBuffer buffer = useDirect ? ByteBuffer.allocateDirect(bufferSize) : ByteBuffer.allocate(bufferSize); buffer.order(java.nio.ByteOrder.BIG_ENDIAN); return buffer; @@ -221,57 +231,66 @@ public void writeL0Records(OrdinalMapper ordinalMapper, var opts = new java.util.HashSet(); opts.add(StandardOpenOption.WRITE); opts.add(StandardOpenOption.READ); - + try (var channel = AsynchronousFileChannel.open(filePath, opts, null)) { List> taskFutures = new ArrayList<>(numTasks); + // Set on interruption or the first failure, so remaining ranges stop instead of running to completion. + var stop = new AtomicBoolean(); + boolean submitted = false; - // Submit range-based tasks - for (int i = 0; i < numTasks; i++) { - int startOrdinal = i * ordinalsPerTask; - int endOrdinal = Math.min(startOrdinal + ordinalsPerTask, totalOrdinals); - - // Skip if range is empty (can happen with final task) - if (startOrdinal >= totalOrdinals) { - break; - } + // Always join submitted tasks before closing the channel or supplier resources. + try { + // Submit range-based tasks + for (int i = 0; i < numTasks; i++) { + int startOrdinal = i * ordinalsPerTask; + int endOrdinal = Math.min(startOrdinal + ordinalsPerTask, totalOrdinals); - final int start = startOrdinal; - final int end = endOrdinal; + // Skip if range is empty (can happen with final task) + if (startOrdinal >= totalOrdinals) { + break; + } - Future future = executor.submit(() -> { - var view = viewPerThread.get(); - var buffer = bufferPerThread.get(); + final int start = startOrdinal; + final int end = endOrdinal; - var task = new NodeRecordTask( - start, // Start of range (inclusive) - end, // End of range (exclusive) - ordinalMapper, - graph, - view, - inlineFeatures, - featureStateSuppliers, - recordSize, - baseOffset, // Base offset (task calculates per-ordinal offsets) - channel, // Async file channel for position-based writes - buffer // Thread-local buffer - ); + Future future = executor.submit(() -> { + if (stop.get()) return null; + try { + var view = viewPerThread.get(); + var buffer = bufferPerThread.get(); - return task.call(); - }); + var task = new NodeRecordTask( + start, // Start of range (inclusive) + end, // End of range (exclusive) + ordinalMapper, + graph, + view, + inlineFeatures, + featureStateSuppliers, + recordSize, + baseOffset, // Base offset (task calculates per-ordinal offsets) + channel, // Async file channel for position-based writes + buffer, // Thread-local buffer + stop::get + ); - taskFutures.add(future); - } + return task.call(); + } catch (Exception | Error failure) { + // Signal immediately: awaitTasks may still be joining an earlier range. + stop.set(true); + throw failure; + } finally { + // External executor threads must not retain this writer's graph views. + viewPerThread.remove(); + } + }); - // Wait for all tasks to complete - each task writes synchronously using writeFully - for (Future taskFuture : taskFutures) { - try { - taskFuture.get(); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new IOException("Interrupted while writing records", e); - } catch (ExecutionException e) { - throw unwrapExecutionException(e); + taskFutures.add(future); } + submitted = true; + } finally { + if (!submitted) stop.set(true); + awaitTasks(taskFutures, stop); } // Force all writes to disk before closing the channel. @@ -283,23 +302,35 @@ public void writeL0Records(OrdinalMapper ordinalMapper, } } - /** - * Unwraps ExecutionException to throw the underlying cause. - * Handles IOException, RuntimeException, and wraps other exceptions. - * - * @param e the execution exception to unwrap - * @return an IOException wrapping the cause - * @throws RuntimeException if the cause is a RuntimeException - */ - private IOException unwrapExecutionException(ExecutionException e) { - Throwable cause = e.getCause(); - if (cause instanceof IOException) { - return (IOException) cause; - } else if (cause instanceof RuntimeException) { - throw (RuntimeException) cause; - } else { - throw new RuntimeException("Error building node record", cause); + /// Join every task, so no worker still uses the channel or suppliers when this returns. On + /// interruption or the first task failure, set `stop` so the remaining tasks skip or abandon + /// their ranges rather than writing them, then rethrow the first failure. + private void awaitTasks(List> futures, AtomicBoolean stop) throws IOException { + Throwable failure = null; + boolean interrupted = false; + for (var future : futures) { + boolean done = false; + while (!done) { + try { + future.get(); + done = true; + } catch (InterruptedException e) { + interrupted = true; + stop.set(true); + if (failure == null) failure = new IOException("Interrupted while writing records", e); + } catch (ExecutionException e) { + stop.set(true); + if (failure == null) failure = e.getCause(); + else if (failure != e.getCause()) failure.addSuppressed(e.getCause()); + done = true; + } + } } + if (interrupted) Thread.currentThread().interrupt(); + if (failure instanceof IOException) throw (IOException) failure; + if (failure instanceof RuntimeException) throw (RuntimeException) failure; + if (failure instanceof Error) throw (Error) failure; + if (failure != null) throw new IOException("Error building node record", failure); } /** @@ -330,6 +361,7 @@ public void close() throws IOException { view.close(); } allViews.clear(); + bufferPerThread.close(); } catch (IOException e) { throw e; } catch (Exception e) { diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/RandomAccessOnDiskGraphIndexWriter.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/RandomAccessOnDiskGraphIndexWriter.java index e3ce70e21..debb2b5bd 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/RandomAccessOnDiskGraphIndexWriter.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/RandomAccessOnDiskGraphIndexWriter.java @@ -141,7 +141,7 @@ protected abstract void writeL0Records(ImmutableGraphIndex.View view, * Both produce an identical on-disk format. * *

Parallel-specific options ({@link #withParallelWorkerThreads}, {@link #withParallelDirectBuffers}, - * {@link #withExecutor}) are accepted unconditionally but silently ignored when the sequential + * {@link #withExecutor}, {@link #withParallelBatchSize}) are accepted unconditionally but silently ignored when the sequential * writer is selected. */ public static class Builder extends AbstractGraphIndexWriter.Builder { @@ -150,6 +150,7 @@ public static class Builder extends AbstractGraphIndexWriter.Builder Arrays.stream(s.split("\\s"))) diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java index d2aa5ceda..23eafdf94 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java @@ -35,6 +35,7 @@ import io.github.jbellis.jvector.example.util.FilteredForkJoinPool; import io.github.jbellis.jvector.example.util.OnDiskGraphIndexCache; import io.github.jbellis.jvector.example.yaml.MetricSelection; +import io.github.jbellis.jvector.management.GraphIndexBuilderConfig; import io.github.jbellis.jvector.graph.ImmutableGraphIndex; import io.github.jbellis.jvector.graph.GraphIndexBuilder; import io.github.jbellis.jvector.graph.VectorAccess; @@ -414,7 +415,6 @@ private static BuildOnDiskResult buildOnDisk(List> feat // use the inline vectors index as the score provider for graph construction Map, RandomAccessOnDiskGraphIndexWriter> writers = new HashMap<>(); - Map, Map>> suppliers = new HashMap<>(); Map, Map, Feature.State>>> inputSuppliers = new HashMap<>(); RandomAccessOnDiskGraphIndexWriter scoringWriter = null; int n = 0; @@ -429,7 +429,6 @@ private static BuildOnDiskResult buildOnDisk(List> feat var bws = builderWithSuppliers(features, builder.getGraph(), graphPath, floatVectors, pq == null ? null : pq.getCompressor(), constructionMetrics); var writer = bws.builder.build(); writers.put(features, writer); - suppliers.put(features, bws.suppliers); inputSuppliers.put(features, bws.inputSuppliers); if (features.contains(FeatureId.INLINE_VECTORS) || features.contains(FeatureId.NVQ_VECTORS)) { scoringWriter = writer; @@ -439,41 +438,45 @@ private static BuildOnDiskResult buildOnDisk(List> feat throw new IllegalStateException("Bench looks for either NVQ_VECTORS or INLINE_VECTORS feature set for scoring compressed builds."); } - // build the graph incrementally + // The parallel writer calls final suppliers from its workers, so original-vector features + // can be encoded after cleanup. The sequential writer calls them on one thread, so it keeps + // encoding on the insertion threads and writing each node's features inline. + boolean deferFeatures = writers.values().stream() + .allMatch(writer -> writer instanceof OnDiskParallelGraphIndexWriter); long startTime = System.nanoTime(); if (ds.useDatasetPrefetching()) { VectorAccess.forEach(floatVectors, PhysicalCoreExecutor.pool(), (node, vector) -> { - writeInlineStates(writers, suppliers, inputSuppliers, node, vector); - // Keep main's original insertion query, not a decoded quantized substitute. + if (!deferFeatures) writeInlineStates(writers, inputSuppliers, node, vector); builder.addGraphNode(node, vector); }); } else { var vv = floatVectors.threadLocalSupplier(); PhysicalCoreExecutor.pool().submit(() -> { IntStream.range(0, floatVectors.size()).parallel().forEach(node -> { - writeInlineStates(writers, suppliers, null, node, null); - builder.addGraphNode(node, vv.get().getVector(node)); + var vector = vv.get().getVector(node); + if (!deferFeatures) writeInlineStates(writers, inputSuppliers, node, vector); + builder.addGraphNode(node, vector); }); }).join(); } builder.cleanup(); - // write the edge lists and close the writers - // if our feature set contains Fused PQ, we need a Fused ADC write-time supplier (as we don't have neighbor information during writeInline) + // With deferred features, supply every feature so the parallel writer emits complete records. writers.entrySet().stream().parallel().forEach(entry -> { var writer = entry.getValue(); var features = entry.getKey(); - Map> writeSuppliers; - if (features.contains(FeatureId.FUSED_PQ)) { - writeSuppliers = new EnumMap<>(FeatureId.class); - var view = builder.getGraph().getView(); - writeSuppliers.put(FeatureId.FUSED_PQ, ordinal -> new FusedPQ.State(view, pq, ordinal)); - } else { - writeSuppliers = Map.of(); - } - try { + try (var vectors = new FinalWriteVectors(floatVectors); + var view = builder.getGraph().getView(); + var closeWriter = writer) { + Map> writeSuppliers = new EnumMap<>(FeatureId.class); + if (deferFeatures) { + inputSuppliers.get(features).forEach((id, supplier) -> + writeSuppliers.put(id, ordinal -> supplier.apply(vectors.get(ordinal)))); + } + if (features.contains(FeatureId.FUSED_PQ)) { + writeSuppliers.put(FeatureId.FUSED_PQ, ordinal -> new FusedPQ.State(view, pq, ordinal)); + } writer.write(writeSuppliers); - writer.close(); } catch (IOException e) { throw new UncheckedIOException(e); } @@ -505,22 +508,76 @@ private static BuildOnDiskResult buildOnDisk(List> feat return new BuildOnDiskResult(indexes, fileSizes); } + /// Write one node's original-vector features to every writer during insertion. The borrowed + /// vector is encoded and written before the caller advances to another vector. private static void writeInlineStates( Map, RandomAccessOnDiskGraphIndexWriter> writers, - Map, Map>> suppliers, Map, Map, Feature.State>>> inputSuppliers, int node, VectorFloat vector) { writers.forEach((features, writer) -> { var states = new EnumMap(FeatureId.class); - if (inputSuppliers == null) - suppliers.get(features).forEach((id, supplier) -> states.put(id, supplier.apply(node))); - else - inputSuppliers.get(features).forEach((id, supplier) -> states.put(id, supplier.apply(vector))); - try { writer.writeInline(node, states); } + inputSuppliers.get(features).forEach((id, supplier) -> states.put(id, supplier.apply(vector))); + try { writer.writeFeaturesInline(node, states); } catch (IOException e) { throw new UncheckedIOException(e); } }); } + /** Worker-local bounded input ranges; a borrowed vector is consumed before advancing. */ + static final class FinalWriteVectors implements AutoCloseable { + private final RandomAccessVectorValues source; + private final java.util.concurrent.ConcurrentLinkedQueue ranges = + new java.util.concurrent.ConcurrentLinkedQueue<>(); + private final ThreadLocal local = ThreadLocal.withInitial(() -> { + var range = new Range(); + ranges.add(range); + return range; + }); + + FinalWriteVectors(RandomAccessVectorValues source) { this.source = source; } + + VectorFloat get(int ordinal) { + Objects.checkIndex(ordinal, source.size()); + var range = local.get(); + if (range.cursor == null || ordinal != range.ordinal) { + if (range.cursor == null || ordinal != range.ordinal + 1 || ordinal == range.end) { + range.close(); + range.end = (int) Math.min(source.size(), (long) ordinal + 64); + range.cursor = VectorAccess.openRange(source, ordinal, range.end); + } + if (!range.cursor.next() || range.cursor.ordinal() != ordinal) + throw new IllegalStateException("Incomplete final-write input range"); + range.ordinal = ordinal; + } + return range.cursor.vector(); + } + + @Override public void close() { + // writer.write joins its workers before cursors and their read-ahead are released. + RuntimeException failure = null; + try { + for (var range : ranges) { + try { range.close(); } + catch (RuntimeException e) { + if (failure == null) failure = e; + else if (failure != e) failure.addSuppressed(e); + } + } + } finally { ranges.clear(); local.remove(); } + if (failure != null) throw failure; + } + + private static final class Range { + io.github.jbellis.jvector.graph.VectorCursor cursor; + int ordinal = -1, end; + void close() { + if (cursor != null) { + try { cursor.close(); } + finally { cursor = null; ordinal = -1; } + } + } + } + } + private static BuilderWithSuppliers builderWithSuppliers(Set features, ImmutableGraphIndex onHeapGraph, Path outPath, @@ -532,6 +589,7 @@ private static BuilderWithSuppliers builderWithSuppliers(Set features var identityMapper = new OrdinalMapper.IdentityMapper(floatVectors.size() - 1); var builder = new RandomAccessOnDiskGraphIndexWriter.Builder(onHeapGraph, outPath); builder.withMapper(identityMapper); + builder.withParallelWorkerThreads(PhysicalCoreExecutor.getPhysicalCoreCount()); Map> suppliers = new EnumMap<>(FeatureId.class); Map, Feature.State>> inputSuppliers = new EnumMap<>(FeatureId.class); @@ -567,6 +625,14 @@ private static BuilderWithSuppliers builderWithSuppliers(Set features return new BuilderWithSuppliers(builder, suppliers, inputSuppliers); } + // Entry-point policy only; the library's writer selection remains unchanged. + static void configureWriterForBenchmark() { + String value = System.getProperty("jvector.parallel_writer", "true").trim(); + if (!value.equalsIgnoreCase("true") && !value.equalsIgnoreCase("false")) + throw new IllegalArgumentException("jvector.parallel_writer must be true or false"); + GraphIndexBuilderConfig.getInstance().setParallelBuild(Boolean.parseBoolean(value)); + } + public static void setDiagnosticLevel(int diagLevel) { diagnostic_level = diagLevel; } diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/IPCService.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/IPCService.java index 3c125fd2b..92782c316 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/IPCService.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/IPCService.java @@ -23,7 +23,11 @@ import io.github.jbellis.jvector.graph.GraphSearcher; import io.github.jbellis.jvector.graph.RandomAccessVectorValues; import io.github.jbellis.jvector.graph.SearchResult; -import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; +import io.github.jbellis.jvector.graph.disk.OnDiskParallelGraphIndexWriter; +import io.github.jbellis.jvector.graph.disk.feature.Feature; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; +import io.github.jbellis.jvector.graph.disk.feature.InlineVectors; +import io.github.jbellis.jvector.util.PhysicalCoreExecutor; import io.github.jbellis.jvector.graph.similarity.DefaultSearchScoreProvider; import io.github.jbellis.jvector.graph.similarity.ScoreFunction; import io.github.jbellis.jvector.quantization.CompressedVectors; @@ -194,7 +198,14 @@ private static ImmutableGraphIndex flushGraphIndex(ImmutableGraphIndex index, Ra var testDirectory = Files.createTempDirectory("BenchGraphDir"); var graphPath = testDirectory.resolve("graph.bin"); - OnDiskGraphIndex.write(index, ravv, graphPath); + var vectors = ravv.threadLocalSupplier(); + try (var writer = new OnDiskParallelGraphIndexWriter.Builder(index, graphPath) + .withParallelWorkerThreads(PhysicalCoreExecutor.getPhysicalCoreCount()) + .with(new InlineVectors(ravv.dimension())) + .build()) { + writer.write(Feature.singleStateFactory(FeatureId.INLINE_VECTORS, + ordinal -> new InlineVectors.State(vectors.get().getVector(ordinal)))); + } return index; } catch (IOException e) { throw new IOError(e); diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/SiftSmall.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/SiftSmall.java index 320b7f1a7..d029262cb 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/SiftSmall.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/SiftSmall.java @@ -34,7 +34,8 @@ import io.github.jbellis.jvector.graph.disk.feature.InlineVectors; import io.github.jbellis.jvector.graph.disk.feature.NVQ; import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; -import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndexWriter; +import io.github.jbellis.jvector.graph.disk.OnDiskParallelGraphIndexWriter; +import io.github.jbellis.jvector.util.PhysicalCoreExecutor; import io.github.jbellis.jvector.graph.disk.OrdinalMapper; import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider; import io.github.jbellis.jvector.graph.similarity.DefaultSearchScoreProvider; @@ -59,7 +60,6 @@ import java.nio.file.Files; import java.nio.file.Path; import java.util.List; -import java.util.Map; import java.util.Random; import java.util.concurrent.ThreadLocalRandom; import java.util.function.Function; @@ -155,7 +155,7 @@ public static void siftPersisted(List> baseVectors, List> baseVectors, List> baseVectors, List> baseVectors, List v = baseVectors.get(ordinal); // compress the new vector and add it to the PQVectors pqv.encodeAndSet(ordinal, v); - // write the full vector to disk - writer.writeInline(ordinal, Feature.singleState(FeatureId.INLINE_VECTORS, new InlineVectors.State(v))); - // now add it to the graph -- the previous steps must be completed first since the PQVectors - // and InlineVectorValues are both used during the search that runs as part of addGraphNode construction + // PQ is the resident construction scorer; FP is written after cleanup. builder.addGraphNode(ordinal, v); } @@ -253,8 +245,16 @@ public static void siftDiskAnnLTM(List> baseVectors, List new InlineVectors.State(vectors.get().getVector(ordinal)))); + } pqv.write(pqWriter); } @@ -292,11 +292,6 @@ public static void siftDiskAnnLTMWithNVQ(List> baseVectors, List< Path pqPath = Files.createTempFile("siftsmall", ".pq"); // Builder creation looks mostly the same try (GraphIndexBuilder builder = new GraphIndexBuilder(bsp, ravv.dimension(), 16, 100, 1.2f, 1.2f, false, true); - // explicit Writer for the first time, this is what's behind OnDiskGraphIndex.write - OnDiskGraphIndexWriter writer = new OnDiskGraphIndexWriter.Builder(builder.getGraph(), indexPath) - .with(new NVQ(nvq)) - .withMapper(new OrdinalMapper.IdentityMapper(baseVectors.size() - 1)) - .build(); // output for the compressed vectors IndexWriter pqWriter = new BufferedRandomAccessWriter(pqPath)) { @@ -305,10 +300,7 @@ public static void siftDiskAnnLTMWithNVQ(List> baseVectors, List< VectorFloat v = baseVectors.get(ordinal); // compress the new vector and add it to the PQVectors pqv.encodeAndSet(ordinal, v); - // write the full vector to disk - writer.writeInline(ordinal, Feature.singleState(FeatureId.NVQ_VECTORS, new NVQ.State(nvq.encode(v)))); - // now add it to the graph -- the previous steps must be completed first since the PQVectors - // and InlineVectorValues are both used during the search that runs as part of addGraphNode construction + // PQ is the resident construction scorer; NVQ is written after cleanup. builder.addGraphNode(ordinal, v); } @@ -316,8 +308,16 @@ public static void siftDiskAnnLTMWithNVQ(List> baseVectors, List< // that we don't need to worry about here builder.cleanup(); - // finish writing the index (by filling in the edge lists) and write our completed PQVectors - writer.write(Map.of()); + // Supply all inline features during final complete-record serialization. + var vectors = ravv.threadLocalSupplier(); + try (var writer = new OnDiskParallelGraphIndexWriter.Builder(builder.getGraph(), indexPath) + .withParallelWorkerThreads(PhysicalCoreExecutor.getPhysicalCoreCount()) + .with(new NVQ(nvq)) + .withMapper(new OrdinalMapper.IdentityMapper(baseVectors.size() - 1)) + .build()) { + writer.write(Feature.singleStateFactory(FeatureId.NVQ_VECTORS, + ordinal -> new NVQ.State(nvq.encode(vectors.get().getVector(ordinal))))); + } pqv.write(pqWriter); } @@ -337,6 +337,17 @@ public static void siftDiskAnnLTMWithNVQ(List> baseVectors, List< } } + private static void writeCompleteGraph(ImmutableGraphIndex graph, RandomAccessVectorValues source, Path path) throws IOException { + var vectors = source.threadLocalSupplier(); + try (var writer = new OnDiskParallelGraphIndexWriter.Builder(graph, path) + .withParallelWorkerThreads(PhysicalCoreExecutor.getPhysicalCoreCount()) + .with(new InlineVectors(source.dimension())) + .build()) { + writer.write(Feature.singleStateFactory(FeatureId.INLINE_VECTORS, + ordinal -> new InlineVectors.State(vectors.get().getVector(ordinal)))); + } + } + // // Utilities and main() harness // diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/DiskIntro.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/DiskIntro.java index cfb70da09..6f238d81f 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/DiskIntro.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/DiskIntro.java @@ -34,7 +34,8 @@ import io.github.jbellis.jvector.graph.RandomAccessVectorValues; import io.github.jbellis.jvector.graph.SearchResult; import io.github.jbellis.jvector.graph.disk.GraphIndexWriter; -import io.github.jbellis.jvector.graph.disk.GraphIndexWriterTypes; +import io.github.jbellis.jvector.graph.disk.OnDiskParallelGraphIndexWriter; +import io.github.jbellis.jvector.util.PhysicalCoreExecutor; import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; import io.github.jbellis.jvector.graph.disk.feature.FeatureId; import io.github.jbellis.jvector.graph.disk.feature.InlineVectors; @@ -45,12 +46,13 @@ import io.github.jbellis.jvector.vector.VectorSimilarityFunction; import io.github.jbellis.jvector.vector.types.VectorFloat; -// Provides the code used in disk-tutorial.md. -// If you edit this file you may also want to edit disk-tutoral.md +// Provides the code used in 2-disk-tutorial.md. +// Keep the tutorial and this example in sync. public class DiskIntro { public static void main(String[] args) throws IOException { // This is a preconfigured dataset that will be downloaded automatically. - DataSet dataset = DataSets.loadDataSet("ada002-100k").orElseThrow(() -> + String datasetName = args.length == 0 ? "ada002-100k" : args[0]; + DataSet dataset = DataSets.loadDataSet(datasetName).orElseThrow(() -> new RuntimeException("Dataset doesn't exist or wasn't configured correctly") ).getDataSet(); @@ -76,22 +78,25 @@ public static void main(String[] args) throws IOException { heapGraph = builder.build(ravv); } + var writerVectors = ravv.threadLocalSupplier(); Path graphPath = Files.createTempFile("jvector-example-graph", null); // or wherever you want to save the graph try ( // Create a writer for the on-heap graph we just built. - GraphIndexWriter writer = GraphIndexWriter.getBuilderFor(GraphIndexWriterTypes.RANDOM_ACCESS_PARALLEL, heapGraph, graphPath) + GraphIndexWriter writer = new OnDiskParallelGraphIndexWriter.Builder(heapGraph, graphPath) + .withParallelWorkerThreads(PhysicalCoreExecutor.getPhysicalCoreCount()) // Let the writer know that we'll also be passing in the actual vector data // to be saved "inline" with the data for each corresponding graph node. .with(new InlineVectors(dim)) .build(); ) { + // Supplying every inline feature enables complete-record batches (default 64). // Supply one map entry for each feature. // The key is a FeatureId enum corresponding to the feature // and the value is a function which generates the feature state for each graph node. writer.write(Map.of( FeatureId.INLINE_VECTORS, // we already have a RAVV, so we'll just use that to supply the writer. - nodeId -> new InlineVectors.State(ravv.getVector(nodeId)))); + nodeId -> new InlineVectors.State(writerVectors.get().getVector(nodeId)))); } // ReaderSupplierFactory automatically picks an available RandomAccessReader implementation diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/LargerThanMemory.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/LargerThanMemory.java index 5f22b1cbc..7afb05da0 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/LargerThanMemory.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/LargerThanMemory.java @@ -17,12 +17,10 @@ package io.github.jbellis.jvector.example.tutorial; import java.io.IOException; -import java.io.UncheckedIOException; import java.nio.file.Files; import java.nio.file.Path; import java.util.ArrayList; import java.util.List; -import java.util.Map; import java.util.stream.IntStream; import io.github.jbellis.jvector.disk.BufferedRandomAccessWriter; @@ -37,7 +35,7 @@ import io.github.jbellis.jvector.graph.RandomAccessVectorValues; import io.github.jbellis.jvector.graph.SearchResult; import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; -import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndexWriter; +import io.github.jbellis.jvector.graph.disk.OnDiskParallelGraphIndexWriter; import io.github.jbellis.jvector.graph.disk.OrdinalMapper; import io.github.jbellis.jvector.graph.disk.feature.Feature; import io.github.jbellis.jvector.graph.disk.feature.FeatureId; @@ -60,11 +58,12 @@ public static void main(String[] args) throws IOException { // The DataSet provided by loadDataSet is in-memory, // but you can apply the same technique even when you don't have // the base vectors in-memory. - DataSet dataset = DataSets.loadDataSet("e5-small-v2-100k").orElseThrow(() -> + String datasetName = args.length == 0 ? "e5-small-v2-100k" : args[0]; + DataSet dataset = DataSets.loadDataSet(datasetName).orElseThrow(() -> new RuntimeException("Dataset doesn't exist or wasn't configured correctly") ).getDataSet(); - // Remember that RAVVs need not be in-memory in the general case. + // RAVVs need not be in-memory, but this source must remain readable through final writing. // We will sample from this RAVV to compute the PQ codebooks. // In general you don't need to have a RAVV over all the vectors to // build PQ codebooks, but you do need a "representative set". @@ -102,36 +101,21 @@ public static void main(String[] args) throws IOException { System.out.println("Building index incrementally..."); try ( - GraphIndexBuilder builder = new GraphIndexBuilder(bsp, dim, M, ef, overflow, alpha, addHierarchy); - // In DiskIntro we created the writer after generating the complete graph and closing the builder, - // but for incremental construction we will build and write in concert. - OnDiskGraphIndexWriter writer = new OnDiskGraphIndexWriter.Builder(builder.getGraph(), graphPath) - .with(new InlineVectors(dim)) - // Since we start with an empty graph, the writer will, by default, - // assume an ordinal mapping of size 0 (which is obviously incorrect). - // This is easy to rectify if you know the number of vectors beforehand, - // if not you may need to implement OrdinalMapper yourself. - .withMapper(new OrdinalMapper.IdentityMapper(ravv.size() - 1)) - .build(); + GraphIndexBuilder builder = new GraphIndexBuilder(bsp, dim, M, ef, overflow, alpha, addHierarchy) ) { + // Independent views also support sources that borrow a scratch vector. + var vectors = ravv.threadLocalSupplier(); // Graph building is best done with threads = number of physical cores // PhysicalCoreExecutor assumes hyperthreading by default, i.e. cores = vCPUs / 2 // If this is not correct, set the system property `jvector.physical_core_count` PhysicalCoreExecutor.pool().submit(() -> { IntStream.range(0, ravv.size()).parallel().forEach(ordinal -> { - VectorFloat v = ravv.getVector(ordinal); + VectorFloat v = vectors.get().getVector(ordinal); // Encode and add the vector to the working set of PQ vectors, // which allows the graph builder to access it through the BuildScoreProvder. pqVectors.encodeAndSet(ordinal, v); - // Write the feature (full-resolution vector) for a single vector instead of all at once. - try { - writer.writeFeaturesInline(ordinal, Feature.singleState(FeatureId.INLINE_VECTORS, new InlineVectors.State(v))); - } catch (IOException e) { - throw new UncheckedIOException(e); - } - builder.addGraphNode(ordinal, v); }); }).join(); @@ -140,9 +124,15 @@ public static void main(String[] args) throws IOException { // Enforces maximum degree constraint among other things. builder.cleanup(); - // No need to pass-in a feature supplier since we wrote the features incrementally - // using writer.writeInLine - writer.write(Map.of()); + // Reread this replayable source for final FP payloads; peers were resident PQ. + try (var writer = new OnDiskParallelGraphIndexWriter.Builder(builder.getGraph(), graphPath) + .withParallelWorkerThreads(PhysicalCoreExecutor.getPhysicalCoreCount()) + .with(new InlineVectors(dim)) + .withMapper(new OrdinalMapper.IdentityMapper(ravv.size() - 1)) + .build()) { + writer.write(Feature.singleStateFactory(FeatureId.INLINE_VECTORS, + ordinal -> new InlineVectors.State(vectors.get().getVector(ordinal)))); + } } // PQ codebooks and vectors also need to be saved somewhere! (we'll not concern ourselves with Fused PQ) diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/NvqExample.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/NvqExample.java index 2f0e9b1eb..661d29cb2 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/NvqExample.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/NvqExample.java @@ -17,11 +17,9 @@ package io.github.jbellis.jvector.example.tutorial; import java.io.IOException; -import java.io.UncheckedIOException; import java.nio.file.Files; import java.nio.file.Path; import java.util.List; -import java.util.Map; import java.util.stream.Collectors; import java.util.stream.IntStream; @@ -33,7 +31,7 @@ import io.github.jbellis.jvector.graph.ImmutableGraphIndex; import io.github.jbellis.jvector.graph.SearchResult; import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; -import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndexWriter; +import io.github.jbellis.jvector.graph.disk.OnDiskParallelGraphIndexWriter; import io.github.jbellis.jvector.graph.disk.OrdinalMapper; import io.github.jbellis.jvector.graph.disk.feature.Feature; import io.github.jbellis.jvector.graph.disk.feature.FeatureId; @@ -52,7 +50,8 @@ public class NvqExample { public static void main(String[] args) throws IOException { // Load a preconfigured dataset - var ds = DataSets.loadDataSet("ada002-100k").orElseThrow(() -> + String datasetName = args.length == 0 ? "ada002-100k" : args[0]; + var ds = DataSets.loadDataSet(datasetName).orElseThrow(() -> new RuntimeException("dataset not found")) .getDataSet(); var dim = ds.getDimension(); @@ -62,7 +61,7 @@ public static void main(String[] args) throws IOException { var numSubVectors = 2; // Setup NVQ parameters. - // The base vectors RAVV instance is used only for computing the global mean + // Compute NVQ statistics; the source remains readable for final encoding. var nvq = NVQuantization.compute(base, numSubVectors); // Use this method instead if you don't have all the vectors up-front but can estimate the mean // var nvq = NVQuantization.create(scaledGlobalMean, numSubVectors); @@ -88,16 +87,11 @@ public static void main(String[] args) throws IOException { Path graphPath = Files.createTempFile("jvector-nvq-graph", null); - System.out.println("Building graph in streaming mode..."); + var vectors = base.threadLocalSupplier(); + System.out.println("Building graph incrementally with resident PQ vectors..."); try ( // Create the graph builder using PQ-based scoring var builder = new GraphIndexBuilder(bsp, dim, M, ef, nOv, alpha, addHierarchy); - // Create the on-disk writer configured with NVQ feature - // This allows us to write both the graph structure and NVQ-compressed vectors - var writer = new OnDiskGraphIndexWriter.Builder(builder.getGraph(), graphPath) - .with(new NVQ(nvq)) - .withMapper(new OrdinalMapper.IdentityMapper(base.size() - 1)) - .build(); var pb = new ProgressBar("Build graph", base.size()); ) { @@ -105,22 +99,12 @@ public static void main(String[] args) throws IOException { IntStream.range(0, base.size()) .parallel() .forEach(ordinal -> { - var vec = base.getVector(ordinal); + var vec = vectors.get().getVector(ordinal); // Encode the PQ vector first, then add the graph node pqv.encodeAndSet(ordinal, vec); builder.addGraphNode(ordinal, vec); - // Encode and write NVQ vectors for later re-ranking - var nvqVec = nvq.encode(vec); - Map featureMap = Map.of( - FeatureId.NVQ_VECTORS, new NVQ.State(nvqVec) - ); - try { - writer.writeFeaturesInline(ordinal, featureMap); - } catch (IOException e) { - throw new UncheckedIOException(e); - } pb.step(); }); }).join(); @@ -129,7 +113,16 @@ public static void main(String[] args) throws IOException { // cleanup System.out.println("Cleanup..."); builder.cleanup(); - writer.write(Map.of()); + // NVQ is reranking data; encode it while writing complete records after cleanup. + System.out.println("Writing graph and NVQ vectors..."); + try (var writer = new OnDiskParallelGraphIndexWriter.Builder(builder.getGraph(), graphPath) + .withParallelWorkerThreads(PhysicalCoreExecutor.getPhysicalCoreCount()) + .with(new NVQ(nvq)) + .withMapper(new OrdinalMapper.IdentityMapper(base.size() - 1)) + .build()) { + writer.write(Feature.singleStateFactory(FeatureId.NVQ_VECTORS, + ordinal -> new NVQ.State(nvq.encode(vectors.get().getVector(ordinal))))); + } } // Search parameters diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/graph/disk/ParallelWriteExample.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/graph/disk/ParallelWriteExample.java index f3728234c..eb3c3cce9 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/graph/disk/ParallelWriteExample.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/graph/disk/ParallelWriteExample.java @@ -31,6 +31,9 @@ import io.github.jbellis.jvector.quantization.NVQuantization; import io.github.jbellis.jvector.quantization.PQVectors; import io.github.jbellis.jvector.quantization.ProductQuantization; +import io.github.jbellis.jvector.util.ExplicitThreadLocal; +import io.github.jbellis.jvector.util.ExceptionUtils; +import io.github.jbellis.jvector.util.PhysicalCoreExecutor; import java.io.IOException; import java.nio.file.Files; @@ -54,9 +57,9 @@ * writer.write(featureSuppliers); * * // Parallel: - * var writer = new OnDiskGraphIndexWriter.Builder(graph, outputPath) - * .with(inlineVectors) - * .withParallelWrites(true) // Enable parallel writes + * var writer = new OnDiskParallelGraphIndexWriter.Builder(graph, outputPath) + * .withParallelWorkerThreads(PhysicalCoreExecutor.getPhysicalCoreCount()) + * .with(inlineVectors) // Supply every inline feature for complete-record batches * .build(); * writer.write(featureSuppliers); * @@ -193,8 +196,8 @@ private static void verifyIndicesIdentical(OnDiskGraphIndex index1, OnDiskGraphI } // Compare vectors if present (only check layer 0) - if (index1.getFeatureSet().contains(FeatureId.INLINE_VECTORS) || - index1.getFeatureSet().contains(FeatureId.NVQ_VECTORS)) { + // NVQ and other feature payloads are checked by verifySerializedBytes below. + if (index1.getFeatureSet().contains(FeatureId.INLINE_VECTORS)) { int vectorsChecked = 0; int maxToCheck = Math.min(100, index1.size(0)); // Check up to 100 vectors as a sample @@ -240,9 +243,10 @@ public static void benchmarkComparison(ImmutableGraphIndex graph, var nvqFeature = new NVQ(nvq); var fusedPQFeature = new FusedPQ(graph.maxDegree(), pq); + var vectors = floatVectors.threadLocalSupplier(); // Build suppliers for inline features (NVQ only - FUSED_ADC needs neighbors) Map> inlineSuppliers = new EnumMap<>(FeatureId.class); - inlineSuppliers.put(FeatureId.NVQ_VECTORS, ordinal -> new NVQ.State(nvq.encode(floatVectors.getVector(ordinal)))); + inlineSuppliers.put(FeatureId.NVQ_VECTORS, ordinal -> new NVQ.State(nvq.encode(vectors.get().getVector(ordinal)))); // FUSED_ADC supplier needs graph view, provided at write time var identityMapper = new OrdinalMapper.IdentityMapper(floatVectors.size() - 1); @@ -254,15 +258,16 @@ public static void benchmarkComparison(ImmutableGraphIndex graph, .with(nvqFeature) .with(fusedPQFeature) .withMapper(identityMapper) - .build()) { + .build(); + var views = ExplicitThreadLocal.withInitial(graph::getView)) { - var view = graph.getView(); Map> writeSuppliers = new EnumMap<>(FeatureId.class); writeSuppliers.put(FeatureId.NVQ_VECTORS, inlineSuppliers.get(FeatureId.NVQ_VECTORS)); - writeSuppliers.put(FeatureId.FUSED_PQ, ordinal -> new FusedPQ.State(view, pqVectors, ordinal)); + writeSuppliers.put(FeatureId.FUSED_PQ, ordinal -> new FusedPQ.State(views.get(), pqVectors, ordinal)); writer.write(writeSuppliers); - view.close(); + } catch (Exception e) { + ExceptionUtils.throwIoException(e); } long sequentialTime = System.nanoTime() - sequentialStart; System.out.printf("Sequential write: %.2f ms%n", sequentialTime / 1_000_000.0); @@ -270,18 +275,20 @@ public static void benchmarkComparison(ImmutableGraphIndex graph, // Parallel write long parallelStart = System.nanoTime(); try (var writer = new OnDiskParallelGraphIndexWriter.Builder(graph, parallelPath) + .withParallelWorkerThreads(PhysicalCoreExecutor.getPhysicalCoreCount()) .with(nvqFeature) .with(fusedPQFeature) .withMapper(identityMapper) - .build()) { + .build(); + var views = ExplicitThreadLocal.withInitial(graph::getView)) { - var view = graph.getView(); Map> writeSuppliers = new EnumMap<>(FeatureId.class); writeSuppliers.put(FeatureId.NVQ_VECTORS, inlineSuppliers.get(FeatureId.NVQ_VECTORS)); - writeSuppliers.put(FeatureId.FUSED_PQ, ordinal -> new FusedPQ.State(view, pqVectors, ordinal)); + writeSuppliers.put(FeatureId.FUSED_PQ, ordinal -> new FusedPQ.State(views.get(), pqVectors, ordinal)); writer.write(writeSuppliers); - view.close(); + } catch (Exception e) { + ExceptionUtils.throwIoException(e); } long parallelTime = System.nanoTime() - parallelStart; @@ -289,6 +296,23 @@ public static void benchmarkComparison(ImmutableGraphIndex graph, System.out.printf("Speedup: %.2fx%n", (double) sequentialTime / parallelTime); } + private static void verifySerializedBytes(Path sequential, Path parallel) throws IOException { + try (var left = Files.newInputStream(sequential); var right = Files.newInputStream(parallel)) { + byte[] a = new byte[64 * 1024], b = new byte[a.length]; + long offset = 0; + while (true) { + int n = left.readNBytes(a, 0, a.length); + int m = right.readNBytes(b, 0, b.length); + if (n != m) throw new AssertionError("Serialized index lengths differ"); + if (n == 0) break; + int mismatch = Arrays.mismatch(a, 0, n, b, 0, n); + if (mismatch >= 0) throw new AssertionError("Serialized indexes differ at byte " + (offset + mismatch)); + offset += n; + } + } + System.out.println("All serialized features match byte for byte."); + } + /** * Main method to run a benchmark test of sequential vs parallel writes. * @@ -365,6 +389,7 @@ public static void main(String[] args) throws IOException { // Verify that both indices are identical verifyIndicesIdentical(sequentialIndex, parallelIndex); + verifySerializedBytes(sequentialPath, parallelPath); // Close the loaded indices sequentialIndex.close(); diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestFinalWriteVectors.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestFinalWriteVectors.java new file mode 100644 index 000000000..5c35708b1 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestFinalWriteVectors.java @@ -0,0 +1,152 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.example; + +import io.github.jbellis.jvector.graph.*; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import org.junit.Test; +import java.util.*; +import java.util.concurrent.*; +import java.util.concurrent.atomic.AtomicInteger; +import static org.junit.Assert.*; + +public class TestFinalWriteVectors { + private static ListRandomAccessVectorValues values() { + var vectors = new ArrayList>(); + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + for (int i = 0; i < 131; i++) vectors.add(vts.createFloatVector(new float[]{i, -i, i * 2})); + return new ListRandomAccessVectorValues(vectors, 3); + } + + @Test + public void rangesAreBoundedAndReleased() { + var resident = values(); + var opens = new AtomicInteger(); + var closes = new AtomicInteger(); + BatchedVectorValues source = new BatchedVectorValues() { + public int size() { return resident.size(); } + public int dimension() { return resident.dimension(); } + public boolean isValueShared() { return false; } + public RandomAccessVectorValues copy() { return this; } + public VectorFloat getVector(int ordinal) { throw new AssertionError("Point access"); } + public VectorCursor openSelection(int[] ordinals, int offset, int count) { + throw new AssertionError("Selection access"); + } + public VectorCursor openRange(int start, int end) { + assertTrue(end - start <= 64); + opens.incrementAndGet(); + var cursor = VectorAccess.openRange(resident, start, end); + return new VectorCursor() { + boolean closed; + public boolean next() { return cursor.next(); } + public int ordinal() { return cursor.ordinal(); } + public VectorFloat vector() { return cursor.vector(); } + public boolean isValueShared() { return cursor.isValueShared(); } + public void close() { + if (!closed) { closed = true; closes.incrementAndGet(); cursor.close(); } + } + }; + } + }; + try (var input = new Grid.FinalWriteVectors(source)) { + for (int i = 0; i < source.size(); i++) { + var v = input.get(i); + assertEquals(i, v.get(0), 0); + assertSame("Features for one record reuse the same input", v, input.get(i)); + } + assertEquals(3, opens.get()); + assertEquals(17, input.get(17).get(0), 0); + assertEquals(0, input.get(0).get(0), 0); + } + assertEquals(opens.get(), closes.get()); + } + + @Test(timeout = 60000) + public void ordinaryBorrowedSourcesHaveIndependentWorkerViews() throws Exception { + var resident = values(); + class Borrowed implements RandomAccessVectorValues { + final VectorFloat scratch = VectorizationProvider.getInstance().getVectorTypeSupport().createFloatVector(3); + public int size() { return resident.size(); } + public int dimension() { return 3; } + public boolean isValueShared() { return true; } + public RandomAccessVectorValues copy() { return new Borrowed(); } + public VectorFloat getVector(int ordinal) { + resident.getVectorInto(ordinal, scratch, 0); + return scratch; + } + } + var pool = Executors.newFixedThreadPool(4); + try (var input = new Grid.FinalWriteVectors(new Borrowed())) { + var tasks = new ArrayList>(); + for (int worker = 0; worker < 8; worker++) tasks.add(() -> { + for (int ordinal = 0; ordinal < resident.size(); ordinal++) { + var vector = input.get(ordinal); + Thread.yield(); + assertEquals(ordinal, vector.get(0), 0); + assertEquals(-ordinal, vector.get(1), 0); + assertSame(vector, input.get(ordinal)); + } + return null; + }); + for (var future : pool.invokeAll(tasks)) future.get(); + } finally { pool.shutdownNow(); } + } + + @Test(timeout = 60000) + public void closeReleasesOtherWorkersWhenOneCursorFails() throws Exception { + var resident = values(); + var closes = new AtomicInteger(); + BatchedVectorValues source = new BatchedVectorValues() { + public int size() { return resident.size(); } + public int dimension() { return resident.dimension(); } + public boolean isValueShared() { return false; } + public RandomAccessVectorValues copy() { return this; } + public VectorFloat getVector(int ordinal) { throw new AssertionError("Point access"); } + public VectorCursor openSelection(int[] ordinals, int offset, int count) { throw new AssertionError("Selection access"); } + public VectorCursor openRange(int start, int end) { + var cursor = VectorAccess.openRange(resident, start, end); + return new VectorCursor() { + public boolean next() { return cursor.next(); } + public int ordinal() { return cursor.ordinal(); } + public VectorFloat vector() { return cursor.vector(); } + public boolean isValueShared() { return false; } + public void close() { + closes.incrementAndGet(); cursor.close(); + if (start == 0) throw new IllegalStateException("cursor close failure"); + } + }; + } + }; + var pool = Executors.newFixedThreadPool(2); + var input = new Grid.FinalWriteVectors(source); + try { + var barrier = new CyclicBarrier(2); + for (var future : pool.invokeAll(List.>of( + () -> { input.get(0); barrier.await(10, TimeUnit.SECONDS); return null; }, + () -> { input.get(64); barrier.await(10, TimeUnit.SECONDS); return null; }))) future.get(); + var failure = assertThrows(IllegalStateException.class, input::close); + assertEquals("cursor close failure", failure.getMessage()); + assertEquals(2, closes.get()); + input.close(); + assertEquals(2, closes.get()); + } finally { + pool.shutdownNow(); pool.awaitTermination(10, TimeUnit.SECONDS); + input.close(); + } + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestGridDeferredFeatureParallelism.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestGridDeferredFeatureParallelism.java new file mode 100644 index 000000000..7a5bbd3b7 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestGridDeferredFeatureParallelism.java @@ -0,0 +1,220 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.example; + +import io.github.jbellis.jvector.TestUtil; +import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet; +import io.github.jbellis.jvector.example.reporting.RunArtifacts; +import io.github.jbellis.jvector.example.util.CompressorParameters; +import io.github.jbellis.jvector.example.util.OnDiskGraphIndexCache; +import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues; +import io.github.jbellis.jvector.graph.RandomAccessVectorValues; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; +import io.github.jbellis.jvector.management.GraphIndexBuilderConfig; +import io.github.jbellis.jvector.util.PhysicalCoreExecutor; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.VectorUtil; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import org.junit.Assume; +import org.junit.Test; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.EnumSet; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.Assert.assertTrue; + +/// Repro for PR #747: with the default writer, Grid serializes all deferred feature work on +/// one thread. +/// +/// ## The bug +/// +/// Before #747, Grid encoded NVQ (and copied full-precision vectors) for each node during +/// insertion, on the parallel insertion threads, and wrote them with `writeInline`. #747 moves +/// that work into the final `writer.write(suppliers)` call for every Grid run. +/// +/// Only the parallel writer calls final suppliers from several threads. It is selected only when +/// `GraphIndexBuilderConfig.isParallelBuild()` is true. The default is `false`, and BenchYAML +/// offers no way to change it, so a plain BenchYAML run (including CI) uses the sequential +/// `OnDiskGraphIndexWriter`. Its `writeL0Records` calls each supplier in turn, on the calling +/// thread. Every vector read and NVQ encode for the whole dataset then runs on one thread after +/// construction. +/// +/// On Cohere-1M (preload, native, 48 workers), build+write went from 74–76 s at `4dcfe917` +/// to 100–104 s at `ec42a94b` with the default writer, about 35% slower. With the parallel writer +/// enabled, #747 is faster than before (78 s against 88–92 s). But only the PR's measurements enable it. +/// +/// ## The test +/// +/// Records the thread of every base-vector read made during final serialization, meaning with a +/// frame from the graph-writer package (`io.github.jbellis.jvector.graph.disk`) on the reading +/// thread's stack. That covers the sequential writer's `writeL0Records` and the parallel writer's +/// `NodeRecordTask`s. No other read in a Grid run goes through that package. The test runs Grid +/// twice on a resident dataset. +/// +/// 1. **Control**, with the parallel writer (`parallelBuild = true`). Grid defers features with +/// this writer, so the final write must read base vectors, from more than one writer thread. +/// This proves the detector sees the final write. If the writers move package or the final +/// write stops reading through the dataset's vectors, the control fails instead of the check +/// below passing vacuously. +/// 2. **Requirement**, with BenchYAML's default writer selection (`parallelBuild = false`). +/// Final-write input must not be confined to one thread. Either no vectors are read during +/// the final write (features were produced during insertion), or they are read from more than +/// one thread. +/// +/// At `ec42a94b` the control passes, and every read in the default run's final write comes from +/// one thread. The control needs at least two writer threads, so it's skipped on a machine with +/// one physical core. +public class TestGridDeferredFeatureParallelism { + private static final int SIZE = 4096, DIMENSIONS = 17; + + /// Package of the graph writers: the sequential writer's `writeL0Records`, and the parallel + /// writer's `NodeRecordTask`s, which run on worker threads whose stacks lack `writeL0Records`. + private static final String WRITER_PACKAGE = "io.github.jbellis.jvector.graph.disk."; + + /// Resident vectors that record which threads read them during final serialization, meaning + /// with a graph-writer frame on the reading thread's stack. + private static final class RecordingVectors implements RandomAccessVectorValues { + final RandomAccessVectorValues resident; + final Set writeThreads = ConcurrentHashMap.newKeySet(); + final AtomicInteger writeReads = new AtomicInteger(); + + RecordingVectors(RandomAccessVectorValues resident) { this.resident = resident; } + + @Override public int size() { return resident.size(); } + @Override public int dimension() { return resident.dimension(); } + @Override public boolean isValueShared() { return false; } + @Override public RandomAccessVectorValues copy() { return this; } + + @Override public VectorFloat getVector(int ordinal) { + boolean inFinalWrite = StackWalker.getInstance().walk(frames -> + frames.anyMatch(f -> f.getClassName().startsWith(WRITER_PACKAGE))); + if (inFinalWrite) { + writeReads.incrementAndGet(); + writeThreads.add(Thread.currentThread().getName()); + } + return resident.getVector(ordinal); + } + } + + /// Run Grid once on `vectors` with the given writer selection, restoring the selection and + /// removing the run's work directory and cached codebook afterwards. + private static void runGrid(List> list, RecordingVectors vectors, boolean parallel) + throws Exception { + String name = "deferred-input-" + UUID.randomUUID(); + DataSet ds = new DataSet() { + public int getDimension() { return DIMENSIONS; } + public String getName() { return name; } + public VectorSimilarityFunction getSimilarityFunction() { return VectorSimilarityFunction.DOT_PRODUCT; } + public RandomAccessVectorValues getBaseRavv() { return vectors; } + public List> getBaseVectors() { return list; } + public List> getQueryVectors() { return List.of(list.get(0)); } + public List> getGroundTruth() { return List.of(List.of(0)); } + }; + var config = GraphIndexBuilderConfig.getInstance(); + boolean previous = config.isParallelBuild(); + var work = Files.createTempDirectory("grid-deferred-input"); + var pqParameters = new CompressorParameters.PQParameters(4, 16, false, -1f); + try { + config.setParallelBuild(parallel); + Grid.runOneGraph(OnDiskGraphIndexCache.initialize(work.resolve("saved")), + List.of(EnumSet.of(FeatureId.NVQ_VECTORS)), 8, 32, 1.2f, true, false, + ignored -> pqParameters, List.of(), Map.of(10, List.of(1.0)), List.of(false), + RunArtifacts.disabled(), ds, work); + } finally { + config.setParallelBuild(previous); + TestUtil.deleteQuietly(work); + // Grid caches codebooks by dataset name in the working directory; remove this run's. + Files.deleteIfExists(Path.of("pq_cache").resolve(pqParameters.idStringFor(ds))); + } + } + + @Test(timeout = 120_000) + public void defaultWriterDoesNotConfineDeferredFeatureInputToOneThread() throws Exception { + Assume.assumeTrue("the control needs at least two writer threads", + PhysicalCoreExecutor.getPhysicalCoreCount() > 1); + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + var random = new Random(42); + var list = new ArrayList>(); + for (int n = 0; n < SIZE; n++) { + var v = vts.createFloatVector(DIMENSIONS); + for (int d = 0; d < DIMENSIONS; d++) v.set(d, random.nextFloat() - 0.5f); + VectorUtil.l2normalize(v); + list.add(v); + } + var resident = new ListRandomAccessVectorValues(list, DIMENSIONS); + + var control = new RecordingVectors(resident); + runGrid(list, control, true); + assertTrue("control: the parallel writer's final write read " + control.writeReads.get() + + " base vectors on threads " + control.writeThreads + + "; the detector no longer sees the final write", + control.writeReads.get() > 0 && control.writeThreads.size() > 1); + + var vectors = new RecordingVectors(resident); + runGrid(list, vectors, false); // Explicit sequential selection; retain efficient insertion-time encoding. + int reads = vectors.writeReads.get(); + assertTrue("final serialization read " + reads + " base vectors, all on one thread " + + vectors.writeThreads + ", with the library's default (sequential) writer", + reads == 0 || vectors.writeThreads.size() > 1); + } + + @Test + public void benchmarkWriterDefaultAndOverrideSelectActualWriterTypes() throws Exception { + var settings = GraphIndexBuilderConfig.getInstance(); + boolean previous = settings.isParallelBuild(); + String property = System.getProperty("jvector.parallel_writer"); + var path = Files.createTempFile("benchmark-writer-selection", ".idx"); + var support = VectorizationProvider.getInstance().getVectorTypeSupport(); + var values = new ListRandomAccessVectorValues(List.of( + support.createFloatVector(new float[] {1f, 0f}), + support.createFloatVector(new float[] {0f, 1f})), 2); + try (var builder = new io.github.jbellis.jvector.graph.GraphIndexBuilder(values, + VectorSimilarityFunction.DOT_PRODUCT, 4, 16, 1.2f, 1.2f, true)) { + var graph = TestUtil.buildSequentially(builder, values); + System.clearProperty("jvector.parallel_writer"); + Grid.configureWriterForBenchmark(); + try (var writer = new io.github.jbellis.jvector.graph.disk.RandomAccessOnDiskGraphIndexWriter.Builder(graph, path) + .with(new io.github.jbellis.jvector.graph.disk.feature.InlineVectors(2)).build()) { + assertTrue(writer instanceof io.github.jbellis.jvector.graph.disk.OnDiskParallelGraphIndexWriter); + } + System.setProperty("jvector.parallel_writer", "false"); + Grid.configureWriterForBenchmark(); + try (var writer = new io.github.jbellis.jvector.graph.disk.RandomAccessOnDiskGraphIndexWriter.Builder(graph, path) + .with(new io.github.jbellis.jvector.graph.disk.feature.InlineVectors(2)).build()) { + assertTrue(writer instanceof io.github.jbellis.jvector.graph.disk.OnDiskGraphIndexWriter); + } + System.setProperty("jvector.parallel_writer", "invalid"); + org.junit.Assert.assertThrows(IllegalArgumentException.class, Grid::configureWriterForBenchmark); + org.junit.Assert.assertFalse(settings.isParallelBuild()); + } finally { + settings.setParallelBuild(previous); + if (property == null) System.clearProperty("jvector.parallel_writer"); + else System.setProperty("jvector.parallel_writer", property); + Files.deleteIfExists(path); + } + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestBufferedGraphCompatibility.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestBufferedGraphCompatibility.java index 73b761d79..fe171b872 100644 --- a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestBufferedGraphCompatibility.java +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestBufferedGraphCompatibility.java @@ -35,10 +35,30 @@ import static org.junit.Assert.*; public class TestBufferedGraphCompatibility { + /// Every level's adjacency, one "level node: neighbors" line per node in node order. + private static List adjacency(ImmutableGraphIndex graph) throws Exception { + var lines = new ArrayList(); + try (var view = graph.getView()) { + for (int level = 0; level <= graph.getMaxLevel(); level++) { + var nodes = new ArrayList(); + for (var it = graph.getNodes(level); it.hasNext(); ) nodes.add(it.nextInt()); + Collections.sort(nodes); + for (int node : nodes) { + var neighbors = new StringBuilder(); + for (var it = view.getNeighborsIterator(level, node); it.hasNext(); ) neighbors.append(it.nextInt()).append(' '); + lines.add(level + " " + node + ": " + neighbors); + } + } + } + return lines; + } + @Test(timeout = 180000) public void exactPqAndBinaryBuildsWithAllMetricsAndRefinementModes() throws Exception { var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); var pool = new ForkJoinPool(3); + // One worker makes insertion order, level draws and cleanup order reproducible. + var single = new ForkJoinPool(1); try { for (int dimension : new int[] {17, 768}) { int count = 257; @@ -68,10 +88,20 @@ public void exactPqAndBinaryBuildsWithAllMetricsAndRefinementModes() throws Exce var bsp = mode == 0 ? BuildScoreProvider.randomAccessScoreProvider(source, metric) : mode == 1 ? BuildScoreProvider.pqBuildScoreProvider(metric, pqv) : BuildScoreProvider.bqBuildScoreProvider(bqv); + var residentBsp = mode == 0 ? BuildScoreProvider.randomAccessScoreProvider(resident, metric) : bsp; + List residentAdjacency; + try (var reference = new GraphIndexBuilder(residentBsp, dimension, 8, 64, 1.2f, 1.2f, + true, refine, single, single)) { + residentAdjacency = adjacency(reference.build(resident)); + } try (var builder = new GraphIndexBuilder(bsp, dimension, 8, 64, 1.2f, 1.2f, - true, refine, pool, pool)) { + true, refine, single, single)) { var graph = builder.build(source); assertEquals(count, graph.size(0)); + // Deterministic builds: the file-backed source must give the resident graph. + assertEquals("dimension=" + dimension + ", build=" + (mode == 0 ? "exact" : mode == 1 ? "PQ" : "BQ") + + ", metric=" + metric + ", refine=" + refine, + residentAdjacency, adjacency(graph)); // A file-backed shared source must not overwrite insertion/diversity queries. var result = GraphSearcher.search(vectors.get(5), 10, resident, metric, graph, Bits.ALL); assertTrue("dimension=" + dimension + ", build=" + (mode == 0 ? "exact" : mode == 1 ? "PQ" : "BQ") @@ -123,6 +153,6 @@ public void exactPqAndBinaryBuildsWithAllMetricsAndRefinementModes() throws Exce } } finally { Files.delete(path); } } - } finally { pool.shutdown(); } + } finally { pool.shutdown(); single.shutdown(); } } } diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestConcurrentBufferedGraphConstruction.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestConcurrentBufferedGraphConstruction.java new file mode 100644 index 000000000..84fefcae4 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestConcurrentBufferedGraphConstruction.java @@ -0,0 +1,145 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph; + +import io.github.jbellis.jvector.disk.FvecFileVectorValues; +import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider; +import io.github.jbellis.jvector.quantization.*; +import io.github.jbellis.jvector.vector.*; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import org.junit.Test; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.file.Files; +import java.util.ArrayList; +import java.util.Random; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ForkJoinPool; +import java.util.concurrent.ForkJoinTask; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.Assert.*; + +public class TestConcurrentBufferedGraphConstruction { + private static final class Reads { + final CountDownLatch rendezvous = new CountDownLatch(2); + final AtomicInteger tickets = new AtomicInteger(); + final AtomicInteger active = new AtomicInteger(); + final AtomicInteger maximum = new AtomicInteger(); + final AtomicInteger copies = new AtomicInteger(); + } + + private static final class CheckedSource implements RandomAccessVectorValues { + final RandomAccessVectorValues source, resident; + final Reads reads; + + CheckedSource(RandomAccessVectorValues source, RandomAccessVectorValues resident, Reads reads) { + this.source = source; + this.resident = resident; + this.reads = reads; + } + public int size() { return source.size(); } + public int dimension() { return source.dimension(); } + public boolean isValueShared() { return source.isValueShared(); } + public RandomAccessVectorValues copy() { + reads.copies.incrementAndGet(); + return new CheckedSource(source.copy(), resident, reads); + } + public VectorFloat getVector(int ordinal) { + reads.maximum.accumulateAndGet(reads.active.incrementAndGet(), Math::max); + try { + var vector = source.getVector(ordinal); + // Force overlapping borrowed reads; an accidentally shared scratch buffer is observable. + if (ForkJoinTask.inForkJoinPool() && reads.tickets.getAndIncrement() < 2) { + reads.rendezvous.countDown(); + try { assertTrue("construction reads did not overlap", reads.rendezvous.await(10, TimeUnit.SECONDS)); } + catch (InterruptedException e) { Thread.currentThread().interrupt(); throw new AssertionError(e); } + } + var expected = resident.getVector(ordinal); + for (int d = 0; d < dimension(); d++) + assertEquals("borrowed vector changed at ordinal=" + ordinal + ", dimension=" + d, + Float.floatToIntBits(expected.get(d)), Float.floatToIntBits(vector.get(d))); + return vector; + } finally { + reads.active.decrementAndGet(); + } + } + } + + @Test(timeout = 180000) + public void exactPqAndBinaryBuildsReadIndependentBorrowedVectorsConcurrently() throws Exception { + var support = VectorizationProvider.getInstance().getVectorTypeSupport(); + var pool = new ForkJoinPool(3); + try { + for (int dimension : new int[] {17, 768}) { + int count = 257; + var vectors = new ArrayList>(); + var bytes = ByteBuffer.allocate(count * (dimension + 1) * 4).order(ByteOrder.LITTLE_ENDIAN); + var random = new Random(42); + for (int n = 0; n < count; n++) { + var vector = support.createFloatVector(dimension); + for (int d = 0; d < dimension; d++) vector.set(d, random.nextFloat() * 2 - 1); + VectorUtil.l2normalize(vector); + vectors.add(vector); + bytes.putInt(dimension); + for (int d = 0; d < dimension; d++) bytes.putFloat(vector.get(d)); + } + var path = Files.createTempFile("concurrent-graph-input", ".fvecs"); + Files.write(path, bytes.array()); + var resident = new ListRandomAccessVectorValues(vectors, dimension); + try (var file = FvecFileVectorValues.open(path)) { + var pq = ProductQuantization.compute(resident, Math.max(1, dimension / 8), 16, false, + KMeansPlusPlusClusterer.UNWEIGHTED, pool, pool); + var pqv = pq.encodeAll(file, pool); + var bqv = (BQVectors) new BinaryQuantization(dimension).encodeAll(file, pool); + for (var metric : VectorSimilarityFunction.values()) { + for (int mode = 0; mode < 3; mode++) { + var reads = new Reads(); + var source = new CheckedSource(file, resident, reads); + var scores = mode == 0 ? BuildScoreProvider.randomAccessScoreProvider(source, metric) + : mode == 1 ? BuildScoreProvider.pqBuildScoreProvider(metric, pqv) + : BuildScoreProvider.bqBuildScoreProvider(bqv); + try (var builder = new GraphIndexBuilder(scores, dimension, 8, 64, 1.2f, 1.2f, + true, true, pool, pool)) { + var graph = builder.build(source); + assertEquals(count, graph.size(0)); + assertTrue("source copies were not used", reads.copies.get() >= 2); + assertTrue("construction was not concurrent", reads.maximum.get() >= 2); + try (var view = graph.getView()) { + for (int level = 0; level <= graph.getMaxLevel(); level++) { + for (var nodes = graph.getNodes(level); nodes.hasNext(); ) { + int node = nodes.nextInt(); + var neighbors = view.getNeighborsIterator(level, node); + assertTrue(neighbors.size() <= graph.getDegree(level)); + while (neighbors.hasNext()) { + int neighbor = neighbors.nextInt(); + assertTrue(neighbor >= 0 && neighbor < count); + assertNotEquals(node, neighbor); + } + } + } + } + } + } + } + } finally { Files.deleteIfExists(path); } + } + } finally { pool.shutdown(); } + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestDeferredGraphWriter.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestDeferredGraphWriter.java new file mode 100644 index 000000000..7a6071830 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestDeferredGraphWriter.java @@ -0,0 +1,197 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.github.jbellis.jvector.graph.disk; + +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; +import io.github.jbellis.jvector.LuceneTestCase; +import io.github.jbellis.jvector.TestUtil; +import io.github.jbellis.jvector.disk.BufferedRandomAccessWriter; +import io.github.jbellis.jvector.graph.GraphIndexBuilder; +import io.github.jbellis.jvector.graph.ImmutableGraphIndex; +import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues; +import io.github.jbellis.jvector.graph.disk.feature.*; +import io.github.jbellis.jvector.management.GraphIndexBuilderConfig; +import io.github.jbellis.jvector.quantization.NVQuantization; +import io.github.jbellis.jvector.quantization.ProductQuantization; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import org.junit.Test; + +import java.nio.file.Files; +import java.util.*; +import java.util.concurrent.*; +import java.util.function.IntFunction; + +import static org.junit.Assert.*; + +@ThreadLeakScope(ThreadLeakScope.Scope.NONE) +public class TestDeferredGraphWriter extends LuceneTestCase { + private ListRandomAccessVectorValues values() { + var vectors = new ArrayList>(); + var support = VectorizationProvider.getInstance().getVectorTypeSupport(); + var random = new Random(42); + for (int n = 0; n < 259; n++) { + var vector = support.createFloatVector(17); + for (int d = 0; d < 17; d++) vector.set(d, random.nextFloat() - 0.5f); + vectors.add(vector); + } + return new ListRandomAccessVectorValues(vectors, 17); + } + + @Test + public void completeAndPrewrittenFeaturesMatchSequentialV6() throws Exception { + var directory = Files.createTempDirectory("deferred-writer"); + var values = values(); + try (var builder = new GraphIndexBuilder(values, VectorSimilarityFunction.DOT_PRODUCT, + 8, 32, 1.2f, 1.2f, true)) { + var graph = TestUtil.buildSequentially(builder, values); + var pq = ProductQuantization.compute(values, 3, 256, false); + var codes = pq.encodeAll(values, ForkJoinPool.commonPool()); + var nvq = NVQuantization.compute(values, 2); + var frozenView = graph.getView(); + try { + for (boolean holes : new boolean[]{false, true}) { + OrdinalMapper mapper = new OrdinalMapper() { + public int maxOrdinal() { return holes ? values.size() * 2 - 1 : values.size() - 1; } + public int oldToNew(int n) { return (values.size() - 1 - n) * (holes ? 2 : 1); } + public int newToOld(int n) { return holes && n % 2 != 0 ? OMITTED : values.size() - 1 - n / (holes ? 2 : 1); } + }; + for (boolean quantized : new boolean[]{false, true}) { + Feature vectors = quantized ? new NVQ(nvq) : new InlineVectors(values.dimension()); + var fused = new FusedPQ(graph.maxDegree(), pq); + Map> suppliers = new EnumMap<>(FeatureId.class); + suppliers.put(vectors.id(), n -> quantized ? new NVQ.State(nvq.encode(values.getVector(n))) + : new InlineVectors.State(values.getVector(n))); + suppliers.put(fused.id(), n -> new FusedPQ.State(frozenView, codes, n)); + byte[] prefix = new byte[83]; + Arrays.fill(prefix, (byte) 0x5a); + var sequential = directory.resolve("sequential"); + Files.write(sequential, prefix); + try (var writer = new OnDiskGraphIndexWriter.Builder(graph, sequential) + .withStartOffset(prefix.length).withVersion(6).withMapper(mapper) + .with(vectors).with(fused).build()) { + writer.write(suppliers); + } + byte[] expected = Files.readAllBytes(sequential); + for (int batch : new int[]{1, 7, 64, 128}) { + for (boolean direct : new boolean[]{false, true}) { + for (boolean prewritten : new boolean[]{false, true}) { + var parallel = directory.resolve("parallel"); + Files.write(parallel, prefix); + try (var writer = new OnDiskParallelGraphIndexWriter.Builder(graph, parallel) + .withParallelWorkerThreads(3).withParallelBatchSize(batch) + .withParallelDirectBuffers(direct).withStartOffset(prefix.length) + .withVersion(6).withMapper(mapper).with(vectors).with(fused).build()) { + if (prewritten) { + for (int ordinal = 0; ordinal < values.size(); ordinal++) { + writer.writeFeaturesInline(mapper.oldToNew(ordinal), + Map.of(vectors.id(), suppliers.get(vectors.id()).apply(ordinal))); + } + writer.write(Map.of(fused.id(), suppliers.get(fused.id()))); + } else writer.write(suppliers); + } + assertArrayEquals("holes=" + holes + " nvq=" + quantized + " batch=" + batch + + " direct=" + direct + " prewritten=" + prewritten, + expected, Files.readAllBytes(parallel)); + } + } + } + } + } + } finally { frozenView.close(); } + } finally { TestUtil.deleteQuietly(directory); } + } + + @Test(timeout = 60000) + public void failedWriteJoinsWorkersAndPreservesExternalExecutor() throws Exception { + var directory = Files.createTempDirectory("failed-deferred-writer"); + var values = values(); + var workers = Executors.newFixedThreadPool(2); + var caller = Executors.newSingleThreadExecutor(); + var release = new CountDownLatch(1); + try (var builder = new GraphIndexBuilder(values, VectorSimilarityFunction.DOT_PRODUCT, + 8, 32, 1.2f, 1.2f, false)) { + var graph = TestUtil.buildSequentially(builder, values); + var file = directory.resolve("records"); + Files.createFile(file); + var blocked = new CountDownLatch(1); + var failed = new CountDownLatch(1); + try (var out = new BufferedRandomAccessWriter(file); + var writer = new ParallelGraphWriter(out, graph, List.of(new InlineVectors(17)), + new ParallelGraphWriter.Config(2, false, 1, 7), file, workers)) { + var writing = caller.submit(() -> { + writer.writeL0Records(new OrdinalMapper.IdentityMapper(values.size() - 1), + List.of(new InlineVectors(17)), Feature.singleStateFactory(FeatureId.INLINE_VECTORS, n -> { + if (n == 0) { + try { if (!blocked.await(10, TimeUnit.SECONDS)) throw new AssertionError("Worker did not start"); } + catch (InterruptedException e) { throw new AssertionError(e); } + failed.countDown(); + throw new IllegalStateException("supplier failure"); + } + blocked.countDown(); + try { if (!release.await(10, TimeUnit.SECONDS)) throw new AssertionError("Writer cleanup timeout"); } + catch (InterruptedException e) { throw new AssertionError(e); } + return new InlineVectors.State(values.getVector(n)); + }), 0); + return null; + }); + assertTrue(failed.await(10, TimeUnit.SECONDS)); + assertTrue(blocked.await(10, TimeUnit.SECONDS)); + assertThrows(TimeoutException.class, () -> writing.get(100, TimeUnit.MILLISECONDS)); + release.countDown(); + var failure = assertThrows(ExecutionException.class, () -> writing.get(10, TimeUnit.SECONDS)); + assertTrue(failure.getCause() instanceof IllegalStateException); + } + assertEquals(42, workers.submit(() -> 42).get().intValue()); + } finally { + release.countDown(); + caller.shutdownNow(); workers.shutdownNow(); + caller.awaitTermination(10, TimeUnit.SECONDS); workers.awaitTermination(10, TimeUnit.SECONDS); + TestUtil.deleteQuietly(directory); + } + } + + @Test + public void batchSizeValidationAndUnifiedBuilderSelection() throws Exception { + var directory = Files.createTempDirectory("writer-batch-option"); + var values = values(); + var settings = GraphIndexBuilderConfig.getInstance(); + boolean wasParallel = settings.isParallelBuild(); + try (var graphBuilder = new GraphIndexBuilder(values, VectorSimilarityFunction.DOT_PRODUCT, + 8, 32, 1.2f, 1.2f, false)) { + var graph = TestUtil.buildSequentially(graphBuilder, values); + assertEquals(64, ParallelGraphWriter.Config.defaultConfig().recordsPerBatch); + assertThrows(IllegalArgumentException.class, () -> new ParallelGraphWriter.Config(1, false, 1, 0)); + for (boolean parallel : new boolean[]{false, true}) { + settings.setParallelBuild(parallel); + var builder = new RandomAccessOnDiskGraphIndexWriter.Builder(graph, directory.resolve("unified")); + assertThrows(IllegalArgumentException.class, () -> builder.withParallelBatchSize(-1)); + try (var writer = builder.withParallelBatchSize(7).withParallelWorkerThreads(2) + .with(new InlineVectors(17)).build()) { + assertEquals(parallel, writer instanceof OnDiskParallelGraphIndexWriter); + writer.write(Feature.singleStateFactory(FeatureId.INLINE_VECTORS, + n -> new InlineVectors.State(values.getVector(n)))); + } + } + var direct = new OnDiskParallelGraphIndexWriter.Builder(graph, directory.resolve("direct")); + assertThrows(IllegalArgumentException.class, () -> direct.withParallelBatchSize(0)); + try (var writer = direct.with(new InlineVectors(17)).build()) { } + } finally { + settings.setParallelBuild(wasParallel); + TestUtil.deleteQuietly(directory); + } + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestNodeRecordTask.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestNodeRecordTask.java new file mode 100644 index 000000000..12b2090c4 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestNodeRecordTask.java @@ -0,0 +1,130 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph.disk; + +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; +import io.github.jbellis.jvector.LuceneTestCase; +import io.github.jbellis.jvector.TestUtil; +import io.github.jbellis.jvector.graph.GraphIndexBuilder; +import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues; +import io.github.jbellis.jvector.graph.disk.feature.Feature; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; +import io.github.jbellis.jvector.graph.disk.feature.InlineVectors; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import org.junit.Test; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.channels.AsynchronousFileChannel; +import java.nio.file.Files; +import java.nio.file.StandardOpenOption; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +@ThreadLeakScope(ThreadLeakScope.Scope.NONE) +public class TestNodeRecordTask extends LuceneTestCase { + @Test + public void testCompleteRecordBatches() throws Exception { + var root = Files.createTempDirectory("node-record-batches"); + var values = new ArrayList>(); + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + for (int n = 0; n < 131; n++) { + var vector = vts.createFloatVector(17); + for (int d = 0; d < 17; d++) { + vector.set(d, (float) Math.sin(n * 17 + d)); + } + values.add(vector); + } + var ravv = new ListRandomAccessVectorValues(values, 17); + try (var builder = new GraphIndexBuilder(ravv, VectorSimilarityFunction.DOT_PRODUCT, + 4, 10, 1.2f, 1.2f, false)) { + var graph = TestUtil.buildSequentially(builder, ravv); + var feature = new InlineVectors(17); + for (boolean holes : new boolean[]{false, true}) { + int count = holes ? 262 : 131; + OrdinalMapper mapper = holes ? new OrdinalMapper() { + public int maxOrdinal() { return 261; } + public int oldToNew(int n) { return 2 * n; } + public int newToOld(int n) { return n % 2 == 0 ? n / 2 : OMITTED; } + } : new OrdinalMapper.IdentityMapper(130); + int payload = Integer.BYTES * (2 + graph.getDegree(0)) + feature.featureSize(); + for (int stride : new int[]{payload}) { + for (boolean direct : new boolean[]{false, true}) { + var target = root.resolve("records-" + holes + "-" + stride + "-" + direct); + // Preserve bytes before L0 and force full batches plus a partial tail. + final int baseOffset = 83; + byte[] prefix = new byte[baseOffset]; + Arrays.fill(prefix, (byte) 0x5a); + Files.write(target, prefix); + var buffer = direct ? ByteBuffer.allocateDirect(stride * 64) + : ByteBuffer.allocate(stride * 64); + buffer.order(ByteOrder.BIG_ENDIAN); + try (var view = graph.getView(); + var channel = AsynchronousFileChannel.open(target, StandardOpenOption.WRITE)) { + var task = new NodeRecordTask(0, count, mapper, graph, view, List.of(feature), + Feature.singleStateFactory(FeatureId.INLINE_VECTORS, + n -> new InlineVectors.State(values.get(n))), + stride, baseOffset, channel, buffer); + var padded = new NodeRecordTask(0, count, mapper, graph, view, List.of(feature), + Feature.singleStateFactory(FeatureId.INLINE_VECTORS, + n -> new InlineVectors.State(values.get(n))), + 4096, baseOffset, channel, buffer); + org.junit.Assert.assertThrows(IllegalArgumentException.class, padded::call); + task.call(); + } + byte[] actual = Files.readAllBytes(target); + assertEquals(baseOffset + count * stride, actual.length); + assertArrayEquals(prefix, Arrays.copyOf(actual, baseOffset)); + try (var view = graph.getView()) { + for (int n = 0; n < count; n++) { + var expected = ByteBuffer.allocate(stride).order(ByteOrder.BIG_ENDIAN); + expected.putInt(mapper.newToOld(n) == OrdinalMapper.OMITTED ? -1 : n); + int old = mapper.newToOld(n); + if (old == OrdinalMapper.OMITTED) { + expected.position(expected.position() + feature.featureSize()); + expected.putInt(0); + for (int d = 0; d < graph.getDegree(0); d++) expected.putInt(-1); + } else { + for (int d = 0; d < 17; d++) expected.putFloat(values.get(old).get(d)); + var neighbors = view.getNeighborsIterator(0, old); + expected.putInt(neighbors.size()); + int d = 0; + while (neighbors.hasNext()) { + expected.putInt(mapper.oldToNew(neighbors.nextInt())); + d++; + } + for (; d < graph.getDegree(0); d++) expected.putInt(-1); + } + assertEquals("Record " + n + " holes=" + holes + " direct=" + direct, -1, + Arrays.mismatch(actual, baseOffset + n * stride, + baseOffset + (n + 1) * stride, expected.array(), 0, stride)); + } + } + } + } + } + } finally { + TestUtil.deleteQuietly(root); + } + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestParallelWriterInterrupt.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestParallelWriterInterrupt.java new file mode 100644 index 000000000..e79a99aa7 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestParallelWriterInterrupt.java @@ -0,0 +1,208 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph.disk; + +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; +import io.github.jbellis.jvector.LuceneTestCase; +import io.github.jbellis.jvector.TestUtil; +import io.github.jbellis.jvector.graph.GraphIndexBuilder; +import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues; +import io.github.jbellis.jvector.graph.disk.feature.Feature; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; +import io.github.jbellis.jvector.graph.disk.feature.InlineVectors; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import org.junit.Test; + +import java.nio.file.Files; +import java.util.ArrayList; +import java.util.Random; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.Executors; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +/// Repro for PR #747: an interrupted parallel write keeps running until every record is written. +/// +/// ## The bug +/// +/// `ParallelGraphWriter.writeL0Records` submits one task per ordinal range, then +/// `awaitTasks` waits on each future in turn. When the waiting thread is interrupted, +/// `awaitTasks` records the interruption and calls `future.get()` again. When a task fails, it +/// keeps the failure and moves on to the next future. Either way it waits for every task. No +/// task is cancelled, so not-yet-started ranges still run and every running supplier runs to +/// the end of its range. +/// +/// Joining before the channel and suppliers are released is right. But nothing stops the work +/// being joined. So an interrupted or already-failed write runs the rest of the final +/// serialization, encoding every remaining record, and only then reports the interruption or +/// failure. For a 10M-node graph that's the whole final write. Before #747, `writeL0Records` +/// threw on interrupt immediately, without joining. +/// +/// ## The test +/// +/// Writes a 6,000-node graph with two writer workers and an `InlineVectors` supplier that takes +/// 1 ms per record, so an uninterrupted write takes about 3 s. After the first 50 records it +/// interrupts the writing thread, then expects the write to fail within 1 s and to have supplied +/// fewer than all records. At `ec42a94b` the write runs to completion, 3.2 s after the +/// interrupt, and supplies every record. +@ThreadLeakScope(ThreadLeakScope.Scope.NONE) +public class TestParallelWriterInterrupt extends LuceneTestCase { + private static final int SIZE = 6000, DIMENSIONS = 17, WORKERS = 2; + private static final long MAX_STOP_MILLIS = 1000; + + @Test(timeout = 120_000) + public void interruptedWriteStopsPromptly() throws Exception { + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + var random = new Random(42); + var list = new ArrayList>(); + for (int n = 0; n < SIZE; n++) { + var v = vts.createFloatVector(DIMENSIONS); + for (int d = 0; d < DIMENSIONS; d++) v.set(d, random.nextFloat() - 0.5f); + list.add(v); + } + var values = new ListRandomAccessVectorValues(list, DIMENSIONS); + var path = Files.createTempFile("parallel-writer-interrupt", ".idx"); + try (var builder = new GraphIndexBuilder(values, VectorSimilarityFunction.DOT_PRODUCT, + 8, 32, 1.2f, 1.2f, true)) { + var graph = TestUtil.buildSequentially(builder, values); + var supplied = new AtomicInteger(); + var started = new CountDownLatch(50); + var failure = new AtomicReference(); + var interruptedAt = new AtomicReference(); + var finishedAt = new AtomicReference(); + + var writing = new Thread(() -> { + try (var writer = new OnDiskParallelGraphIndexWriter.Builder(graph, path) + .withParallelWorkerThreads(WORKERS) + .with(new InlineVectors(DIMENSIONS)) + .build()) { + writer.write(Feature.singleStateFactory(FeatureId.INLINE_VECTORS, ordinal -> { + try { + Thread.sleep(1); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException("supplier interrupted", e); + } + supplied.incrementAndGet(); + started.countDown(); + return new InlineVectors.State(values.getVector(ordinal)); + })); + } catch (Throwable t) { + failure.set(t); + } finally { + finishedAt.set(System.nanoTime()); + } + }, "parallel-writer-under-test"); + writing.start(); + assertTrue("writer did not start supplying records", started.await(30, TimeUnit.SECONDS)); + interruptedAt.set(System.nanoTime()); + writing.interrupt(); + writing.join(60_000); + + assertNotNull("write did not finish", finishedAt.get()); + long stopMillis = TimeUnit.NANOSECONDS.toMillis(finishedAt.get() - interruptedAt.get()); + assertTrue("interrupted write took " + stopMillis + " ms to stop and supplied " + + supplied.get() + " of " + SIZE + " records", + stopMillis <= MAX_STOP_MILLIS && supplied.get() < SIZE); + assertNotNull("interrupted write reported success", failure.get()); + } finally { + Files.deleteIfExists(path); + } + } + + @Test(timeout = 60_000) + public void laterWorkerFailureStopsQueuedRangesBeforeEarlierTaskFinishes() throws Exception { + var vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + var list = new ArrayList>(); + for (int n = 0; n < 512; n++) list.add(vts.createFloatVector(new float[] {n + 1f, 1f})); + var values = new ListRandomAccessVectorValues(list, 2); + int tasks = Math.min(Runtime.getRuntime().availableProcessors(), values.size()); + // An earlier blocked task and a later failed task require at least two ranges. + org.junit.Assume.assumeTrue(tasks > 1); + int chunk = (values.size() + tasks - 1) / tasks; + int submittedTasks = (values.size() + chunk - 1) / chunk; + var submitted = new CountDownLatch(submittedTasks); + var workers = new ThreadPoolExecutor(2, 2, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<>()) { + @Override public void execute(Runnable task) { + super.execute(task); + submitted.countDown(); + } + }; + var caller = Executors.newSingleThreadExecutor(); + var release = new CountDownLatch(1); + var blocked = new CountDownLatch(1); + var failed = new CountDownLatch(1); + var laterRecords = new AtomicInteger(); + var path = Files.createTempFile("parallel-later-failure", ".idx"); + try (var builder = new GraphIndexBuilder(values, VectorSimilarityFunction.DOT_PRODUCT, + 4, 16, 1.2f, 1.2f, false)) { + var graph = TestUtil.buildSequentially(builder, values); + try (var out = new io.github.jbellis.jvector.disk.BufferedRandomAccessWriter(path); + var writer = new ParallelGraphWriter(out, graph, java.util.List.of(new InlineVectors(2)), + new ParallelGraphWriter.Config(2, false, 1, 64), path, workers)) { + var writing = caller.submit(() -> { + writer.writeL0Records(new OrdinalMapper.IdentityMapper(values.size() - 1), + java.util.List.of(new InlineVectors(2)), + Feature.singleStateFactory(FeatureId.INLINE_VECTORS, ordinal -> { + if (ordinal == 0) { + blocked.countDown(); + try { if (!release.await(20, TimeUnit.SECONDS)) throw new AssertionError("release timeout"); } + catch (InterruptedException e) { throw new AssertionError(e); } + } else if (ordinal == chunk) { + try { if (!blocked.await(20, TimeUnit.SECONDS)) throw new AssertionError("first task did not block"); } + catch (InterruptedException e) { throw new AssertionError(e); } + failed.countDown(); + throw new IllegalStateException("later worker failure"); + } else if (failed.getCount() == 0) { + laterRecords.incrementAndGet(); + } + return new InlineVectors.State(values.getVector(ordinal)); + }), 0); + return null; + }); + assertTrue(failed.await(20, TimeUnit.SECONDS)); + assertTrue(submitted.await(20, TimeUnit.SECONDS)); + // The free worker reaches this sentinel only after draining the queued ranges. + assertEquals(42, workers.submit(() -> 42).get(20, TimeUnit.SECONDS).intValue()); + assertEquals("queued ranges encoded after another worker failed", 0, laterRecords.get()); + assertFalse("writer released resources while the first supplier was still running", writing.isDone()); + release.countDown(); + var error = assertThrows(ExecutionException.class, () -> writing.get(20, TimeUnit.SECONDS)); + assertTrue(error.getCause() instanceof IllegalStateException); + assertEquals("later worker failure", error.getCause().getMessage()); + assertEquals(42, workers.submit(() -> 42).get().intValue()); + } + } finally { + release.countDown(); + caller.shutdownNow(); workers.shutdownNow(); + caller.awaitTermination(20, TimeUnit.SECONDS); workers.awaitTermination(20, TimeUnit.SECONDS); + Files.deleteIfExists(path); + } + } +}