diff --git a/docs/release notes/4.1.0/695.performance.md b/docs/release notes/4.1.0/695.performance.md new file mode 100644 index 000000000..87c019168 --- /dev/null +++ b/docs/release notes/4.1.0/695.performance.md @@ -0,0 +1,98 @@ +### Fully Asynchronous Parallel Graph Index Writes + +**Description** +`OnDiskParallelGraphIndexWriter` (introduced in #608) parallelizes serialization of +Level-0 (L0) node records to disk using Java's `AsynchronousFileChannel`. Although it +already used the async channel API, each write was previously submitted and immediately +blocked on via `Future.get()` before the next one began — writing a node's ordinal, then +each feature, then its neighbor list as a sequence of blocking round-trips rather than a +truly asynchronous pipeline. This PR removes that bottleneck so parallel writes take full +advantage of async I/O: + +- **Fast path** (no pre-written features — the common case): each task now packs its + entire ordinal range into a single contiguous `ByteBuffer` and issues one + `channel.write()` call for the whole range, instead of one blocking write per node field. +- **Legacy path** (some features already placed on disk via `writeFeaturesInline()`): each + task identifies the contiguous byte spans it still owns, submits every write for its + entire range up front, and only then waits on the collected futures — letting the OS + schedule the full I/O workload instead of alternating write-then-wait per span. +- The old per-thread scratch `ByteBuffer`, sized to a single record, forced writes to + serialize at the buffer level regardless of channel concurrency. It has been removed; + each task now allocates its own range-sized (fast path) or per-region (legacy path) + buffer instead. + +**Also in this PR:** a `parallelGraphConstruction` boolean was added to +`ConstructionParameters`, letting BenchYAML / AutoBenchYAML test configs opt into +`OnDiskParallelGraphIndexWriter` for index construction instead of the default serial +`OnDiskGraphIndexWriter`, without any code changes. + +**Performance** +Example run writing with NVQ + FUSED_ADC features. Before this change, parallel writes +were ~4x-8x faster than sequential: +``` +Sequential write: 6074.18 ms +Parallel write: 1373.73 ms +Speedup: 4.42x +``` +After this change, the same comparison shows ~24x-32x speedups: +``` +Sequential write: 20147.29 ms +Parallel write: 627.52 ms +Speedup: 32.11x +``` + +**How to Enable** + +*Programmatic API* — no change; `OnDiskParallelGraphIndexWriter.Builder` (available since +#608) is used the same way as before. The throughput improvement is automatic: + +```java +// Serial (existing) path +var writer = new OnDiskGraphIndexWriter.Builder(graph, outputPath) + .withMapper(ordinalMapper) + .with(new InlineVectors(dimension)) + .build(); + +// Parallel path +var writer = new OnDiskParallelGraphIndexWriter.Builder(graph, outputPath) + .withMapper(ordinalMapper) + .with(new InlineVectors(dimension)) + // Optional tuning: + .withParallelWorkerThreads(0) // 0 = use available processors + .withParallelDirectBuffers(false) // true = off-heap ByteBuffers + .build(); + +writer.write(featureStateSuppliers); +writer.close(); +``` + +To supply a shared, externally-managed executor (e.g. to bound total thread count across +concurrent builds): + +```java +ExecutorService ioPool = Executors.newFixedThreadPool(16); +var writer = new OnDiskParallelGraphIndexWriter.Builder(graph, outputPath) + .withMapper(ordinalMapper) + .with(new InlineVectors(dimension)) + .withExecutor(ioPool) // caller is responsible for shutdown + .build(); +``` + +*BenchYAML / AutoBenchYAML* — new in this PR: add the following field under +`construction` in any index-parameter YAML config file: + +```yaml +construction: + parallelGraphConstruction: Yes # default: No +``` + +When set to `Yes`, `Grid` uses `OnDiskParallelGraphIndexWriter` for the on-disk build path. +The field is optional and defaults to `No` (serial writes) if omitted, so existing config +files require no changes. + +**Notes** +- The writer produces output in the same on-disk format as `OnDiskGraphIndexWriter`; indexes + written with either class are interchangeable and loaded with `OnDiskGraphIndex.load()`. + This PR changes only the internal write scheduling, not the on-disk format. +- Write tasks perform blocking file I/O. For best performance supply an I/O-sized thread pool + (thread count ≥ logical cores) rather than a compute-sized pool when using `withExecutor()`. 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..f58a9deb5 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 @@ -24,35 +24,62 @@ import java.io.IOException; import java.nio.ByteBuffer; import java.nio.channels.AsynchronousFileChannel; +import java.util.ArrayList; import java.util.List; import java.util.Map; import java.util.concurrent.Callable; -import java.util.concurrent.ExecutionException; +import java.util.concurrent.Future; import java.util.function.IntFunction; /** - * A task that writes L0 records for a range of nodes directly to disk using synchronous position-based writes. + * A task that writes L0 records for a range of nodes to disk via an AsynchronousFileChannel. *

- * This task is designed to be executed in a thread pool, with each worker thread - * owning its own ImmutableGraphIndex.View for thread-safe neighbor iteration. - * Each task processes a contiguous range of ordinals to reduce task creation overhead. + * Each task processes a contiguous range of ordinals. Two execution paths exist: *

- * 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. + * Fast path (no pre-written features): all records in the range are built into one + * contiguous {@link ByteBuffer} and written with a single {@code channel.write()} call. + *

+ * Legacy path (pre-written features present): some feature regions in each node record + * were written to disk ahead of time via {@code writeFeaturesInline()}. Those byte ranges must + * not be overwritten. {@link #callLegacy()} builds one growing in-memory buffer for the whole + * task range, but skips writing anything into it for a pre-written (gap) feature — the buffer + * ends up holding only owned bytes, packed contiguously. A run of owned bytes is sliced off and + * flushed as a single non-blocking write only when a gap is reached (or at the end of the task), + * so a run commonly spans many nodes rather than being flushed per-node: the seam + * between one node's trailing neighbor section and the next node's leading ordinal is never + * itself a gap, so runs only break where a {@link FeatureId} is actually pre-written, regardless + * of node boundaries. All writes for the entire task are submitted before any + * {@code Future.get()} call, so the OS sees the full I/O workload and can schedule it efficiently. + *

+ * Understanding {@code hasPrewrittenFeatures}: this flag is derived from the + * {@code featureStateSuppliers} map passed to {@code write()}. It does not involve + * any read-before-write. The mechanism is purely contractual: a client that calls + * {@code writeFeaturesInline(ordinal, stateMap)} before {@code write(featureStateSuppliers)} + * simply omits those {@link FeatureId}s from the suppliers map. + * {@code featureStateSuppliers.get(featureId) == null} is the signal "those bytes are already + * on disk — do not touch them." The flag is computed once at construction time so the + * per-node hot path pays no overhead checking it. + *

+ * FUTURE IMPROVEMENT: when {@code writeFeaturesInline} support is removed, {@code hasPrewrittenFeatures} + * will always be {@code false}, the legacy path ({@link #callLegacy()}) and all related helpers + * can be deleted, and this class becomes the clean single-path fast path only. */ class NodeRecordTask implements Callable { - private final int startOrdinal; // Inclusive - private final int endOrdinal; // Exclusive + private final int startOrdinal; + private final int endOrdinal; private final OrdinalMapper ordinalMapper; private final ImmutableGraphIndex graph; private final ImmutableGraphIndex.View view; private final List inlineFeatures; private final Map> featureStateSuppliers; private final int recordSize; - private final long baseOffset; // Base file offset for L0 (offsets calculated per-ordinal) + private final long baseOffset; private final AsynchronousFileChannel channel; - private final ByteBuffer buffer; // Thread-local buffer for building record components + private final boolean useDirectBuffers; + + // FUTURE IMPROVEMENT: when writeFeaturesInline is removed this flag is always false, + // the callLegacy() branch disappears, and this field can be deleted entirely. + private final boolean hasPrewrittenFeatures; NodeRecordTask(int startOrdinal, int endOrdinal, @@ -64,7 +91,7 @@ class NodeRecordTask implements Callable { int recordSize, long baseOffset, AsynchronousFileChannel channel, - ByteBuffer buffer) { + boolean useDirectBuffers) { this.startOrdinal = startOrdinal; this.endOrdinal = endOrdinal; this.ordinalMapper = ordinalMapper; @@ -75,128 +102,299 @@ class NodeRecordTask implements Callable { this.recordSize = recordSize; this.baseOffset = baseOffset; this.channel = channel; - this.buffer = buffer; + this.useDirectBuffers = useDirectBuffers; + // Null supplier for any inline feature means the caller omitted it from the write() + // suppliers map, signalling that writeFeaturesInline() already placed that data on disk. + // FUTURE IMPROVEMENT: remove this field once writeFeaturesInline support is dropped. + this.hasPrewrittenFeatures = inlineFeatures.stream() + .anyMatch(f -> featureStateSuppliers.get(f.id()) == null); } - /** - * Writes a buffer fully to the channel at the specified position. - * Ensures all bytes are written by looping until the buffer is empty. - * This is critical for correctness as AsynchronousFileChannel.write() may not write all bytes in one call. - * - * @param channel the channel to write to - * @param buffer the buffer to write (will be fully consumed) - * @param position the file position to write at - * @throws IOException if an I/O error occurs - * @throws ExecutionException if the write operation fails - * @throws InterruptedException if interrupted while waiting for write completion - */ - private static void writeFully(AsynchronousFileChannel channel, ByteBuffer buffer, long position) - throws IOException, ExecutionException, InterruptedException { - long currentPosition = position; - while (buffer.hasRemaining()) { - int written = channel.write(buffer, currentPosition).get(); - if (written < 0) { - throw new IOException("Channel closed while writing"); - } - currentPosition += written; + @Override + public Void call() throws Exception { + // FUTURE IMPROVEMENT: once writeFeaturesInline is removed this dispatch goes away — + // callBatched() becomes the only execution path. + if (hasPrewrittenFeatures) { + callLegacy(); + } else { + callBatched(); } + return null; } - @Override - public Void call() throws Exception { - // Reuse writer and buffer across all ordinals in this range - var writer = new ByteBufferIndexWriter(buffer); + // ------------------------------------------------------------------------- + // Fast path: one contiguous buffer for the entire ordinal range, one write. + // ------------------------------------------------------------------------- + + private void callBatched() throws Exception { + int rangeSize = endOrdinal - startOrdinal; + ByteBuffer rangeBuffer = useDirectBuffers + ? ByteBuffer.allocateDirect(rangeSize * recordSize) + : ByteBuffer.allocate(rangeSize * recordSize); + rangeBuffer.order(java.nio.ByteOrder.BIG_ENDIAN); + + // ByteBufferIndexWriter clears the buffer on construction; since it was just + // allocated that is a no-op, but it sets initialPosition = 0 as required. + var writer = new ByteBufferIndexWriter(rangeBuffer); for (int newOrdinal = startOrdinal; newOrdinal < endOrdinal; newOrdinal++) { - var originalOrdinal = ordinalMapper.newToOld(newOrdinal); - long recordOffset = baseOffset + (long) newOrdinal * recordSize; - long currentPosition = recordOffset; + buildFullRecord(writer, newOrdinal); + } - // Reset buffer for this ordinal - writer.reset(); + // One channel.write() for the entire task range — one syscall, one OS I/O request + // in the common case. writeAllFully() guarantees the buffer is fully drained even + // if the OS reports a short write (e.g. disk full). + rangeBuffer.flip(); + writeAllFully(List.of(new PendingWrite(rangeBuffer, baseOffset + (long) startOrdinal * recordSize))); + } - // Write node ordinal - writer.writeInt(newOrdinal); - ByteBuffer ordinalData = writer.cloneBuffer(); - writeFully(channel, ordinalData, currentPosition); - currentPosition += Integer.BYTES; + // ------------------------------------------------------------------------- + // Legacy path: handle pre-written feature regions. + // + // Pre-written bytes must not be overwritten. A single growing buffer is built for the + // whole task range, same as callBatched(), except nothing is written into it for a + // pre-written (gap) feature -- the buffer ends up holding only owned bytes, packed + // contiguously. A run of owned bytes is sliced off the buffer and queued as a single + // non-blocking write only when a gap is reached (or at the end of the task), so runs + // commonly span many nodes rather than being flushed per-node: the seam between one + // node's trailing neighbor section and the next node's leading ordinal is never itself + // a gap, so runs only break where a FeatureId is actually pre-written. ALL writes for + // the entire task are submitted before any Future.get() call, letting the OS pipeline + // them. + // + // FUTURE IMPROVEMENT: delete this method entirely once writeFeaturesInline support is + // removed. The fast path handles everything. + // ------------------------------------------------------------------------- + + private void callLegacy() throws Exception { + int rangeSize = endOrdinal - startOrdinal; + ByteBuffer rangeBuffer = useDirectBuffers + ? ByteBuffer.allocateDirect(rangeSize * recordSize) + : ByteBuffer.allocate(rangeSize * recordSize); + rangeBuffer.order(java.nio.ByteOrder.BIG_ENDIAN); + var writer = new ByteBufferIndexWriter(rangeBuffer); + + List pending = new ArrayList<>(); + // Buffer-relative start of the run currently being accumulated, and the file + // position that buffer offset corresponds to. Advanced past both the flushed run + // and the skipped gap every time a gap is hit. + int runStart = 0; + long runFilePosition = baseOffset + (long) startOrdinal * recordSize; + + for (int newOrdinal = startOrdinal; newOrdinal < endOrdinal; newOrdinal++) { + var originalOrdinal = ordinalMapper.newToOld(newOrdinal); - // Handle OMITTED nodes (holes in ordinal space) if (originalOrdinal == OrdinalMapper.OMITTED) { - // Write placeholder: zeros for features and empty neighbor list - writer.reset(); + // OMITTED nodes are holes in the ordinal space. writeFeaturesInline() is + // never called for them, so nothing here is a gap: the whole record extends + // the current run with no flush needed. + writer.writeInt(newOrdinal); for (var feature : inlineFeatures) { - // Write zeros for missing features - for (int i = 0; i < feature.featureSize(); i++) { - writer.writeByte(0); - } + for (int i = 0; i < feature.featureSize(); i++) writer.writeByte(0); } - ByteBuffer featureData = writer.cloneBuffer(); - writeFully(channel, featureData, currentPosition); - currentPosition += featureData.remaining(); - - // Write empty neighbor list - writer.reset(); writer.writeInt(0); // neighbor count - for (int n = 0; n < graph.getDegree(0); n++) { - writer.writeInt(-1); // padding - } - ByteBuffer neighborData = writer.cloneBuffer(); - writeFully(channel, neighborData, currentPosition); - } else { - // Validate node exists - if (!graph.containsNode(originalOrdinal)) { - throw new IllegalStateException( - String.format("Ordinal mapper mapped new ordinal %s to non-existing node %s", - newOrdinal, originalOrdinal)); - } + for (int n = 0; n < graph.getDegree(0); n++) writer.writeInt(-1); + continue; + } - // Write inline features (skip if supplier is null - feature was pre-written) - for (var feature : inlineFeatures) { - var supplier = featureStateSuppliers.get(feature.id()); - if (supplier != null) { - // Feature not pre-written, write it now - writer.reset(); - feature.writeInline(writer, supplier.apply(originalOrdinal)); - ByteBuffer featureData = writer.cloneBuffer(); - writeFully(channel, featureData, currentPosition); - } - // Skip to next feature position (whether we wrote it or not) - currentPosition += feature.featureSize(); - } + if (!graph.containsNode(originalOrdinal)) { + throw new IllegalStateException(String.format( + "Ordinal mapper mapped new ordinal %d to non-existing node %d", + newOrdinal, originalOrdinal)); + } - // Write neighbors - writer.reset(); - var neighbors = view.getNeighborsIterator(0, originalOrdinal); - if (neighbors.size() > graph.getDegree(0)) { - throw new IllegalStateException( - String.format("Node %d has more neighbors %d than the graph's max degree %d -- run Builder.cleanup()!", - originalOrdinal, neighbors.size(), graph.getDegree(0))); - } + // Ordinal: always owned. + writer.writeInt(newOrdinal); - writer.writeInt(neighbors.size()); - int n = 0; - for (; n < neighbors.size(); n++) { - var newNeighborOrdinal = ordinalMapper.oldToNew(neighbors.nextInt()); - if (newNeighborOrdinal < 0 || newNeighborOrdinal > ordinalMapper.maxOrdinal()) { - throw new IllegalStateException( - String.format("Neighbor ordinal out of bounds: %d/%d", - newNeighborOrdinal, ordinalMapper.maxOrdinal())); + for (var feature : inlineFeatures) { + var supplier = featureStateSuppliers.get(feature.id()); + if (supplier != null) { + // Owned: extend the current run directly. + feature.writeInline(writer, supplier.apply(originalOrdinal)); + } else { + // Pre-written gap: flush the run accumulated so far (which may span + // multiple earlier nodes), then skip the gap in the file without + // writing anything into rangeBuffer for it. + int runEnd = rangeBuffer.position(); + if (runEnd > runStart) { + pending.add(new PendingWrite(sliceRange(rangeBuffer, runStart, runEnd), runFilePosition)); } - writer.writeInt(newNeighborOrdinal); + runFilePosition += (runEnd - runStart) + feature.featureSize(); + runStart = runEnd; } + } - // Pad to max degree - for (; n < graph.getDegree(0); n++) { - writer.writeInt(-1); + // Neighbor section: always owned — extends the current run. + var neighbors = view.getNeighborsIterator(0, originalOrdinal); + if (neighbors.size() > graph.getDegree(0)) { + throw new IllegalStateException(String.format( + "Node %d has more neighbors %d than max degree %d -- run Builder.cleanup()!", + originalOrdinal, neighbors.size(), graph.getDegree(0))); + } + writer.writeInt(neighbors.size()); + int n = 0; + for (; n < neighbors.size(); n++) { + int newNeighbor = ordinalMapper.oldToNew(neighbors.nextInt()); + if (newNeighbor < 0 || newNeighbor > ordinalMapper.maxOrdinal()) { + throw new IllegalStateException(String.format( + "Neighbor ordinal out of bounds: %d/%d", + newNeighbor, ordinalMapper.maxOrdinal())); } + writer.writeInt(newNeighbor); + } + for (; n < graph.getDegree(0); n++) writer.writeInt(-1); + } + + // Final trailing run. + int runEnd = rangeBuffer.position(); + if (runEnd > runStart) { + pending.add(new PendingWrite(sliceRange(rangeBuffer, runStart, runEnd), runFilePosition)); + } - ByteBuffer neighborData = writer.cloneBuffer(); - writeFully(channel, neighborData, currentPosition); + writeAllFully(pending); + } + + /** + * Returns an independent, ready-to-read view over {@code buf}'s backing storage covering + * {@code [start, end)}. The view shares memory with {@code buf} but has its own position + * and limit, so later writes into {@code buf} at other offsets don't affect it — safe to + * hand off to a concurrent {@code channel.write()} while {@code buf} keeps being appended to. + */ + private static ByteBuffer sliceRange(ByteBuffer buf, int start, int end) { + ByteBuffer dup = buf.duplicate(); + dup.limit(end); + dup.position(start); + return dup.slice(); + } + + // ------------------------------------------------------------------------- + // Shared helpers + // ------------------------------------------------------------------------- + + /** + * Writes a complete node record (ordinal + all features + neighbors) sequentially + * into {@code writer}. Called only from {@link #callBatched()}, where all feature + * suppliers are guaranteed non-null. + */ + private void buildFullRecord(ByteBufferIndexWriter writer, int newOrdinal) throws Exception { + var originalOrdinal = ordinalMapper.newToOld(newOrdinal); + 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(String.format( + "Ordinal mapper mapped new ordinal %d to non-existing node %d", + newOrdinal, 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(String.format( + "Node %d has more neighbors %d than max degree %d -- run Builder.cleanup()!", + originalOrdinal, neighbors.size(), graph.getDegree(0))); + } + writer.writeInt(neighbors.size()); + int n = 0; + for (; n < neighbors.size(); n++) { + int newNeighbor = ordinalMapper.oldToNew(neighbors.nextInt()); + if (newNeighbor < 0 || newNeighbor > ordinalMapper.maxOrdinal()) { + throw new IllegalStateException(String.format( + "Neighbor ordinal out of bounds: %d/%d", + newNeighbor, ordinalMapper.maxOrdinal())); + } + writer.writeInt(newNeighbor); + } + for (; n < graph.getDegree(0); n++) writer.writeInt(-1); } + } - return null; + /** A not-yet-fully-written buffer destined for a fixed file offset. */ + private static final class PendingWrite { + final ByteBuffer buffer; + final long position; + + PendingWrite(ByteBuffer buffer, long position) { + this.buffer = buffer; + this.position = position; + } } -} + /** + * Caps how many writes this task submits to the channel before joining any of them. + *

+ * On platforms without a native async file-I/O backend wired into NIO2 (macOS and most + * POSIX systems get {@code sun.nio.ch.SimpleAsynchronousFileChannelImpl}; only Windows gets + * true IOCP), {@code AsynchronousFileChannel.write()} is emulated by handing the write to an + * executor that is observed to spin up a fresh native thread per outstanding call rather + * than reusing a small, bounded pool. {@link #callLegacy()}'s per-task write count can be + * O(nodes) — unlike {@link #callBatched()}'s O(1) — so submitting an entire task's writes + * before joining any of them can create thousands of simultaneously outstanding writes per + * task, times however many tasks are running concurrently. In practice this has been + * observed to exhaust the OS thread limit ({@code OutOfMemoryError: unable to create native + * thread}, {@code pthread_create failed (EAGAIN)}) on a large write. This bound keeps the + * number of writes any single task has outstanding at once small and constant, regardless + * of how many nodes the task covers. + */ + private static final int MAX_IN_FLIGHT_WRITES = 32; + + /** + * Submits {@code wave} to the channel in chunks of at most {@link #MAX_IN_FLIGHT_WRITES}, + * joining each chunk before submitting the next. Within a chunk, every write is submitted + * before any is joined, preserving pipelining at a bounded scale; see + * {@link #MAX_IN_FLIGHT_WRITES} for why an unbounded submit-everything-first approach is + * unsafe on some platforms. + */ + private void writeAllFully(List wave) throws Exception { + for (int chunkStart = 0; chunkStart < wave.size(); chunkStart += MAX_IN_FLIGHT_WRITES) { + int chunkEnd = Math.min(chunkStart + MAX_IN_FLIGHT_WRITES, wave.size()); + writeChunkFully(wave.subList(chunkStart, chunkEnd)); + } + } + + /** + * Submits every write in {@code chunk} without blocking, then joins them all. + * {@code AsynchronousFileChannel.write()} is only guaranteed to write "up to" the + * buffer's remaining bytes in a single call (see + * {@link AsynchronousFileChannel#write(ByteBuffer, long)}); a short write is rare for a + * regular file but can happen (disk full, quota limits, an + * oversized single write). Any write that completes short is resubmitted for its + * remaining bytes — which the channel has already advanced the buffer's position past — + * as a follow-up chunk. + *

+ * In the common case (no short writes) this runs exactly one submit-all/join-all round, + * so the pipelining {@link #callLegacy()} and {@link #callBatched()} rely on is preserved; + * the retry loop only does extra work on the rare short-write path. + */ + private void writeChunkFully(List chunk) throws Exception { + while (!chunk.isEmpty()) { + List> futures = new ArrayList<>(chunk.size()); + for (var w : chunk) { + futures.add(channel.write(w.buffer, w.position)); + } + + List retry = new ArrayList<>(); + for (int i = 0; i < chunk.size(); i++) { + int written = futures.get(i).get(); + if (written < 0) { + throw new IOException("Channel closed during write"); + } + var w = chunk.get(i); + if (w.buffer.hasRemaining()) { + if (written == 0) { + throw new IOException("Channel made no progress writing at position " + w.position); + } + retry.add(new PendingWrite(w.buffer, w.position + written)); + } + } + chunk = retry; + } + } +} 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..bc6cd3d2c 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 @@ -22,7 +22,6 @@ import io.github.jbellis.jvector.graph.disk.feature.FeatureId; import java.io.IOException; -import java.nio.ByteBuffer; import java.nio.channels.AsynchronousFileChannel; import java.nio.file.Path; import java.nio.file.StandardOpenOption; @@ -67,11 +66,14 @@ class ParallelGraphWriter implements AutoCloseable { private final ExecutorService executor; private final boolean ownsExecutor; private final ThreadLocal viewPerThread; - private final ThreadLocal bufferPerThread; private final CopyOnWriteArrayList allViews = new CopyOnWriteArrayList<>(); private final int recordSize; private final Path filePath; private final int taskMultiplier; + // Passed through to NodeRecordTask so each task allocates the right buffer type. + // FUTURE IMPROVEMENT: when the legacy path is removed, each task allocates exactly + // one range-sized buffer, making the direct/heap distinction more impactful and cleaner. + private final boolean useDirectBuffers; private static final AtomicInteger threadCounter = new AtomicInteger(0); /** @@ -143,6 +145,7 @@ public ParallelGraphWriter(RandomAccessWriter writer, this.graph = graph; this.filePath = Objects.requireNonNull(filePath); this.taskMultiplier = config.taskMultiplier; + this.useDirectBuffers = config.useDirectBuffers; if (externalExecutor != null) { this.executor = externalExecutor; this.ownsExecutor = false; @@ -163,23 +166,17 @@ public ParallelGraphWriter(RandomAccessWriter writer, + Integer.BYTES // neighbor count + graph.getDegree(0) * Integer.BYTES; // neighbors + padding - // Thread-local views for safe neighbor iteration - // CopyOnWriteArrayList handles concurrent additions safely + // Thread-local views for safe neighbor iteration. + // CopyOnWriteArrayList handles concurrent additions safely. this.viewPerThread = ThreadLocal.withInitial(() -> { var view = graph.getView(); allViews.add(view); return view; }); - - // Thread-local buffers to avoid allocation overhead - // Use BIG_ENDIAN to match Java DataOutput specification - final int bufferSize = recordSize; - final boolean useDirect = config.useDirectBuffers; - this.bufferPerThread = ThreadLocal.withInitial(() -> { - ByteBuffer buffer = useDirect ? ByteBuffer.allocateDirect(bufferSize) : ByteBuffer.allocate(bufferSize); - buffer.order(java.nio.ByteOrder.BIG_ENDIAN); - return buffer; - }); + // FUTURE IMPROVEMENT: the old per-thread scratch buffer (bufferPerThread) is removed. + // Each task now allocates its own range-sized buffer in the fast path, or per-region + // buffers in the legacy path. The thread-local was sized to a single record and forced + // sub-record writes; the new approach eliminates that bottleneck entirely. } /** @@ -239,23 +236,19 @@ public void writeL0Records(OrdinalMapper ordinalMapper, final int end = endOrdinal; Future future = executor.submit(() -> { - var view = viewPerThread.get(); - var buffer = bufferPerThread.get(); - var task = new NodeRecordTask( - start, // Start of range (inclusive) - end, // End of range (exclusive) + start, // range start (inclusive) + end, // range end (exclusive) ordinalMapper, graph, - view, + viewPerThread.get(), inlineFeatures, featureStateSuppliers, recordSize, - baseOffset, // Base offset (task calculates per-ordinal offsets) - channel, // Async file channel for position-based writes - buffer // Thread-local buffer + baseOffset, + channel, + useDirectBuffers // each task allocates its own buffer(s) ); - return task.call(); }); 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 24c39ae47..ef1d36b89 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 @@ -148,15 +148,16 @@ public static void main(String[] args) throws IOException { List datasetResults = Grid.runAllAndCollectResults(ds, config.construction.useSavedIndexIfExists, - config.construction.outDegree, + config.construction.isParallelGraphConstruction(), + config.construction.outDegree, config.construction.efConstruction, - config.construction.neighborOverflow, + config.construction.neighborOverflow, config.construction.addHierarchy, config.construction.refineFinalGraph, - config.construction.getFeatureSets(), + config.construction.getFeatureSets(), config.construction.getCompressorParameters(), - config.search.getCompressorParameters(), - config.search.topKOverquery, + config.search.getCompressorParameters(), + config.search.topKOverquery, config.search.useSearchPruning); results.addAll(datasetResults); 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 e066a34dc..1ae5bb439 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 @@ -126,6 +126,7 @@ public static void main(String[] args) throws IOException { Grid.runAll(ds, config.construction.useSavedIndexIfExists, + config.construction.isParallelGraphConstruction(), config.construction.outDegree, config.construction.efConstruction, config.construction.neighborOverflow, 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 8f45df2a0..b9287fffe 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 @@ -97,6 +97,7 @@ public static Double getIndexBuildTimeSeconds(String datasetName) { static void runAll(DataSet ds, boolean enableIndexCache, + boolean useParallelConstruction, List mGrid, List efConstructionGrid, List neighborOverflowGrid, @@ -129,7 +130,7 @@ static void runAll(DataSet ds, for (int efC : efConstructionGrid) { for (var bc : buildCompressors) { runOneGraph(cache, featureSets, M, efC, neighborOverflow, addHierarchy, refineFinalGraph, - bc, compressionGrid, topKGrid, usePruningGrid, artifacts, ds, workDir); + bc, compressionGrid, topKGrid, usePruningGrid, artifacts, ds, workDir, useParallelConstruction); } } } @@ -174,6 +175,7 @@ static void runAll(DataSet ds, { runAll(ds, enableIndexCache, + false, // legacy callers always use serial construction mGrid, efConstructionGrid, neighborOverflowGrid, @@ -225,7 +227,8 @@ static void runOneGraph(OnDiskGraphIndexCache cache, List usePruningGrid, RunArtifacts artifacts, DataSet ds, - Path workDirectory) throws IOException + Path workDirectory, + boolean useParallelConstruction) throws IOException { // Prepare to collect index construction metrics for reporting.... var constructionMetrics = new ConstructionMetrics(); @@ -288,7 +291,7 @@ static void runOneGraph(OnDiskGraphIndexCache cache, // At least one index needs to be built (b/c not in cache or cache is disabled) // We pass the handles map so buildOnDisk knows exactly where to write var newIndexes = buildOnDisk(missing, M, efConstruction, neighborOverflow, addHierarchy, refineFinalGraph, - ds, outputDir, buildCompressorObj, handles, constructionMetrics); + ds, outputDir, buildCompressorObj, handles, constructionMetrics, useParallelConstruction); indexes.putAll(newIndexes); } } @@ -374,7 +377,8 @@ private static Map, ImmutableGraphIndex> buildOnDisk(List buildCompressor, Map, OnDiskGraphIndexCache.WriteHandle> handles, - ConstructionMetrics constructionMetrics) throws IOException + ConstructionMetrics constructionMetrics, + boolean useParallelConstruction) throws IOException { Files.createDirectories(outputDir); @@ -388,9 +392,9 @@ private static Map, ImmutableGraphIndex> buildOnDisk(List, OnDiskGraphIndexWriter> writers = new HashMap<>(); + Map, RandomAccessOnDiskGraphIndexWriter> writers = new HashMap<>(); Map, Map>> suppliers = new HashMap<>(); - OnDiskGraphIndexWriter scoringWriter = null; + RandomAccessOnDiskGraphIndexWriter scoringWriter = null; int n = 0; for (var features : featureSets) { // if we are using index caching, use cache names instead of tmp names for index files.... @@ -400,8 +404,8 @@ private static Map, ImmutableGraphIndex> buildOnDisk(List features Path outPath, RandomAccessVectorValues floatVectors, ProductQuantization pq, - ConstructionMetrics constructionMetrics) - throws FileNotFoundException + ConstructionMetrics constructionMetrics, + boolean useParallelConstruction) + throws IOException { var identityMapper = new OrdinalMapper.IdentityMapper(floatVectors.size() - 1); - var builder = new OnDiskGraphIndexWriter.Builder(onHeapGraph, outPath); - builder.withMapper(identityMapper); - Map> suppliers = new EnumMap<>(FeatureId.class); + + RandomAccessOnDiskGraphIndexWriter writer; + if (useParallelConstruction) { + var builder = new OnDiskParallelGraphIndexWriter.Builder(onHeapGraph, outPath); + builder.withMapper(identityMapper); + addFeaturesToBuilder(builder, features, onHeapGraph, floatVectors, pq, constructionMetrics, suppliers); + writer = builder.build(); + } else { + var builder = new OnDiskGraphIndexWriter.Builder(onHeapGraph, outPath); + builder.withMapper(identityMapper); + addFeaturesToBuilder(builder, features, onHeapGraph, floatVectors, pq, constructionMetrics, suppliers); + writer = builder.build(); + } + return new BuilderWithSuppliers(writer, suppliers); + } + + private static void addFeaturesToBuilder(AbstractGraphIndexWriter.Builder builder, + Set features, + ImmutableGraphIndex onHeapGraph, + RandomAccessVectorValues floatVectors, + ProductQuantization pq, + ConstructionMetrics constructionMetrics, + Map> suppliers) { for (var featureId : features) { switch (featureId) { case INLINE_VECTORS: @@ -513,10 +538,8 @@ private static BuilderWithSuppliers builderWithSuppliers(Set features builder.with(new NVQ(nvq)); suppliers.put(FeatureId.NVQ_VECTORS, ordinal -> new NVQ.State(nvq.encode(floatVectors.getVector(ordinal)))); break; - } } - return new BuilderWithSuppliers(builder, suppliers); } public static void setDiagnosticLevel(int diagLevel) { @@ -539,11 +562,11 @@ private static DiagnosticLevel getDiagnosticLevel() { } private static class BuilderWithSuppliers { - public final OnDiskGraphIndexWriter.Builder builder; + public final RandomAccessOnDiskGraphIndexWriter writer; public final Map> suppliers; - public BuilderWithSuppliers(OnDiskGraphIndexWriter.Builder builder, Map> suppliers) { - this.builder = builder; + public BuilderWithSuppliers(RandomAccessOnDiskGraphIndexWriter writer, Map> suppliers) { + this.writer = writer; this.suppliers = suppliers; } } @@ -594,8 +617,8 @@ private static Map, ImmutableGraphIndex> buildInMemory(List ordered(Object... kv) { public static List runAllAndCollectResults( DataSet ds, boolean enableIndexCache, + boolean useParallelConstruction, List mGrid, List efConstructionGrid, List neighborOverflowGrid, @@ -904,7 +928,7 @@ public static List runAllAndCollectResults( // At least one index needs to be built (b/c not in cache or cache is disabled) // We pass the handles map so buildOnDisk knows exactly where to write var newIndexes = buildOnDisk(missing, m, ef, neighborOverflow, addHierarchy, refineFinalGraph, - ds, outputDir, compressor, handles, null); + ds, outputDir, compressor, handles, null, useParallelConstruction); indexes.putAll(newIndexes); } diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/HelloVectorWorld.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/HelloVectorWorld.java index ea4752e4b..9bcf6843e 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/HelloVectorWorld.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/HelloVectorWorld.java @@ -47,6 +47,7 @@ public static void main(String[] args) throws IOException { // Run Grid.runAll(ds, config.construction.useSavedIndexIfExists, + config.construction.isParallelGraphConstruction(), config.construction.outDegree, config.construction.efConstruction, config.construction.neighborOverflow, diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/ConstructionParameters.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/ConstructionParameters.java index 5177fdf4a..24b61f51c 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/ConstructionParameters.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/ConstructionParameters.java @@ -32,6 +32,11 @@ public class ConstructionParameters extends CommonParameters { public List reranking; public List fusedGraph; public Boolean useSavedIndexIfExists; + public Boolean parallelGraphConstruction; + + public boolean isParallelGraphConstruction() { + return Boolean.TRUE.equals(parallelGraphConstruction); + } public List> getFeatureSets() { List> featureSets = null; 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..596bfe7a2 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,35 +31,78 @@ 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.vector.VectorSimilarityFunction; import java.io.IOException; +import java.io.UncheckedIOException; import java.nio.file.Files; import java.nio.file.Path; import java.util.Arrays; import java.util.EnumMap; import java.util.Map; import java.util.function.IntFunction; +import java.util.stream.IntStream; import static io.github.jbellis.jvector.quantization.KMeansPlusPlusClusterer.UNWEIGHTED; /** - * Example demonstrating how to use parallel writes with OnDiskGraphIndexWriter. + * Example demonstrating how to use {@link OnDiskParallelGraphIndexWriter} and comparing its + * two internal write strategies against the sequential {@link OnDiskGraphIndexWriter}. *

- * Usage patterns: + * {@code OnDiskParallelGraphIndexWriter} parallelizes serialization of Level-0 node records + * using {@link java.nio.channels.AsynchronousFileChannel}. Which of its two write paths runs + * (see {@link NodeRecordTask}) is determined entirely by whether every feature supplier is + * still present in the map passed to {@code write()}: * 

- * // Sequential (default):
+ * // Sequential (default) — single-threaded baseline:
  * var writer = new OnDiskGraphIndexWriter.Builder(graph, outputPath)
  *     .with(inlineVectors)
  *     .build();
  * writer.write(featureSuppliers);
  *
- * // Parallel:
- * var writer = new OnDiskGraphIndexWriter.Builder(graph, outputPath)
+ * // Parallel, batched path — every feature supplier is passed directly to write();
+ * // each task packs its whole node range into one buffer and issues one channel write:
+ * var writer = new OnDiskParallelGraphIndexWriter.Builder(graph, outputPath)
  *     .with(inlineVectors)
- *     .withParallelWrites(true)  // Enable parallel writes
  *     .build();
  * writer.write(featureSuppliers);
+ *
+ * // Parallel, legacy path — a plain (non-fused) feature is pre-written per node via
+ * // writeFeaturesInline() and then omitted from the map passed to write(); each task then
+ * // owns only the bytes not already on disk and issues one non-blocking write per owned span:
+ * var writer = new OnDiskParallelGraphIndexWriter.Builder(graph, outputPath)
+ *     .with(plainInlineFeature)
+ *     .build();
+ * for (int ordinal = 0; ordinal < graph.size(0); ordinal++) {
+ *     writer.writeFeaturesInline(ordinal, Map.of(plainInlineFeature.id(), stateFor(ordinal)));
+ * }
+ * writer.write(Map.of());
  *
+ *

+ * Note: a fused feature (one whose {@code Feature.isFused()} is {@code true}, e.g. + * {@code FusedPQ}) can never be fully pre-written this way. {@code write()} always requires a + * supplier for it regardless of what NodeRecordTask needs, because + * {@code AbstractGraphIndexWriter.writeSparseLevels()} writes the fused source feature for the + * hierarchy's higher layers directly — a separate code path from the L0 records NodeRecordTask + * owns, with no {@code writeFeaturesInline()} equivalent. Omitting a fused feature's supplier + * throws {@code IllegalStateException: Supplier for feature ... not found}. The benchmarks below + * demonstrate the correct pattern: pre-write the plain feature (NVQ_VECTORS) and keep supplying + * the fused one (FUSED_PQ) to {@code write()}. + *

+ * Two benchmark methods exercise this: + *

*/ public class ParallelWriteExample { @@ -223,12 +266,19 @@ private static void verifyIndicesIdentical(OnDiskGraphIndex index1, OnDiskGraphI } /** - * Benchmark comparison between sequential and parallel writes using NVQ + FUSED_ADC features. - * This matches the configuration used in Grid.buildOnDisk for realistic performance testing. + * Benchmark comparison on an already-fully-built graph using NVQ + FUSED_ADC features: + * sequential writes via {@link OnDiskGraphIndexWriter} vs. parallel writes via + * {@link OnDiskParallelGraphIndexWriter} taking its batched path (see {@link NodeRecordTask}): + * every feature supplier is passed straight to {@code write()}, so each task builds its + * whole node range into one buffer and issues a single {@code channel.write()} call. + *

+ * Both writers see the graph only after it's completely built — this isolates the cost of + * the write strategy itself from graph construction. Contrast with + * {@link #benchmarkInterleavedWrites}, which measures the legacy/pre-write pattern. */ - public static void benchmarkComparison(ImmutableGraphIndex graph, + public static void benchmarkPlainWrites(ImmutableGraphIndex graph, Path sequentialPath, - Path parallelPath, + Path parallelBatchedPath, RandomAccessVectorValues floatVectors, PQVectors pqVectors) throws IOException { @@ -265,11 +315,13 @@ public static void benchmarkComparison(ImmutableGraphIndex graph, view.close(); } long sequentialTime = System.nanoTime() - sequentialStart; - System.out.printf("Sequential write: %.2f ms%n", sequentialTime / 1_000_000.0); + System.out.printf("Sequential write: %.2f ms%n", sequentialTime / 1_000_000.0); - // Parallel write - long parallelStart = System.nanoTime(); - try (var writer = new OnDiskParallelGraphIndexWriter.Builder(graph, parallelPath) + // Parallel write — batched path: every feature supplier is passed directly to write(), + // so NodeRecordTask.callBatched() packs each task's whole node range into one buffer + // and issues a single channel.write() for it. + long parallelBatchedStart = System.nanoTime(); + try (var writer = new OnDiskParallelGraphIndexWriter.Builder(graph, parallelBatchedPath) .with(nvqFeature) .with(fusedPQFeature) .withMapper(identityMapper) @@ -283,23 +335,124 @@ public static void benchmarkComparison(ImmutableGraphIndex graph, writer.write(writeSuppliers); view.close(); } - long parallelTime = System.nanoTime() - parallelStart; + long parallelBatchedTime = System.nanoTime() - parallelBatchedStart; + System.out.printf("Parallel write (batched): %.2f ms%n", parallelBatchedTime / 1_000_000.0); + System.out.printf("Speedup (batched vs sequential): %.2fx%n", (double) sequentialTime / parallelBatchedTime); + } + + /** + * Benchmark comparison of the legacy/pre-write pattern, built the way {@code Grid.buildOnDisk} + * actually uses it: {@code NVQ_VECTORS} is written via {@code writeFeaturesInline()} + * interleaved with graph construction itself, not as an isolated pass afterward. + *

+ * One {@link GraphIndexBuilder} is shared by both a sequential ({@link OnDiskGraphIndexWriter}) + * and a parallel-legacy ({@link OnDiskParallelGraphIndexWriter}) writer. Both writers are + * registered before construction starts (mirroring {@code Grid.builderWithSuppliers} being + * called before the incremental-build loop), and both receive a {@code writeFeaturesInline()} + * call for every node inside the same parallel construction stream — mirroring + * {@code Grid.buildOnDisk}'s {@code writers.forEach(...)} pattern, which feeds every writer + * regardless of whether it will finalize sequentially or asynchronously. {@code FUSED_PQ} + * can't be pre-written (see the class Javadoc), so it stays in the map passed to the final + * {@code write()} call for both writers. + *

+ * Since both writers are fed from the same construction pass, the graph-build-plus-pre-write + * cost is a single, genuinely shared number — it's added once to each writer's total rather + * than measured or estimated separately. The two {@code write()} finalize calls are timed one + * after the other (not concurrently), so each writer's finalize cost is cleanly attributable + * instead of contending with the other for I/O. + */ + public static void benchmarkInterleavedWrites(RandomAccessVectorValues floatVectors, + VectorSimilarityFunction similarityFunction, + PQVectors pqVectors, + int M, + int efConstruction, + float neighborOverflow, + float alpha, + boolean addHierarchy, + boolean refineFinalGraph, + Path sequentialPath, + Path parallelLegacyPath) throws IOException { + + int nSubVectors = floatVectors.dimension() == 2 ? 1 : 2; + var nvq = NVQuantization.compute(floatVectors, nSubVectors); + var pq = pqVectors.getCompressor(); + var identityMapper = new OrdinalMapper.IdentityMapper(floatVectors.size() - 1); - System.out.printf("Parallel write: %.2f ms%n", parallelTime / 1_000_000.0); - System.out.printf("Speedup: %.2fx%n", (double) sequentialTime / parallelTime); + var bsp = BuildScoreProvider.pqBuildScoreProvider(similarityFunction, pqVectors); + var builder = new GraphIndexBuilder(bsp, floatVectors.dimension(), M, efConstruction, + neighborOverflow, alpha, addHierarchy, refineFinalGraph); + var onHeapGraph = builder.getGraph(); + + try (var sequentialWriter = new OnDiskGraphIndexWriter.Builder(onHeapGraph, sequentialPath) + .with(new NVQ(nvq)) + .with(new FusedPQ(onHeapGraph.maxDegree(), pq)) + .withMapper(identityMapper) + .build(); + var parallelLegacyWriter = new OnDiskParallelGraphIndexWriter.Builder(onHeapGraph, parallelLegacyPath) + .with(new NVQ(nvq)) + .with(new FusedPQ(onHeapGraph.maxDegree(), pq)) + .withMapper(identityMapper) + .build()) { + + // Interleave NVQ pre-writes for BOTH writers into the same construction pass, + // mirroring Grid.buildOnDisk's writers.forEach(...) per-node loop. + System.out.println("Building graph with interleaved feature writes (Grid.buildOnDisk pattern)..."); + long interleavedBuildStart = System.nanoTime(); + var vv = floatVectors.threadLocalSupplier(); + IntStream.range(0, floatVectors.size()).parallel().forEach(node -> { + Feature.State nvqState = new NVQ.State(nvq.encode(floatVectors.getVector(node))); + Map stateMap = Map.of(FeatureId.NVQ_VECTORS, nvqState); + try { + sequentialWriter.writeFeaturesInline(node, stateMap); + parallelLegacyWriter.writeFeaturesInline(node, stateMap); + } catch (IOException e) { + throw new UncheckedIOException(e); + } + builder.addGraphNode(node, vv.get().getVector(node)); + }); + builder.cleanup(); + long interleavedBuildTime = System.nanoTime() - interleavedBuildStart; + System.out.printf("Graph build + interleaved pre-write: %.2f ms%n", interleavedBuildTime / 1_000_000.0); + + // Finalize each writer separately (sequentially, not concurrently) so each one's + // write() time is cleanly attributable rather than contending with the other for I/O. + try (var view = onHeapGraph.getView()) { + long seqWriteStart = System.nanoTime(); + sequentialWriter.write(Map.of(FeatureId.FUSED_PQ, ordinal -> new FusedPQ.State(view, pqVectors, ordinal))); + long seqWriteTime = System.nanoTime() - seqWriteStart; + System.out.printf("Sequential write() (finalize): %.2f ms%n", seqWriteTime / 1_000_000.0); + + long legacyWriteStart = System.nanoTime(); + parallelLegacyWriter.write(Map.of(FeatureId.FUSED_PQ, ordinal -> new FusedPQ.State(view, pqVectors, ordinal))); + long legacyWriteTime = System.nanoTime() - legacyWriteStart; + System.out.printf("Parallel (legacy) write() (finalize): %.2f ms%n", legacyWriteTime / 1_000_000.0); + + long sequentialTotal = interleavedBuildTime + seqWriteTime; + long legacyTotal = interleavedBuildTime + legacyWriteTime; + System.out.printf("%nSequential total (build+prewrite+write): %.2f ms%n", sequentialTotal / 1_000_000.0); + System.out.printf("Parallel (legacy) total (build+prewrite+write): %.2f ms%n", legacyTotal / 1_000_000.0); + System.out.printf("Speedup (legacy finalize vs sequential finalize): %.2fx%n", (double) seqWriteTime / legacyWriteTime); + System.out.printf("Speedup (legacy total vs sequential total): %.2fx%n", (double) sequentialTotal / legacyTotal); + } + } finally { + builder.close(); + } } /** - * Main method to run a benchmark test of sequential vs parallel writes. - * + * Main method to run two benchmark comparisons: sequential vs. parallel-batched writes on + * an already-built graph ({@link #benchmarkPlainWrites}), and sequential vs. parallel-legacy + * writes with feature data interleaved into graph construction, Grid.java-style + * ({@link #benchmarkInterleavedWrites}). + *

* Usage: java ParallelWriteExample [dataset-name] - * - * Example: java ParallelWriteExample cohere-english-v3-100k - * - * If no dataset is provided, uses "cohere-english-v3-100k" by default. + *

+ * Example: java ParallelWriteExample cohere-english-v3-1M + *

+ * If no dataset is provided, uses "cohere-english-v3-1M" by default. */ public static void main(String[] args) throws IOException { - String datasetName = args.length > 0 ? args[0] : "cohere-english-v3-100k"; + String datasetName = args.length > 0 ? args[0] : "cohere-english-v3-1M"; System.out.println("Loading dataset: " + datasetName); DataSet ds = DataSets.loadDataSet(datasetName).orElseThrow( @@ -312,7 +465,7 @@ public static void main(String[] args) throws IOException { // Build PQ compression (matching Grid.buildOnDisk pattern) System.out.println("Computing PQ compression..."); int pqM = floatVectors.dimension() / 8; // m = dimension / 8 - boolean centerData = ds.getSimilarityFunction() == io.github.jbellis.jvector.vector.VectorSimilarityFunction.EUCLIDEAN; + boolean centerData = ds.getSimilarityFunction() == VectorSimilarityFunction.EUCLIDEAN; var pq = ProductQuantization.compute(floatVectors, pqM, 256, centerData, UNWEIGHTED); var pqVectors = (PQVectors) pq.encodeAll(floatVectors); System.out.printf("PQ compression: %d subspaces, 256 clusters%n", pqM); @@ -325,59 +478,67 @@ public static void main(String[] args) throws IOException { boolean addHierarchy = true; boolean refineFinalGraph = true; - System.out.printf("Building graph with PQ-compressed vectors (M=%d, efConstruction=%d)...%n", M, efConstruction); - long buildStart = System.nanoTime(); - - var bsp = BuildScoreProvider.pqBuildScoreProvider(ds.getSimilarityFunction(), pqVectors); - var builder = new GraphIndexBuilder(bsp, floatVectors.dimension(), M, efConstruction, - neighborOverflow, alpha, addHierarchy, refineFinalGraph); - - // Build graph using parallel construction for much better performance - var graph = builder.build(floatVectors); - long buildTime = System.nanoTime() - buildStart; - System.out.printf("Graph built in %.2fs%n", buildTime / 1_000_000_000.0); - System.out.printf("Graph has %d nodes%n", graph.size(0)); - - // Create temporary paths for writing Path tempDir = Files.createTempDirectory("parallel-write-test"); Path sequentialPath = tempDir.resolve("graph-sequential"); - Path parallelPath = tempDir.resolve("graph-parallel"); + Path parallelBatchedPath = tempDir.resolve("graph-parallel-batched"); + Path sequentialInterleavedPath = tempDir.resolve("graph-sequential-interleaved"); + Path parallelLegacyInterleavedPath = tempDir.resolve("graph-parallel-legacy-interleaved"); try { - System.out.println("\n=== Testing Write Performance ==="); + // === Graph A: built once, plain (no writer involved during construction) === + System.out.printf("%nBuilding Graph A (plain) with PQ-compressed vectors (M=%d, efConstruction=%d)...%n", M, efConstruction); + long buildStart = System.nanoTime(); + var bsp = BuildScoreProvider.pqBuildScoreProvider(ds.getSimilarityFunction(), pqVectors); + var graphABuilder = new GraphIndexBuilder(bsp, floatVectors.dimension(), M, efConstruction, + neighborOverflow, alpha, addHierarchy, refineFinalGraph); + var graphA = graphABuilder.build(floatVectors); + long buildTime = System.nanoTime() - buildStart; + System.out.printf("Graph A built in %.2fs (%d nodes)%n", buildTime / 1_000_000_000.0, graphA.size(0)); + graphABuilder.close(); + + System.out.println("\n=== Graph A: sequential vs. parallel (batched) ==="); + benchmarkPlainWrites(graphA, sequentialPath, parallelBatchedPath, floatVectors, pqVectors); - // Run benchmark comparison - benchmarkComparison(graph, sequentialPath, parallelPath, floatVectors, pqVectors); - - // Report file sizes long seqSize = Files.size(sequentialPath); - long parSize = Files.size(parallelPath); - System.out.printf("%nFile sizes: Sequential=%.2f MB, Parallel=%.2f MB%n", - seqSize / 1024.0 / 1024.0, - parSize / 1024.0 / 1024.0); - - // === Read Phase: Load and verify both indices === - System.out.println("\n=== Testing Read Correctness ==="); - System.out.println("Loading sequential index..."); - OnDiskGraphIndex sequentialIndex = OnDiskGraphIndex.load(ReaderSupplierFactory.open(sequentialPath)); - System.out.println("Loading parallel index..."); - OnDiskGraphIndex parallelIndex = OnDiskGraphIndex.load(ReaderSupplierFactory.open(parallelPath)); - - // Verify that both indices are identical - verifyIndicesIdentical(sequentialIndex, parallelIndex); + long batchedSize = Files.size(parallelBatchedPath); + System.out.printf("%nFile sizes: Sequential=%.2f MB, Parallel(batched)=%.2f MB%n", + seqSize / 1024.0 / 1024.0, batchedSize / 1024.0 / 1024.0); + + System.out.println("\n=== Testing Read Correctness (Graph A) ==="); + try (var sequentialIndex = OnDiskGraphIndex.load(ReaderSupplierFactory.open(sequentialPath)); + var parallelBatchedIndex = OnDiskGraphIndex.load(ReaderSupplierFactory.open(parallelBatchedPath))) { + verifyIndicesIdentical(sequentialIndex, parallelBatchedIndex); + } - // Close the loaded indices - sequentialIndex.close(); - parallelIndex.close(); + // === Graph B: built once via feature writes interleaved into construction, fed to + // both writers from the same pass (Grid.buildOnDisk pattern) === + System.out.println("\n=== Graph B: sequential vs. parallel (legacy), interleaved pre-write ==="); + benchmarkInterleavedWrites(floatVectors, ds.getSimilarityFunction(), pqVectors, + M, efConstruction, neighborOverflow, alpha, addHierarchy, refineFinalGraph, + sequentialInterleavedPath, parallelLegacyInterleavedPath); + + long seqInterleavedSize = Files.size(sequentialInterleavedPath); + long legacyInterleavedSize = Files.size(parallelLegacyInterleavedPath); + System.out.printf("%nFile sizes: Sequential(interleaved)=%.2f MB, Parallel(legacy,interleaved)=%.2f MB%n", + seqInterleavedSize / 1024.0 / 1024.0, legacyInterleavedSize / 1024.0 / 1024.0); + + // Both Graph B writers serialized the same shared graph object (see + // benchmarkInterleavedWrites), so the strict structural check applies here too. + System.out.println("\n=== Testing Read Correctness (Graph B) ==="); + try (var sequentialInterleavedIndex = OnDiskGraphIndex.load(ReaderSupplierFactory.open(sequentialInterleavedPath)); + var parallelLegacyInterleavedIndex = OnDiskGraphIndex.load(ReaderSupplierFactory.open(parallelLegacyInterleavedPath))) { + verifyIndicesIdentical(sequentialInterleavedIndex, parallelLegacyInterleavedIndex); + } } finally { // Cleanup - builder.close(); Files.deleteIfExists(sequentialPath); - Files.deleteIfExists(parallelPath); + Files.deleteIfExists(parallelBatchedPath); + Files.deleteIfExists(sequentialInterleavedPath); + Files.deleteIfExists(parallelLegacyInterleavedPath); Files.deleteIfExists(tempDir); } - System.out.println("\n✅ Test complete - sequential and parallel writes produce identical results!"); + System.out.println("\n✅ Test complete - sequential and parallel (batched and legacy) writes produce identical results!"); } } diff --git a/jvector-examples/yaml-configs/index-parameters/cap-10M.yml b/jvector-examples/yaml-configs/index-parameters/cap-6M.yml similarity index 98% rename from jvector-examples/yaml-configs/index-parameters/cap-10M.yml rename to jvector-examples/yaml-configs/index-parameters/cap-6M.yml index c4c285d18..98157a7d4 100644 --- a/jvector-examples/yaml-configs/index-parameters/cap-10M.yml +++ b/jvector-examples/yaml-configs/index-parameters/cap-6M.yml @@ -1,7 +1,7 @@ yamlSchemaVersion: 1 onDiskIndexVersion: 6 -dataset: cap-10M +dataset: cap-6M construction: outDegree: [32] diff --git a/jvector-examples/yaml-configs/index-parameters/default.yml b/jvector-examples/yaml-configs/index-parameters/default.yml index b56e27ed0..a82d80191 100644 --- a/jvector-examples/yaml-configs/index-parameters/default.yml +++ b/jvector-examples/yaml-configs/index-parameters/default.yml @@ -21,6 +21,7 @@ construction: reranking: - NVQ useSavedIndexIfExists: No # If yes, the system will attempt to locate a cached copy of the index instead of constructing + parallelGraphConstruction: No # If yes, uses OnDiskParallelGraphIndexWriter for L0 parallel writes (experimental) search: topKOverquery: diff --git a/jvector-examples/yaml-configs/index-parameters/dpr-gemma-10M.yml b/jvector-examples/yaml-configs/index-parameters/dpr-gemma-10M.yml index f2370a063..52a65f3d7 100644 --- a/jvector-examples/yaml-configs/index-parameters/dpr-gemma-10M.yml +++ b/jvector-examples/yaml-configs/index-parameters/dpr-gemma-10M.yml @@ -1,7 +1,7 @@ yamlSchemaVersion: 1 onDiskIndexVersion: 6 -dataset: dpr-gemma-10M +dataset: dpr-gemma-10m construction: outDegree: [32] diff --git a/jvector-examples/yaml-configs/index-parameters/dpr-gemma-1M.yml b/jvector-examples/yaml-configs/index-parameters/dpr-gemma-1M.yml index 96e92556e..2b9c1f2c1 100644 --- a/jvector-examples/yaml-configs/index-parameters/dpr-gemma-1M.yml +++ b/jvector-examples/yaml-configs/index-parameters/dpr-gemma-1M.yml @@ -1,7 +1,7 @@ yamlSchemaVersion: 1 onDiskIndexVersion: 6 -dataset: dpr-gemma-1M +dataset: dpr-gemma-1m construction: outDegree: [32] @@ -21,6 +21,7 @@ construction: reranking: - NVQ useSavedIndexIfExists: No + parallelGraphConstruction: Yes search: topKOverquery: