From f3687e6a4c9c47a185d970af8818fc03c8cec8c9 Mon Sep 17 00:00:00 2001 From: Ted Willke Date: Mon, 12 Oct 2026 14:51:04 +0000 Subject: [PATCH 01/20] Integrate deferred complete-record batches with reviewed dataset prefetching --- .../jvector/graph/disk/NodeRecordTask.java | 52 +++++++- .../graph/disk/ParallelGraphWriter.java | 4 +- .../github/jbellis/jvector/example/Grid.java | 94 ++++++++----- .../example/TestFinalWriteVectors.java | 109 +++++++++++++++ .../graph/disk/TestNodeRecordTask.java | 125 ++++++++++++++++++ 5 files changed, 345 insertions(+), 39 deletions(-) create mode 100644 jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestFinalWriteVectors.java create mode 100644 jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestNodeRecordTask.java 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..f4c52ea0d 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 @@ -38,8 +38,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 @@ -107,11 +107,54 @@ 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(); + int padding = recordSize - payloadBytes; + 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++) { 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(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); + } + for (int i = 0; i < padding; i++) writer.writeByte(0); + 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(); @@ -194,6 +237,11 @@ public Void call() throws Exception { ByteBuffer neighborData = writer.cloneBuffer(); writeFully(channel, neighborData, currentPosition); } + if (padding > 0) { + writer.reset(); + for (int i = 0; i < padding; i++) writer.writeByte(0); + writeFully(channel, writer.cloneBuffer(), recordOffset + payloadBytes); + } } return null; diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/ParallelGraphWriter.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/ParallelGraphWriter.java index 3475a4660..8968a1ee6 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/ParallelGraphWriter.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/ParallelGraphWriter.java @@ -62,6 +62,8 @@ * */ class ParallelGraphWriter implements AutoCloseable { + // Amortize positioned writes with bounded worker-local buffers. + private static final int RECORDS_PER_BATCH = 64; private final RandomAccessWriter writer; private final ImmutableGraphIndex graph; private final ExecutorService executor; @@ -173,7 +175,7 @@ public ParallelGraphWriter(RandomAccessWriter writer, // Thread-local buffers to avoid allocation overhead // Use BIG_ENDIAN to match Java DataOutput specification - final int bufferSize = recordSize; + final int bufferSize = Math.multiplyExact(recordSize, RECORDS_PER_BATCH); final boolean useDirect = config.useDirectBuffers; this.bufferPerThread = ThreadLocal.withInitial(() -> { ByteBuffer buffer = useDirect ? ByteBuffer.allocateDirect(bufferSize) : ByteBuffer.allocate(bufferSize); 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..e719f3e83 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 @@ -414,7 +414,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 +428,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,39 +437,31 @@ 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 + // Quantized peers are resident; serialize original-vector features after cleanup. 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. - builder.addGraphNode(node, vector); - }); + VectorAccess.forEach(floatVectors, PhysicalCoreExecutor.pool(), builder::addGraphNode); } 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)); - }); + IntStream.range(0, floatVectors.size()).parallel().forEach(node -> + builder.addGraphNode(node, vv.get().getVector(node))); }).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) + // Supply every feature so the parallel writer can emit 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)) { + Map> writeSuppliers = new EnumMap<>(FeatureId.class); + inputSuppliers.get(features).forEach((id, supplier) -> + writeSuppliers.put(id, ordinal -> supplier.apply(vectors.get(ordinal)))); + if (features.contains(FeatureId.FUSED_PQ)) { + var view = builder.getGraph().getView(); + writeSuppliers.put(FeatureId.FUSED_PQ, ordinal -> new FusedPQ.State(view, pq, ordinal)); + } writer.write(writeSuppliers); writer.close(); } catch (IOException e) { @@ -505,20 +495,51 @@ private static BuildOnDiskResult buildOnDisk(List> feat return new BuildOnDiskResult(indexes, fileSizes); } - 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); } - 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. + try { ranges.forEach(Range::close); } + finally { ranges.clear(); local.remove(); } + } + + 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, @@ -532,6 +553,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); 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..3be9bc0e6 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestFinalWriteVectors.java @@ -0,0 +1,109 @@ +/* + * 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(); } + } +} 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..63b37a622 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestNodeRecordTask.java @@ -0,0 +1,125 @@ +/* + * 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, 4096}) { + 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); + 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(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); + } + } +} From a1262e15173f72291ef7572d253331bcb5f0cb78 Mon Sep 17 00:00:00 2001 From: Ted Willke Date: Mon, 12 Oct 2026 20:03:06 +0000 Subject: [PATCH 02/20] Finish deferred record batching configuration and resource lifecycle --- docs/deferred-batch-writer.md | 62 ++++++ .../jvector/graph/disk/NodeRecordTask.java | 12 +- .../disk/OnDiskParallelGraphIndexWriter.java | 36 +++- .../graph/disk/ParallelGraphWriter.java | 146 +++++++------ .../RandomAccessOnDiskGraphIndexWriter.java | 15 +- .../github/jbellis/jvector/example/Grid.java | 19 +- .../example/TestFinalWriteVectors.java | 43 ++++ .../graph/disk/TestDeferredGraphWriter.java | 192 ++++++++++++++++++ .../graph/disk/TestNodeRecordTask.java | 2 +- 9 files changed, 446 insertions(+), 81 deletions(-) create mode 100644 docs/deferred-batch-writer.md create mode 100644 jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestDeferredGraphWriter.java diff --git a/docs/deferred-batch-writer.md b/docs/deferred-batch-writer.md new file mode 100644 index 000000000..124fa52a6 --- /dev/null +++ b/docs/deferred-batch-writer.md @@ -0,0 +1,62 @@ +# 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. + +## Supplying features at final write + +Provide the final feature suppliers to `write` after graph construction and cleanup. +Feature encoding then happens during serialization, rather than through a +synchronized `writeFeaturesInline` call on each insertion. This applies to NVQ, +full-precision inline vectors, fused PQ, and other features implementing the existing +`Feature` interface. The core writer does not choose a quantizer or retain input +vectors on the application's behalf. + +```java +try (var writer = new OnDiskParallelGraphIndexWriter.Builder(graph, outputPath) + .withParallelWorkerThreads(workerCount) + .with(new InlineVectors(vectors.dimension())) + .build()) { + var localVectors = vectors.threadLocalSupplier(); + writer.write(Feature.singleStateFactory(FeatureId.INLINE_VECTORS, + ordinal -> new InlineVectors.State(localVectors.get().getVector(ordinal)))); +} +``` + +The source must remain available until `write` finishes. Suppliers must support +concurrent calls; a source that reuses a scratch vector needs an independent view +per worker. Feature state is consumed while serializing its record. + +Applications can continue using `writeFeaturesInline` before the final write. +If a feature supplier is absent, its existing bytes are preserved and that write +uses the component-write path. Complete-record batching requires all inline feature +suppliers; it does not overwrite prewritten features to enable batching. + +## 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 + +Grid defers its inline NVQ or FP feature serialization until after cleanup. Its final +serialization suppliers consume worker-local ranges through `VectorAccess.openRange`. +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. 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 f4c52ea0d..d5cd90a01 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 @@ -95,8 +95,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; } @@ -121,7 +121,7 @@ public Void call() throws Exception { if (completeRecord) { if (buffer.position() == buffer.limit()) writer.reset(); int recordBegin = buffer.position(); - writer.writeInt(newOrdinal); + 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); @@ -159,7 +159,7 @@ public Void call() throws Exception { 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; @@ -175,8 +175,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 viewPerThread; - private final ThreadLocal bufferPerThread; + private final ExplicitThreadLocal bufferPerThread; private final CopyOnWriteArrayList allViews = new CopyOnWriteArrayList<>(); private final int recordSize; private final Path filePath; @@ -83,6 +84,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) @@ -91,6 +93,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; @@ -145,6 +155,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; @@ -159,12 +173,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(() -> { @@ -175,9 +183,8 @@ public ParallelGraphWriter(RandomAccessWriter writer, // Thread-local buffers to avoid allocation overhead // Use BIG_ENDIAN to match Java DataOutput specification - final int bufferSize = Math.multiplyExact(recordSize, RECORDS_PER_BATCH); 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; @@ -227,53 +234,52 @@ public void writeL0Records(OrdinalMapper ordinalMapper, try (var channel = AsynchronousFileChannel.open(filePath, opts, null)) { List> taskFutures = new ArrayList<>(numTasks); - // 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(() -> { + 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 + ); - taskFutures.add(future); - } + try { + return task.call(); + } 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); } + + } finally { + awaitTasks(taskFutures); } // Force all writes to disk before closing the channel. @@ -285,23 +291,30 @@ 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); + private void awaitTasks(List> futures) 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; + if (failure == null) failure = new IOException("Interrupted while writing records", e); + } catch (ExecutionException e) { + 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); } /** @@ -332,6 +345,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> feat writers.entrySet().stream().parallel().forEach(entry -> { var writer = entry.getValue(); var features = entry.getKey(); - try (var vectors = new FinalWriteVectors(floatVectors)) { + try (var vectors = new FinalWriteVectors(floatVectors); + var view = builder.getGraph().getView(); + var closeWriter = writer) { Map> writeSuppliers = new EnumMap<>(FeatureId.class); inputSuppliers.get(features).forEach((id, supplier) -> writeSuppliers.put(id, ordinal -> supplier.apply(vectors.get(ordinal)))); if (features.contains(FeatureId.FUSED_PQ)) { - var view = builder.getGraph().getView(); writeSuppliers.put(FeatureId.FUSED_PQ, ordinal -> new FusedPQ.State(view, pq, ordinal)); } writer.write(writeSuppliers); - writer.close(); } catch (IOException e) { throw new UncheckedIOException(e); } @@ -526,8 +526,17 @@ VectorFloat get(int ordinal) { @Override public void close() { // writer.write joins its workers before cursors and their read-ahead are released. - try { ranges.forEach(Range::close); } - finally { ranges.clear(); local.remove(); } + 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 { 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 index 3be9bc0e6..5c35708b1 100644 --- 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 @@ -106,4 +106,47 @@ public VectorFloat getVector(int ordinal) { 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/graph/disk/TestDeferredGraphWriter.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestDeferredGraphWriter.java new file mode 100644 index 000000000..3951b4f3e --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestDeferredGraphWriter.java @@ -0,0 +1,192 @@ +/* + * 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) { 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 index 63b37a622..e21e888de 100644 --- 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 @@ -93,7 +93,7 @@ public void testCompleteRecordBatches() throws Exception { try (var view = graph.getView()) { for (int n = 0; n < count; n++) { var expected = ByteBuffer.allocate(stride).order(ByteOrder.BIG_ENDIAN); - expected.putInt(n); + expected.putInt(mapper.newToOld(n) == OrdinalMapper.OMITTED ? -1 : n); int old = mapper.newToOld(n); if (old == OrdinalMapper.OMITTED) { expected.position(expected.position() + feature.featureSize()); From a3fc971254435939b65c20a6fc7a1b5659bf045c Mon Sep 17 00:00:00 2001 From: Ted Willke Date: Mon, 12 Oct 2026 20:10:24 +0000 Subject: [PATCH 03/20] Document existing parallel writer selection for Grid --- docs/deferred-batch-writer.md | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/docs/deferred-batch-writer.md b/docs/deferred-batch-writer.md index 124fa52a6..4d75d4934 100644 --- a/docs/deferred-batch-writer.md +++ b/docs/deferred-batch-writer.md @@ -60,3 +60,17 @@ This integrates with #744's optional `useDatasetPrefetching` setting. It does no 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 Grid parallel-write path + +Grid keeps the existing writer-selection setting. To exercise complete-record +parallel batches through BenchYAML, enable it before invoking the benchmark: + +```java +GraphIndexBuilderConfig.getInstance().setParallelBuild(true); +BenchYAML.main(new String[]{"your-index-parameters.yml"}); +``` + +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. From 2135f9561dc17b888766f8be5bf33711642875a9 Mon Sep 17 00:00:00 2001 From: Ted Willke Date: Mon, 12 Oct 2026 20:11:33 +0000 Subject: [PATCH 04/20] Add deferred batch writer release note for PR 747 --- docs/release notes/4.0.2/747.performance.md | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) create mode 100644 docs/release notes/4.0.2/747.performance.md 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..af048473c --- /dev/null +++ b/docs/release notes/4.0.2/747.performance.md @@ -0,0 +1,17 @@ +### 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. 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. + +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. The existing parallel-writer selection 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. These are separate fresh runs, +not a universal throughput or memory guarantee. See +[Deferred batch serialization](../../deferred-batch-writer.md) for usage. From 7ae3d808d5e7dce386674c58299f5782ab266426 Mon Sep 17 00:00:00 2001 From: Ted Willke Date: Mon, 12 Oct 2026 20:27:43 +0000 Subject: [PATCH 05/20] Clarify loader-independent production use of deferred serialization --- docs/deferred-batch-writer.md | 15 +++++++++++++++ docs/release notes/4.0.2/747.performance.md | 4 ++++ 2 files changed, 19 insertions(+) diff --git a/docs/deferred-batch-writer.md b/docs/deferred-batch-writer.md index 4d75d4934..ada33cea2 100644 --- a/docs/deferred-batch-writer.md +++ b/docs/deferred-batch-writer.md @@ -6,6 +6,21 @@ and adjacency list. By default, a worker writes up to 64 complete records togeth 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. + ## Supplying features at final write Provide the final feature suppliers to `write` after graph construction and cleanup. diff --git a/docs/release notes/4.0.2/747.performance.md b/docs/release notes/4.0.2/747.performance.md index af048473c..bec8775a4 100644 --- a/docs/release notes/4.0.2/747.performance.md +++ b/docs/release notes/4.0.2/747.performance.md @@ -5,6 +5,10 @@ batches of 64, with an optional positive batch-size override. Grid supplies inli NVQ or full-precision features during final serialization, after graph cleanup, and uses buffered vector ranges when the input source supports them. +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. + 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 From 29587063413ff542e0184392bacd12b1ce19b1cb Mon Sep 17 00:00:00 2001 From: Ted Willke Date: Mon, 12 Oct 2026 20:34:28 +0000 Subject: [PATCH 06/20] Document old and deferred caller serialization patterns --- docs/deferred-batch-writer.md | 74 +++++++++++++++++++++++++---------- 1 file changed, 54 insertions(+), 20 deletions(-) diff --git a/docs/deferred-batch-writer.md b/docs/deferred-batch-writer.md index ada33cea2..272b39dd8 100644 --- a/docs/deferred-batch-writer.md +++ b/docs/deferred-batch-writer.md @@ -21,34 +21,68 @@ benefits come from deferring per-insertion feature writes and grouping final wri independently of how input vectors are loaded. Buffered range access is an optional Grid integration, not a prerequisite for the core writer. -## Supplying features at final write +## Caller migration: early writes versus deferred batches -Provide the final feature suppliers to `write` after graph construction and cleanup. -Feature encoding then happens during serialization, rather than through a -synchronized `writeFeaturesInline` call on each insertion. This applies to NVQ, -full-precision inline vectors, fused PQ, and other features implementing the existing -`Feature` interface. The core writer does not choose a quantizer or retain input -vectors on the application's behalf. +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(graph, outputPath) - .withParallelWorkerThreads(workerCount) - .with(new InlineVectors(vectors.dimension())) +try (var writer = new OnDiskParallelGraphIndexWriter.Builder(builder.getGraph(), outputPath) + .withParallelWorkerThreads(writerWorkers) + .withMapper(new OrdinalMapper.IdentityMapper(count - 1)) + .with(new NVQ(nvq)) .build()) { - var localVectors = vectors.threadLocalSupplier(); - writer.write(Feature.singleStateFactory(FeatureId.INLINE_VECTORS, - ordinal -> new InlineVectors.State(localVectors.get().getVector(ordinal)))); + 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. } ``` -The source must remain available until `write` finishes. Suppliers must support -concurrent calls; a source that reuses a scratch vector needs an independent view -per worker. Feature state is consumed while serializing its record. +### 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))))); +} +``` -Applications can continue using `writeFeaturesInline` before the final write. -If a feature supplier is absent, its existing bytes are preserved and that write -uses the component-write path. Complete-record batching requires all inline feature -suppliers; it does not overwrite prewritten features to enable batching. +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 From b538ad14d488f0c30a1f5f1eb2b2e8e75da2a3be Mon Sep 17 00:00:00 2001 From: Ted Willke Date: Mon, 12 Oct 2026 21:20:38 +0000 Subject: [PATCH 07/20] Refresh construction examples for deferred complete-record writes --- README.md | 2 + docs/benchmarking.md | 7 +++ docs/legacy/jvector-step-by-step.md | 4 ++ docs/release notes/4.0.2/747.performance.md | 8 ++- docs/tutorials/2-disk-tutorial.md | 14 ++++- .../3-larger-than-memory-tutorial.md | 56 +++++++++-------- .../jbellis/jvector/example/IPCService.java | 15 ++++- .../benchmarks/datasets/SiftSmall.java | 63 +++++++++++-------- .../jvector/example/tutorial/DiskIntro.java | 17 +++-- .../example/tutorial/LargerThanMemory.java | 44 +++++-------- .../jvector/example/tutorial/NvqExample.java | 41 +++++------- .../graph/disk/ParallelWriteExample.java | 53 +++++++++++----- 12 files changed, 195 insertions(+), 129 deletions(-) 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/docs/benchmarking.md b/docs/benchmarking.md index 79ddba82e..ccb97764a 100644 --- a/docs/benchmarking.md +++ b/docs/benchmarking.md @@ -115,3 +115,10 @@ 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 + +Grid defers inline FP/NVQ serialization until after graph cleanup. The existing +parallel-writer selection enables complete-record batches (default 64); 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/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 index bec8775a4..aa3e6c6ab 100644 --- a/docs/release notes/4.0.2/747.performance.md +++ b/docs/release notes/4.0.2/747.performance.md @@ -16,6 +16,12 @@ finishes. The existing parallel-writer selection 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. These are separate fresh runs, +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-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(); From 01779a6e40f0fa50bb06a1540cf6953b8869a808 Mon Sep 17 00:00:00 2001 From: Ted Willke Date: Mon, 12 Oct 2026 21:20:38 +0000 Subject: [PATCH 08/20] Document missing hardware provenance for historical timings --- benchmarks-jmh/README.md | 4 +++- docs/release notes/4.0.0-RC.9/4.0.0-RC.9.md | 4 ++++ docs/release notes/4.0.0-RC.9/634.testing.md | 6 +++++- docs/release notes/4.0.0-RC.9/647.testing.md | 6 ++++-- docs/release notes/4.0.0-RC.9/667.testing.md | 6 +++++- docs/release notes/4.0.0-RC.9/671.testing.md | 5 ++++- 6 files changed, 25 insertions(+), 6 deletions(-) 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/release notes/4.0.0-RC.9/4.0.0-RC.9.md b/docs/release notes/4.0.0-RC.9/4.0.0-RC.9.md index 1de283f45..7176d7acc 100644 --- a/docs/release notes/4.0.0-RC.9/4.0.0-RC.9.md +++ b/docs/release notes/4.0.0-RC.9/4.0.0-RC.9.md @@ -638,3 +638,7 @@ Provide a possibility to store and load an empty graph from/to disk - - Helps to preserve graph metadata (similarity function, features, dimensions, even if the graph is empty) + +Hardware context for the timing examples and performance table above.[^historical-machine] + +[^historical-machine]: The original reports did not record the machine type or core count; these historical timings are not a hardware baseline. diff --git a/docs/release notes/4.0.0-RC.9/634.testing.md b/docs/release notes/4.0.0-RC.9/634.testing.md index 6bebd23b9..4045ae3dc 100644 --- a/docs/release notes/4.0.0-RC.9/634.testing.md +++ b/docs/release notes/4.0.0-RC.9/634.testing.md @@ -18,4 +18,8 @@ Disk Usage Summary Graph Index Build: Total Files: 1 Net Change: 1.61 GB, +1 files Index build time: 956.781921 seconds -``` \ No newline at end of file +``` + +Example output.[^machine] + +[^machine]: The original report did not record the machine type or core count; these historical timings are not a hardware baseline. diff --git a/docs/release notes/4.0.0-RC.9/647.testing.md b/docs/release notes/4.0.0-RC.9/647.testing.md index abeaad9fb..632826489 100644 --- a/docs/release notes/4.0.0-RC.9/647.testing.md +++ b/docs/release notes/4.0.0-RC.9/647.testing.md @@ -21,7 +21,7 @@ This PR optimizes `AccuracyMetrics` to improve the performance and readability o - **Duplicate Handling:** Confirmed that duplicate IDs in search results are handled via the "seen" Set, preventing artificial score inflation while maintaining $O(1)$ lookup speed. - **Exception Parity:** Explicitly verified that `IllegalArgumentException` strings remain 1:1 identical to the original implementation to prevent breaking downstream error-parsing. -## Performance results +## Performance results[^machine] Benchmarks were conducted using 10,000 queries on the `gecko-100k` dataset (768 dimensions). @@ -34,4 +34,6 @@ Benchmarks were conducted using 10,000 queries on the `gecko-100k` dataset (768 While the benefit is negligible for $K=10$, the optimization becomes critical as $K$ increases. The reduction in AP measurement time is particularly significant as it eliminates the $O(K^2)$ complexity previously caused by nested sublist scans. ## Notes / limitations -- **Autoboxing:** While this PR eliminates intermediate `List` allocations, it still utilizes `HashSet`. Further gains could be achieved using primitive-specific collections (like `IntHashSet`) if necessary in the future. \ No newline at end of file +- **Autoboxing:** While this PR eliminates intermediate `List` allocations, it still utilizes `HashSet`. Further gains could be achieved using primitive-specific collections (like `IntHashSet`) if necessary in the future. + +[^machine]: The original report did not record the machine type or core count; these historical timings are not a hardware baseline. diff --git a/docs/release notes/4.0.0-RC.9/667.testing.md b/docs/release notes/4.0.0-RC.9/667.testing.md index 5ec5261d2..d2a2696f2 100644 --- a/docs/release notes/4.0.0-RC.9/667.testing.md +++ b/docs/release notes/4.0.0-RC.9/667.testing.md @@ -12,4 +12,8 @@ Index build time: null seconds After: ``` Index build time: 32.192336 seconds -``` \ No newline at end of file +``` + +Example output.[^machine] + +[^machine]: The original report did not record the machine type or core count; these historical timings are not a hardware baseline. diff --git a/docs/release notes/4.0.0-RC.9/671.testing.md b/docs/release notes/4.0.0-RC.9/671.testing.md index 3a168ac4d..960beec5e 100644 --- a/docs/release notes/4.0.0-RC.9/671.testing.md +++ b/docs/release notes/4.0.0-RC.9/671.testing.md @@ -42,10 +42,13 @@ Overquery Avg QPS ± Std Dev CV % Mean Latency Avg Visit 5.00 28832.6 264.1 0.9 0.306 951.1 0.98 834.3 116.4 10.00 17835.6 1134.1 6.4 0.504 1536.7 0.99 786.6 116.4 - Overquery Avg QPS ± Std Dev CV % Mean Latency Avg Visited Recall@100 Max heap Max offheap (of 3) (ms) usage (MB) usage (MB) ---------------------------------------------------------------------------------------------------------------------------------- 1.00 18204.6 119.4 0.7 0.522 1536.7 0.85 803.5 116.4 2.00 11269.1 83.9 0.7 0.846 2495.2 0.97 803.5 116.4 ``` + +Example output.[^machine] + +[^machine]: The original report did not record the machine type or core count; these historical timings are not a hardware baseline. From 551ff6c32b31c2b2ff3bfaa71c047a70d637105d Mon Sep 17 00:00:00 2001 From: Ted Willke Date: Mon, 12 Oct 2026 22:05:40 +0000 Subject: [PATCH 09/20] Remove historical release-note edits from deferred writer PR --- docs/release notes/4.0.0-RC.9/4.0.0-RC.9.md | 4 ---- docs/release notes/4.0.0-RC.9/634.testing.md | 6 +----- docs/release notes/4.0.0-RC.9/647.testing.md | 6 ++---- docs/release notes/4.0.0-RC.9/667.testing.md | 6 +----- docs/release notes/4.0.0-RC.9/671.testing.md | 5 +---- 5 files changed, 5 insertions(+), 22 deletions(-) diff --git a/docs/release notes/4.0.0-RC.9/4.0.0-RC.9.md b/docs/release notes/4.0.0-RC.9/4.0.0-RC.9.md index 7176d7acc..1de283f45 100644 --- a/docs/release notes/4.0.0-RC.9/4.0.0-RC.9.md +++ b/docs/release notes/4.0.0-RC.9/4.0.0-RC.9.md @@ -638,7 +638,3 @@ Provide a possibility to store and load an empty graph from/to disk - - Helps to preserve graph metadata (similarity function, features, dimensions, even if the graph is empty) - -Hardware context for the timing examples and performance table above.[^historical-machine] - -[^historical-machine]: The original reports did not record the machine type or core count; these historical timings are not a hardware baseline. diff --git a/docs/release notes/4.0.0-RC.9/634.testing.md b/docs/release notes/4.0.0-RC.9/634.testing.md index 4045ae3dc..6bebd23b9 100644 --- a/docs/release notes/4.0.0-RC.9/634.testing.md +++ b/docs/release notes/4.0.0-RC.9/634.testing.md @@ -18,8 +18,4 @@ Disk Usage Summary Graph Index Build: Total Files: 1 Net Change: 1.61 GB, +1 files Index build time: 956.781921 seconds -``` - -Example output.[^machine] - -[^machine]: The original report did not record the machine type or core count; these historical timings are not a hardware baseline. +``` \ No newline at end of file diff --git a/docs/release notes/4.0.0-RC.9/647.testing.md b/docs/release notes/4.0.0-RC.9/647.testing.md index 632826489..abeaad9fb 100644 --- a/docs/release notes/4.0.0-RC.9/647.testing.md +++ b/docs/release notes/4.0.0-RC.9/647.testing.md @@ -21,7 +21,7 @@ This PR optimizes `AccuracyMetrics` to improve the performance and readability o - **Duplicate Handling:** Confirmed that duplicate IDs in search results are handled via the "seen" Set, preventing artificial score inflation while maintaining $O(1)$ lookup speed. - **Exception Parity:** Explicitly verified that `IllegalArgumentException` strings remain 1:1 identical to the original implementation to prevent breaking downstream error-parsing. -## Performance results[^machine] +## Performance results Benchmarks were conducted using 10,000 queries on the `gecko-100k` dataset (768 dimensions). @@ -34,6 +34,4 @@ Benchmarks were conducted using 10,000 queries on the `gecko-100k` dataset (768 While the benefit is negligible for $K=10$, the optimization becomes critical as $K$ increases. The reduction in AP measurement time is particularly significant as it eliminates the $O(K^2)$ complexity previously caused by nested sublist scans. ## Notes / limitations -- **Autoboxing:** While this PR eliminates intermediate `List` allocations, it still utilizes `HashSet`. Further gains could be achieved using primitive-specific collections (like `IntHashSet`) if necessary in the future. - -[^machine]: The original report did not record the machine type or core count; these historical timings are not a hardware baseline. +- **Autoboxing:** While this PR eliminates intermediate `List` allocations, it still utilizes `HashSet`. Further gains could be achieved using primitive-specific collections (like `IntHashSet`) if necessary in the future. \ No newline at end of file diff --git a/docs/release notes/4.0.0-RC.9/667.testing.md b/docs/release notes/4.0.0-RC.9/667.testing.md index d2a2696f2..5ec5261d2 100644 --- a/docs/release notes/4.0.0-RC.9/667.testing.md +++ b/docs/release notes/4.0.0-RC.9/667.testing.md @@ -12,8 +12,4 @@ Index build time: null seconds After: ``` Index build time: 32.192336 seconds -``` - -Example output.[^machine] - -[^machine]: The original report did not record the machine type or core count; these historical timings are not a hardware baseline. +``` \ No newline at end of file diff --git a/docs/release notes/4.0.0-RC.9/671.testing.md b/docs/release notes/4.0.0-RC.9/671.testing.md index 960beec5e..3a168ac4d 100644 --- a/docs/release notes/4.0.0-RC.9/671.testing.md +++ b/docs/release notes/4.0.0-RC.9/671.testing.md @@ -42,13 +42,10 @@ Overquery Avg QPS ± Std Dev CV % Mean Latency Avg Visit 5.00 28832.6 264.1 0.9 0.306 951.1 0.98 834.3 116.4 10.00 17835.6 1134.1 6.4 0.504 1536.7 0.99 786.6 116.4 + Overquery Avg QPS ± Std Dev CV % Mean Latency Avg Visited Recall@100 Max heap Max offheap (of 3) (ms) usage (MB) usage (MB) ---------------------------------------------------------------------------------------------------------------------------------- 1.00 18204.6 119.4 0.7 0.522 1536.7 0.85 803.5 116.4 2.00 11269.1 83.9 0.7 0.846 2495.2 0.97 803.5 116.4 ``` - -Example output.[^machine] - -[^machine]: The original report did not record the machine type or core count; these historical timings are not a hardware baseline. From eda0cfa1df9e316fb558d8780d3b2d09c582cb50 Mon Sep 17 00:00:00 2001 From: Jonathan Shook Date: Mon, 12 Oct 2026 23:07:43 +0000 Subject: [PATCH 10/20] Add failing repro: default Grid writer serializes deferred features on one thread Grid now supplies NVQ/FP features at the final write for every run, but only the parallel writer calls those suppliers from several threads. BenchYAML's default (sequential) writer calls them on one thread, so all final-write vector reads and NVQ encoding run single-threaded after construction. On Cohere-1M build+write went from 74-76 s to 100-104 s, and on OpenAI-3072-1M from 165 s to 264 s. Fails at ec42a94b (PR #747). --- .../TestGridDeferredFeatureParallelism.java | 184 ++++++++++++++++++ 1 file changed, 184 insertions(+) create mode 100644 jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestGridDeferredFeatureParallelism.java 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..39522b8ff --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestGridDeferredFeatureParallelism.java @@ -0,0 +1,184 @@ +/* + * 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); // BenchYAML's default writer selection + int reads = vectors.writeReads.get(); + assertTrue("final serialization read " + reads + " base vectors, all on one thread " + + vectors.writeThreads + ", with the default (sequential) writer", + reads == 0 || vectors.writeThreads.size() > 1); + } +} From cee3f9ec5676f75b3b9efcbc3bcde2282356da51 Mon Sep 17 00:00:00 2001 From: Jonathan Shook Date: Mon, 12 Oct 2026 23:07:43 +0000 Subject: [PATCH 11/20] Add failing repro: interrupted parallel write runs to completion ParallelGraphWriter.awaitTasks joins every task but never stops them, so an interrupted or already-failed write still serializes every remaining record before reporting. A 3 s write interrupted after 50 records took 3.2 s to stop and supplied all 6000 records. Fails at ec42a94b (PR #747). --- .../disk/TestParallelWriterInterrupt.java | 131 ++++++++++++++++++ 1 file changed, 131 insertions(+) create mode 100644 jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestParallelWriterInterrupt.java 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..b1db13976 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestParallelWriterInterrupt.java @@ -0,0 +1,131 @@ +/* + * 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.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.Assert.assertNotNull; +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); + } + } +} From cce40cec6ccd5bd3a9fa1d1c839562ef805e69f4 Mon Sep 17 00:00:00 2001 From: Jonathan Shook Date: Mon, 12 Oct 2026 23:08:14 +0000 Subject: [PATCH 12/20] Defer Grid feature serialization only when the parallel writer runs Grid moved NVQ/FP feature encoding into the final write for every run, but the sequential writer (BenchYAML's default) calls final suppliers on one thread. Grid now defers features only when every writer is an OnDiskParallelGraphIndexWriter; otherwise it encodes them on the insertion threads and writes them with writeFeaturesInline, as before #747. The docs and release note say so. Fixes TestGridDeferredFeatureParallelism. --- docs/deferred-batch-writer.md | 6 ++- docs/release notes/4.0.2/747.performance.md | 8 ++-- .../github/jbellis/jvector/example/Grid.java | 40 +++++++++++++++---- 3 files changed, 42 insertions(+), 12 deletions(-) diff --git a/docs/deferred-batch-writer.md b/docs/deferred-batch-writer.md index 272b39dd8..05aa26f8e 100644 --- a/docs/deferred-batch-writer.md +++ b/docs/deferred-batch-writer.md @@ -99,8 +99,10 @@ remains caller-owned and must outlive the write. ## Dataset prefetching in Grid -Grid defers its inline NVQ or FP feature serialization until after cleanup. Its final -serialization suppliers consume worker-local ranges through `VectorAccess.openRange`. +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 default 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. diff --git a/docs/release notes/4.0.2/747.performance.md b/docs/release notes/4.0.2/747.performance.md index aa3e6c6ab..b1d74dfce 100644 --- a/docs/release notes/4.0.2/747.performance.md +++ b/docs/release notes/4.0.2/747.performance.md @@ -1,9 +1,11 @@ ### 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. 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. +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 default sequential writer, 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 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 13571ef62..50bcf2631 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 @@ -437,20 +437,30 @@ private static BuildOnDiskResult buildOnDisk(List> feat throw new IllegalStateException("Bench looks for either NVQ_VECTORS or INLINE_VECTORS feature set for scoring compressed builds."); } - // Quantized peers are resident; serialize original-vector features after cleanup. + // 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(), builder::addGraphNode); + VectorAccess.forEach(floatVectors, PhysicalCoreExecutor.pool(), (node, vector) -> { + 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 -> - builder.addGraphNode(node, vv.get().getVector(node))); + IntStream.range(0, floatVectors.size()).parallel().forEach(node -> { + var vector = vv.get().getVector(node); + if (!deferFeatures) writeInlineStates(writers, inputSuppliers, node, vector); + builder.addGraphNode(node, vector); + }); }).join(); } builder.cleanup(); - // Supply every feature so the parallel writer can emit complete records. + // 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(); @@ -458,8 +468,10 @@ private static BuildOnDiskResult buildOnDisk(List> feat var view = builder.getGraph().getView(); var closeWriter = writer) { Map> writeSuppliers = new EnumMap<>(FeatureId.class); - inputSuppliers.get(features).forEach((id, supplier) -> - writeSuppliers.put(id, ordinal -> supplier.apply(vectors.get(ordinal)))); + 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)); } @@ -495,6 +507,20 @@ 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, Feature.State>>> inputSuppliers, + int node, VectorFloat vector) { + writers.forEach((features, writer) -> { + var states = new EnumMap(FeatureId.class); + 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; From eb560d88950662b21ed93064fdfbbb1b0850f9a0 Mon Sep 17 00:00:00 2001 From: Jonathan Shook Date: Mon, 12 Oct 2026 23:08:58 +0000 Subject: [PATCH 13/20] Stop remaining parallel write tasks on interrupt or first failure awaitTasks joined every task but never stopped them, so an interrupted or failed write serialized every remaining record before reporting. writeL0Records now shares a stop flag with its tasks: awaitTasks sets it on interruption or the first task failure, unstarted tasks return at once, and running tasks stop before their next record. Every task is still joined before the channel and suppliers are released. NodeRecordTask keeps its existing constructor. Fixes TestParallelWriterInterrupt. --- .../jvector/graph/disk/NodeRecordTask.java | 23 +++++++++++++++++++ .../graph/disk/ParallelGraphWriter.java | 20 ++++++++++++---- 2 files changed, 39 insertions(+), 4 deletions(-) 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 d5cd90a01..5872147da 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; /** @@ -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; } /** @@ -113,6 +135,7 @@ public Void call() throws Exception { 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; diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/ParallelGraphWriter.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/ParallelGraphWriter.java index db40f38e6..3498ec17a 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/ParallelGraphWriter.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/ParallelGraphWriter.java @@ -39,6 +39,7 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.ExecutionException; import java.util.concurrent.Future; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.IntFunction; @@ -233,6 +234,9 @@ public void writeL0Records(OrdinalMapper ordinalMapper, 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; // Always join submitted tasks before closing the channel or supplier resources. try { @@ -250,6 +254,7 @@ public void writeL0Records(OrdinalMapper ordinalMapper, final int end = endOrdinal; Future future = executor.submit(() -> { + if (stop.get()) return null; var view = viewPerThread.get(); var buffer = bufferPerThread.get(); @@ -264,7 +269,8 @@ public void writeL0Records(OrdinalMapper ordinalMapper, recordSize, baseOffset, // Base offset (task calculates per-ordinal offsets) channel, // Async file channel for position-based writes - buffer // Thread-local buffer + buffer, // Thread-local buffer + stop::get ); try { @@ -277,9 +283,10 @@ public void writeL0Records(OrdinalMapper ordinalMapper, taskFutures.add(future); } - + submitted = true; } finally { - awaitTasks(taskFutures); + if (!submitted) stop.set(true); + awaitTasks(taskFutures, stop); } // Force all writes to disk before closing the channel. @@ -291,7 +298,10 @@ public void writeL0Records(OrdinalMapper ordinalMapper, } } - private void awaitTasks(List> futures) throws IOException { + /// 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) { @@ -302,8 +312,10 @@ private void awaitTasks(List> futures) throws IOException { 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; From 477e9b2c465ec9dda71e23e95ea9ca73c4b8cfa5 Mon Sep 17 00:00:00 2001 From: Jonathan Shook Date: Tue, 13 Oct 2026 00:18:21 +0000 Subject: [PATCH 14/20] Build TestBufferedGraphCompatibility graphs deterministically The test built each graph on a 3-thread ForkJoinPool, so insertion order, hierarchy-level draws and refinement cleanup varied with thread timing. Its self-match check failed intermittently under full-suite load for the BQ, 768-dimension, DOT_PRODUCT, refined configuration (node 5 missing from its own top 10). Graphs are now built through builder.build(source) with single-worker executors, and the file-backed graph's adjacency must equal a resident-built reference graph at every level. That checks the test's purpose (a shared file source must not overwrite insertion or diversity queries) edge for edge. The self-match check is kept and is now deterministic for the fixed data. --- .../graph/TestBufferedGraphCompatibility.java | 34 +++++++++++++++++-- 1 file changed, 32 insertions(+), 2 deletions(-) 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(); } } } From b8300f667b7fc1f55f885eaa728a45fecd62e77f Mon Sep 17 00:00:00 2001 From: Ted Willke Date: Tue, 13 Oct 2026 05:05:59 +0000 Subject: [PATCH 15/20] Signal parallel writer failure before joining earlier ranges --- .../graph/disk/ParallelGraphWriter.java | 40 +++++----- .../graph/disk/TestDeferredGraphWriter.java | 7 +- .../disk/TestParallelWriterInterrupt.java | 77 +++++++++++++++++++ 3 files changed, 105 insertions(+), 19 deletions(-) diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/ParallelGraphWriter.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/ParallelGraphWriter.java index 3498ec17a..ce0d97e55 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/ParallelGraphWriter.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/ParallelGraphWriter.java @@ -231,7 +231,7 @@ 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. @@ -255,26 +255,30 @@ public void writeL0Records(OrdinalMapper ordinalMapper, Future future = executor.submit(() -> { if (stop.get()) return null; - var view = viewPerThread.get(); - var buffer = bufferPerThread.get(); + try { + var view = viewPerThread.get(); + var buffer = bufferPerThread.get(); - 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 - ); + 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 + ); - try { 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(); 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 index 3951b4f3e..7a6071830 100644 --- 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 @@ -135,7 +135,12 @@ public void failedWriteJoinsWorkersAndPreservesExternalExecutor() throws Excepti 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) { failed.countDown(); throw new IllegalStateException("supplier failure"); } + 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); } 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 index b1db13976..e79a99aa7 100644 --- 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 @@ -33,11 +33,18 @@ 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. @@ -128,4 +135,74 @@ public void interruptedWriteStopsPromptly() throws Exception { 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); + } + } } From 1f6c84b0c9b5345ef41c7542176050c9546919fc Mon Sep 17 00:00:00 2001 From: Ted Willke Date: Tue, 13 Oct 2026 05:05:59 +0000 Subject: [PATCH 16/20] Remove unused node-record padding from V6 serialization --- .../jbellis/jvector/graph/disk/NodeRecordTask.java | 10 +++------- .../jbellis/jvector/graph/disk/TestNodeRecordTask.java | 7 ++++++- 2 files changed, 9 insertions(+), 8 deletions(-) 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 5872147da..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 @@ -131,7 +131,9 @@ public Void call() throws Exception { int payloadBytes = Integer.BYTES * (2 + graph.getDegree(0)) + inlineFeatures.stream().mapToInt(Feature::featureSize).sum(); - int padding = recordSize - payloadBytes; + 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++) { @@ -167,7 +169,6 @@ public Void call() throws Exception { } for (; n < graph.getDegree(0); n++) writer.writeInt(-1); } - for (int i = 0; i < padding; i++) writer.writeByte(0); 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; @@ -262,11 +263,6 @@ public Void call() throws Exception { ByteBuffer neighborData = writer.cloneBuffer(); writeFully(channel, neighborData, currentPosition); } - if (padding > 0) { - writer.reset(); - for (int i = 0; i < padding; i++) writer.writeByte(0); - writeFully(channel, writer.cloneBuffer(), recordOffset + payloadBytes); - } } return null; 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 index e21e888de..12b2090c4 100644 --- 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 @@ -68,7 +68,7 @@ public void testCompleteRecordBatches() throws Exception { 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, 4096}) { + 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. @@ -85,6 +85,11 @@ public void testCompleteRecordBatches() throws Exception { 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); From ae8a3a8819e2397382936530b27844aad3273360 Mon Sep 17 00:00:00 2001 From: Ted Willke Date: Tue, 13 Oct 2026 05:05:59 +0000 Subject: [PATCH 17/20] Cover concurrent buffered vector access during graph construction --- ...stConcurrentBufferedGraphConstruction.java | 145 ++++++++++++++++++ 1 file changed, 145 insertions(+) create mode 100644 jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestConcurrentBufferedGraphConstruction.java 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(); } + } +} From adf14e29e91c624f7a435f15cc5715ae02c20972 Mon Sep 17 00:00:00 2001 From: Ted Willke Date: Tue, 13 Oct 2026 05:05:59 +0000 Subject: [PATCH 18/20] Default benchmark entry points to parallel graph serialization --- docs/benchmarking.md | 7 ++-- docs/deferred-batch-writer.md | 20 +++++----- docs/release notes/4.0.2/747.performance.md | 6 ++- .../jvector/example/AutoBenchYAML.java | 2 + .../jbellis/jvector/example/BenchYAML.java | 2 + .../github/jbellis/jvector/example/Grid.java | 9 +++++ .../TestGridDeferredFeatureParallelism.java | 40 ++++++++++++++++++- 7 files changed, 70 insertions(+), 16 deletions(-) diff --git a/docs/benchmarking.md b/docs/benchmarking.md index ccb97764a..5809b9152 100644 --- a/docs/benchmarking.md +++ b/docs/benchmarking.md @@ -118,7 +118,8 @@ For remote datasets, use `base_url` to specify where files should be downloaded ## Final index serialization -Grid defers inline FP/NVQ serialization until after graph cleanup. The existing -parallel-writer selection enables complete-record batches (default 64); dataset -prefetching is a separate option and is not required. See +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 index 05aa26f8e..df0129b47 100644 --- a/docs/deferred-batch-writer.md +++ b/docs/deferred-batch-writer.md @@ -101,7 +101,7 @@ remains caller-owned and must outlive the write. 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 default sequential writer, which calls suppliers on one +`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 @@ -112,16 +112,18 @@ change construction scoring, quantization parameters, or the query path. Product 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 Grid parallel-write path +## Running the benchmark writer paths -Grid keeps the existing writer-selection setting. To exercise complete-record -parallel batches through BenchYAML, enable it before invoking the benchmark: - -```java -GraphIndexBuilderConfig.getInstance().setParallelBuild(true); -BenchYAML.main(new String[]{"your-index-parameters.yml"}); -``` +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/release notes/4.0.2/747.performance.md b/docs/release notes/4.0.2/747.performance.md index b1d74dfce..8198583ca 100644 --- a/docs/release notes/4.0.2/747.performance.md +++ b/docs/release notes/4.0.2/747.performance.md @@ -4,7 +4,7 @@ Parallel graph serialization can now write complete layer-zero records in bounde 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 default sequential writer, Grid keeps encoding those +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 @@ -14,7 +14,9 @@ to fully preloaded inputs; the source does not have to preload the whole dataset 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. The existing parallel-writer selection remains unchanged. +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 diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/AutoBenchYAML.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/AutoBenchYAML.java index a3d1a80bc..eeac4c9bd 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/AutoBenchYAML.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/AutoBenchYAML.java @@ -66,6 +66,8 @@ public static void main(String[] args) throws IOException { System.exit(1); } + Grid.configureWriterForBenchmark(); + logger.info("Heap space available is {}", Runtime.getRuntime().maxMemory()); // Initialize checkpoint manager diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/BenchYAML.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/BenchYAML.java index e5ab3d970..01b74bc4d 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/BenchYAML.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/BenchYAML.java @@ -43,6 +43,8 @@ public static void main(String[] args) throws IOException { if (args == null) { throw new InvalidParameterException("argv[] is null, check your maven exec config"); } + Grid.configureWriterForBenchmark(); + String regex = Arrays.stream(args) .filter(Objects::nonNull) .flatMap(s -> 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 50bcf2631..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; @@ -624,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-tests/src/test/java/io/github/jbellis/jvector/example/TestGridDeferredFeatureParallelism.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestGridDeferredFeatureParallelism.java index 39522b8ff..7a5bbd3b7 100644 --- 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 @@ -175,10 +175,46 @@ public void defaultWriterDoesNotConfineDeferredFeatureInputToOneThread() throws control.writeReads.get() > 0 && control.writeThreads.size() > 1); var vectors = new RecordingVectors(resident); - runGrid(list, vectors, false); // BenchYAML's default writer selection + 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 default (sequential) writer", + + 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); + } + } } From 74ea1bf2dae3888453ac3fdbee91b6bf59e04062 Mon Sep 17 00:00:00 2001 From: Ted Willke Date: Tue, 13 Oct 2026 05:13:40 +0000 Subject: [PATCH 19/20] Clarify parallel writer requirement for deferred encoding --- docs/deferred-batch-writer.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/docs/deferred-batch-writer.md b/docs/deferred-batch-writer.md index df0129b47..b36015074 100644 --- a/docs/deferred-batch-writer.md +++ b/docs/deferred-batch-writer.md @@ -23,6 +23,11 @@ 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 From 201dd576a02ec30955db8a615ff622bf76bbf21a Mon Sep 17 00:00:00 2001 From: Ted Willke Date: Tue, 13 Oct 2026 05:14:13 +0000 Subject: [PATCH 20/20] Document parallel encoding requirement for integrators in release note --- docs/release notes/4.0.2/747.performance.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/docs/release notes/4.0.2/747.performance.md b/docs/release notes/4.0.2/747.performance.md index 8198583ca..3c36111ee 100644 --- a/docs/release notes/4.0.2/747.performance.md +++ b/docs/release notes/4.0.2/747.performance.md @@ -11,6 +11,11 @@ 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