From c4bb0174d88c8bee6bfcea1508dd5275d1201a66 Mon Sep 17 00:00:00 2001 From: Mark Wolters Date: Fri, 12 Jun 2026 17:03:37 -0400 Subject: [PATCH 1/7] make parallel writes truly asyn --- .../jvector/graph/disk/NodeRecordTask.java | 370 ++++++++++++------ .../graph/disk/ParallelGraphWriter.java | 41 +- 2 files changed, 269 insertions(+), 142 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 7e756dab9..ccb9b6e5e 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,57 @@ 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. Each node's owned byte spans (ordinal, non-null features, + * neighbor section) are identified and written as a minimal set of non-blocking channel writes. + * 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()}, {@link #collectNodeWrites}) + * 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 +86,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 +97,240 @@ 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(); + @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; + } + + // ------------------------------------------------------------------------- + // 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++) { + buildFullRecord(writer, newOrdinal); + } + + // One channel.write() for the entire task range — one syscall, one OS I/O request. + rangeBuffer.flip(); + channel.write(rangeBuffer, baseOffset + (long) startOrdinal * recordSize).get(); + } + + // ------------------------------------------------------------------------- + // Legacy path: handle pre-written feature regions. + // + // Pre-written bytes must not be overwritten. For each node we identify the + // contiguous owned spans (ordinal + consecutive non-null features + neighbor + // section) and fire one non-blocking channel write per span. ALL writes for + // the entire task are submitted before any Future.get() call, letting the OS + // pipeline them. + // + // FUTURE IMPROVEMENT: delete this method and collectNodeWrites() entirely once + // writeFeaturesInline support is removed. The fast path handles everything. + // ------------------------------------------------------------------------- + + private void callLegacy() throws Exception { + List> pendingWrites = new ArrayList<>(); + + for (int newOrdinal = startOrdinal; newOrdinal < endOrdinal; newOrdinal++) { + collectNodeWrites(newOrdinal, pendingWrites); + } + + // All writes are in-flight; wait once rather than blocking after each one. + for (var f : pendingWrites) { + int written = f.get(); if (written < 0) { - throw new IOException("Channel closed while writing"); + throw new IOException("Channel closed during write"); } - currentPosition += written; } } - @Override - public Void call() throws Exception { - // Reuse writer and buffer across all ordinals in this range - var writer = new ByteBufferIndexWriter(buffer); + /** + * Determines the owned byte spans for one node record and appends non-blocking + * channel writes for each span to {@code pending}. + *

+ * The record layout is: + * 

+     *   [ordinal: 4] [feature_0: F0] [feature_1: F1] ... [neighbor_count: 4] [neighbors: D×4]
+     *
+ * The ordinal and neighbor section are always owned. Features are owned when their + * supplier is non-null; null-supplier features are pre-written gaps we skip over. + * Consecutive owned bytes are merged into a single write to minimise I/O calls. + *

+ * FUTURE IMPROVEMENT: when writeFeaturesInline is removed, null suppliers cannot exist. + * This method disappears and every node is handled by buildFullRecord() in callBatched(). + */ + private void collectNodeWrites(int newOrdinal, List> pending) throws Exception { + var originalOrdinal = ordinalMapper.newToOld(newOrdinal); + long recordBase = baseOffset + (long) newOrdinal * recordSize; - for (int newOrdinal = startOrdinal; newOrdinal < endOrdinal; newOrdinal++) { - var originalOrdinal = ordinalMapper.newToOld(newOrdinal); - long recordOffset = baseOffset + (long) newOrdinal * recordSize; - long currentPosition = recordOffset; - - // Reset buffer for this ordinal - writer.reset(); - - // Write node ordinal - writer.writeInt(newOrdinal); - ByteBuffer ordinalData = writer.cloneBuffer(); - writeFully(channel, ordinalData, currentPosition); - currentPosition += Integer.BYTES; - - // Handle OMITTED nodes (holes in ordinal space) - if (originalOrdinal == OrdinalMapper.OMITTED) { - // Write placeholder: zeros for features and empty neighbor list - writer.reset(); - for (var feature : inlineFeatures) { - // Write zeros for missing features - 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); + if (originalOrdinal == OrdinalMapper.OMITTED) { + // OMITTED nodes are holes in the ordinal space. writeFeaturesInline() is never + // called for them, so the entire record is ours — write it in one shot. + ByteBuffer buf = allocRegion(recordSize); + buf.putInt(newOrdinal); + for (var feature : inlineFeatures) { + for (int i = 0; i < feature.featureSize(); i++) buf.put((byte) 0); + } + buf.putInt(0); // neighbor count + for (int n = 0; n < graph.getDegree(0); n++) buf.putInt(-1); + buf.flip(); + pending.add(channel.write(buf, recordBase)); + return; + } + + if (!graph.containsNode(originalOrdinal)) { + throw new IllegalStateException(String.format( + "Ordinal mapper mapped new ordinal %d to non-existing node %d", + newOrdinal, originalOrdinal)); + } + + // Accumulate contiguous owned bytes into a region buffer. When a null supplier + // (pre-written gap) is encountered, flush the current region and start a new one + // positioned past the gap. + // + // FUTURE IMPROVEMENT: this entire accumulation loop simplifies to a straight + // linear write through all features when null suppliers are impossible. + ByteBuffer region = allocRegion(recordSize); + long regionStart = recordBase; + + // Ordinal: always owned. + region.putInt(newOrdinal); + + for (var feature : inlineFeatures) { + var supplier = featureStateSuppliers.get(feature.id()); + if (supplier != null) { + // Owned: extend the current region directly. + // ByteBufferIndexWriter(region, false) picks up at region.position() + // without clearing, so it appends to whatever we have already written. + var fw = new ByteBufferIndexWriter(region, false); + feature.writeInline(fw, supplier.apply(originalOrdinal)); } 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)); + // Pre-written gap: flush the current region (if any bytes are pending) + // then advance regionStart past both the flushed bytes and the gap. + // + // FUTURE IMPROVEMENT: this branch is only needed because writeFeaturesInline + // can place data at arbitrary sub-record offsets. Remove it with that API. + if (region.position() > 0) { + int owned = region.position(); + region.flip(); + pending.add(channel.write(region, regionStart)); + regionStart += owned + feature.featureSize(); + } else { + // Nothing accumulated yet in this region — just skip the gap. + regionStart += feature.featureSize(); } + region = allocRegion(recordSize); + } + } - // 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(); - } + // Neighbor section: always owned — append to the current region. + 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))); + } + var nw = new ByteBufferIndexWriter(region, false); + nw.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())); + } + nw.writeInt(newNeighbor); + } + for (; n < graph.getDegree(0); n++) nw.writeInt(-1); - // 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))); - } + if (region.position() > 0) { + region.flip(); + pending.add(channel.write(region, regionStart)); + } + } - 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())); - } - writer.writeInt(newNeighborOrdinal); - } + // ------------------------------------------------------------------------- + // Shared helpers + // ------------------------------------------------------------------------- - // Pad to max degree - for (; n < graph.getDegree(0); n++) { - writer.writeInt(-1); - } + /** + * 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); - ByteBuffer neighborData = writer.cloneBuffer(); - writeFully(channel, neighborData, currentPosition); + 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; + /** Allocates a fresh BIG_ENDIAN ByteBuffer for a write region. */ + private ByteBuffer allocRegion(int capacity) { + var buf = useDirectBuffers + ? ByteBuffer.allocateDirect(capacity) + : ByteBuffer.allocate(capacity); + buf.order(java.nio.ByteOrder.BIG_ENDIAN); + return buf; } } - 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 230c290da..44f838935 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; @@ -66,11 +65,14 @@ class ParallelGraphWriter implements AutoCloseable { private final ImmutableGraphIndex graph; private final ExecutorService executor; 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); /** @@ -122,6 +124,7 @@ public ParallelGraphWriter(RandomAccessWriter writer, this.graph = graph; this.filePath = Objects.requireNonNull(filePath); this.taskMultiplier = config.taskMultiplier; + this.useDirectBuffers = config.useDirectBuffers; this.executor = Executors.newFixedThreadPool(config.workerThreads, r -> { Thread t = new Thread(r); @@ -136,23 +139,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. } /** @@ -212,23 +209,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(); }); From 26bc437e1cace9d9042a219d858796343dba3ffd Mon Sep 17 00:00:00 2001 From: Mark Wolters Date: Thu, 9 Jul 2026 10:43:30 -0400 Subject: [PATCH 2/7] merge fix --- .../jbellis/jvector/graph/disk/ParallelGraphWriter.java | 7 ------- 1 file changed, 7 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 992a3fb38..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 @@ -146,13 +146,6 @@ public ParallelGraphWriter(RandomAccessWriter writer, this.filePath = Objects.requireNonNull(filePath); this.taskMultiplier = config.taskMultiplier; this.useDirectBuffers = config.useDirectBuffers; - this.executor = Executors.newFixedThreadPool(config.workerThreads, - r -> { - Thread t = new Thread(r); - t.setName("ParallelGraphWriter-Worker-" + threadCounter.getAndIncrement()); - t.setDaemon(false); - return t; - }); if (externalExecutor != null) { this.executor = externalExecutor; this.ownsExecutor = false; From 219aa730243a6bf353fbb1aba0f13fa5ab1d3225 Mon Sep 17 00:00:00 2001 From: Mark Wolters Date: Thu, 9 Jul 2026 17:00:03 -0400 Subject: [PATCH 3/7] make parallel writes configurable option for bench tests --- .../jvector/example/AutoBenchYAML.java | 11 ++-- .../jbellis/jvector/example/BenchYAML.java | 1 + .../github/jbellis/jvector/example/Grid.java | 66 +++++++++++++------ .../jvector/example/HelloVectorWorld.java | 1 + .../example/yaml/ConstructionParameters.java | 5 ++ .../{cap-10M.yml => cap-6M.yml} | 2 +- .../yaml-configs/index-parameters/default.yml | 1 + .../index-parameters/dpr-gemma-10M.yml | 2 +- .../index-parameters/dpr-gemma-1M.yml | 3 +- 9 files changed, 63 insertions(+), 29 deletions(-) rename jvector-examples/yaml-configs/index-parameters/{cap-10M.yml => cap-6M.yml} (98%) 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/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: From 76e3edb6beab5762ef8f0347fd343a1d9bf90611 Mon Sep 17 00:00:00 2001 From: Mark Wolters Date: Wed, 15 Jul 2026 14:34:27 -0400 Subject: [PATCH 4/7] adding release notes --- docs/release notes/4.1.0/695.performance.md | 86 +++++++++++++++++++++ 1 file changed, 86 insertions(+) create mode 100644 docs/release notes/4.1.0/695.performance.md 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..79cced200 --- /dev/null +++ b/docs/release notes/4.1.0/695.performance.md @@ -0,0 +1,86 @@ +### Parallel Graph Index Writing via AsynchronousFileChannel + +**Description** +Introduces `OnDiskParallelGraphIndexWriter`, a new writer that parallelizes the serialization +of Level-0 (L0) node records using Java's `AsynchronousFileChannel`. In the standard +`OnDiskGraphIndexWriter`, all L0 records are written sequentially by a single thread. At +large scale — tens of millions of vectors or high-dimensional embeddings — this sequential +write step becomes the dominant bottleneck after graph construction, negating much of the +benefit of parallel HNSW build. + +The parallel writer divides the L0 ordinal space into tasks, each covering a contiguous range +of nodes. Two write strategies are used depending on whether inline features have been +pre-written via `writeInline()`: + +- **Fast path** (no pre-written features): all records in the range are packed into a single + contiguous `ByteBuffer` and committed with one `channel.write()` call, minimizing syscall + overhead. +- **Legacy path** (features pre-written): each node's non-pre-written byte spans (ordinal, + neighbor section) are identified and issued as individual non-blocking async writes. All + futures for the task are submitted before any are joined, allowing the OS to schedule the + full I/O workload efficiently. + +All tasks are submitted to a configurable thread pool before results are collected; the +`BufferedRandomAccessWriter`'s cache is invalidated after parallel writes complete to prevent +stale buffered reads. + +**How to Enable** + +*Programmatic API* — substitute `OnDiskParallelGraphIndexWriter.Builder` for +`OnDiskGraphIndexWriter.Builder` when writing the index to disk: + +```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* — 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** +- `OnDiskParallelGraphIndexWriter` is annotated `@Experimental`. Its API may change in a + future release. +- The writer produces output in the same on-disk format as `OnDiskGraphIndexWriter`; indexes + written with either class are interchangeable and loaded with `OnDiskGraphIndex.load()`. +- 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()`. +- The per-thread scratch `ByteBuffer` that previously limited throughput to one record at a + time has been removed. Each task now allocates a range-sized buffer, eliminating the + sub-record write bottleneck at high core counts. +- Related issue: #695. From 51f573229b7dfc8bbaf8996a86ee37fd0bad993a Mon Sep 17 00:00:00 2001 From: Mark Wolters Date: Thu, 13 Aug 2026 13:46:10 -0400 Subject: [PATCH 5/7] updating release notes to correct erroneous identification as new feature --- docs/release notes/4.1.0/695.performance.md | 72 ++++++++++++--------- 1 file changed, 42 insertions(+), 30 deletions(-) diff --git a/docs/release notes/4.1.0/695.performance.md b/docs/release notes/4.1.0/695.performance.md index 79cced200..87c019168 100644 --- a/docs/release notes/4.1.0/695.performance.md +++ b/docs/release notes/4.1.0/695.performance.md @@ -1,33 +1,50 @@ -### Parallel Graph Index Writing via AsynchronousFileChannel +### Fully Asynchronous Parallel Graph Index Writes **Description** -Introduces `OnDiskParallelGraphIndexWriter`, a new writer that parallelizes the serialization -of Level-0 (L0) node records using Java's `AsynchronousFileChannel`. In the standard -`OnDiskGraphIndexWriter`, all L0 records are written sequentially by a single thread. At -large scale — tens of millions of vectors or high-dimensional embeddings — this sequential -write step becomes the dominant bottleneck after graph construction, negating much of the -benefit of parallel HNSW build. +`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: -The parallel writer divides the L0 ordinal space into tasks, each covering a contiguous range -of nodes. Two write strategies are used depending on whether inline features have been -pre-written via `writeInline()`: +- **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. -- **Fast path** (no pre-written features): all records in the range are packed into a single - contiguous `ByteBuffer` and committed with one `channel.write()` call, minimizing syscall - overhead. -- **Legacy path** (features pre-written): each node's non-pre-written byte spans (ordinal, - neighbor section) are identified and issued as individual non-blocking async writes. All - futures for the task are submitted before any are joined, allowing the OS to schedule the - full I/O workload efficiently. +**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. -All tasks are submitted to a configurable thread pool before results are collected; the -`BufferedRandomAccessWriter`'s cache is invalidated after parallel writes complete to prevent -stale buffered reads. +**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* — substitute `OnDiskParallelGraphIndexWriter.Builder` for -`OnDiskGraphIndexWriter.Builder` when writing the index to disk: +*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 @@ -61,8 +78,8 @@ var writer = new OnDiskParallelGraphIndexWriter.Builder(graph, outputPath) .build(); ``` -*BenchYAML / AutoBenchYAML* — add the following field under `construction` in any -index-parameter YAML config file: +*BenchYAML / AutoBenchYAML* — new in this PR: add the following field under +`construction` in any index-parameter YAML config file: ```yaml construction: @@ -74,13 +91,8 @@ The field is optional and defaults to `No` (serial writes) if omitted, so existi files require no changes. **Notes** -- `OnDiskParallelGraphIndexWriter` is annotated `@Experimental`. Its API may change in a - future release. - 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()`. -- The per-thread scratch `ByteBuffer` that previously limited throughput to one record at a - time has been removed. Each task now allocates a range-sized buffer, eliminating the - sub-record write bottleneck at high core counts. -- Related issue: #695. From bc076be5cbbd8350448157f3a8b90458019cdcda Mon Sep 17 00:00:00 2001 From: Mark Wolters Date: Fri, 14 Aug 2026 16:36:01 -0400 Subject: [PATCH 6/7] add protection against partial writes --- .../jvector/graph/disk/NodeRecordTask.java | 74 ++++- .../graph/disk/ParallelWriteExample.java | 291 ++++++++++++++---- 2 files changed, 286 insertions(+), 79 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 ccb9b6e5e..66a37774e 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 @@ -136,9 +136,11 @@ private void callBatched() throws Exception { buildFullRecord(writer, newOrdinal); } - // One channel.write() for the entire task range — one syscall, one OS I/O request. + // 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(); - channel.write(rangeBuffer, baseOffset + (long) startOrdinal * recordSize).get(); + writeAllFully(List.of(new PendingWrite(rangeBuffer, baseOffset + (long) startOrdinal * recordSize))); } // ------------------------------------------------------------------------- @@ -155,19 +157,13 @@ private void callBatched() throws Exception { // ------------------------------------------------------------------------- private void callLegacy() throws Exception { - List> pendingWrites = new ArrayList<>(); + List pendingWrites = new ArrayList<>(); for (int newOrdinal = startOrdinal; newOrdinal < endOrdinal; newOrdinal++) { collectNodeWrites(newOrdinal, pendingWrites); } - // All writes are in-flight; wait once rather than blocking after each one. - for (var f : pendingWrites) { - int written = f.get(); - if (written < 0) { - throw new IOException("Channel closed during write"); - } - } + writeAllFully(pendingWrites); } /** @@ -185,7 +181,7 @@ private void callLegacy() throws Exception { * FUTURE IMPROVEMENT: when writeFeaturesInline is removed, null suppliers cannot exist. * This method disappears and every node is handled by buildFullRecord() in callBatched(). */ - private void collectNodeWrites(int newOrdinal, List> pending) throws Exception { + private void collectNodeWrites(int newOrdinal, List pending) throws Exception { var originalOrdinal = ordinalMapper.newToOld(newOrdinal); long recordBase = baseOffset + (long) newOrdinal * recordSize; @@ -200,7 +196,7 @@ private void collectNodeWrites(int newOrdinal, List> pending) th buf.putInt(0); // neighbor count for (int n = 0; n < graph.getDegree(0); n++) buf.putInt(-1); buf.flip(); - pending.add(channel.write(buf, recordBase)); + pending.add(new PendingWrite(buf, recordBase)); return; } @@ -239,7 +235,7 @@ private void collectNodeWrites(int newOrdinal, List> pending) th if (region.position() > 0) { int owned = region.position(); region.flip(); - pending.add(channel.write(region, regionStart)); + pending.add(new PendingWrite(region, regionStart)); regionStart += owned + feature.featureSize(); } else { // Nothing accumulated yet in this region — just skip the gap. @@ -272,7 +268,7 @@ private void collectNodeWrites(int newOrdinal, List> pending) th if (region.position() > 0) { region.flip(); - pending.add(channel.write(region, regionStart)); + pending.add(new PendingWrite(region, regionStart)); } } @@ -333,4 +329,54 @@ private ByteBuffer allocRegion(int capacity) { buf.order(java.nio.ByteOrder.BIG_ENDIAN); return buf; } + + /** 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; + } + } + + /** + * Submits every write in {@code wave} 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 wave. + *

+ * 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 writeAllFully(List wave) throws Exception { + while (!wave.isEmpty()) { + List> futures = new ArrayList<>(wave.size()); + for (var w : wave) { + futures.add(channel.write(w.buffer, w.position)); + } + + List retry = new ArrayList<>(); + for (int i = 0; i < wave.size(); i++) { + int written = futures.get(i).get(); + if (written < 0) { + throw new IOException("Channel closed during write"); + } + var w = wave.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)); + } + } + wave = retry; + } + } } 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!"); } } From 07d3ace5bcba86d5c27bb47ab1309327022163b2 Mon Sep 17 00:00:00 2001 From: Mark Wolters Date: Fri, 14 Aug 2026 18:18:14 -0400 Subject: [PATCH 7/7] speedup for the legacy path --- .../jvector/graph/disk/NodeRecordTask.java | 276 ++++++++++-------- 1 file changed, 147 insertions(+), 129 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 66a37774e..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 @@ -41,10 +41,15 @@ *

* 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. Each node's owned byte spans (ordinal, non-null features, - * neighbor section) are identified and written as a minimal set of non-blocking channel writes. - * 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. + * 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 @@ -56,8 +61,8 @@ * 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()}, {@link #collectNodeWrites}) - * and all related helpers can be deleted, and this class becomes the clean single-path fast path only. + * 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; @@ -146,130 +151,120 @@ private void callBatched() throws Exception { // ------------------------------------------------------------------------- // Legacy path: handle pre-written feature regions. // - // Pre-written bytes must not be overwritten. For each node we identify the - // contiguous owned spans (ordinal + consecutive non-null features + neighbor - // section) and fire one non-blocking channel write per span. ALL writes for - // the entire task are submitted before any Future.get() call, letting the OS - // pipeline them. + // 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 and collectNodeWrites() entirely once - // writeFeaturesInline support is removed. The fast path handles everything. + // FUTURE IMPROVEMENT: delete this method entirely once writeFeaturesInline support is + // removed. The fast path handles everything. // ------------------------------------------------------------------------- private void callLegacy() throws Exception { - List pendingWrites = new ArrayList<>(); - - for (int newOrdinal = startOrdinal; newOrdinal < endOrdinal; newOrdinal++) { - collectNodeWrites(newOrdinal, pendingWrites); - } + 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); - writeAllFully(pendingWrites); - } + 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; - /** - * Determines the owned byte spans for one node record and appends non-blocking - * channel writes for each span to {@code pending}. - *

- * The record layout is: - * 

-     *   [ordinal: 4] [feature_0: F0] [feature_1: F1] ... [neighbor_count: 4] [neighbors: D×4]
-     *
- * The ordinal and neighbor section are always owned. Features are owned when their - * supplier is non-null; null-supplier features are pre-written gaps we skip over. - * Consecutive owned bytes are merged into a single write to minimise I/O calls. - *

- * FUTURE IMPROVEMENT: when writeFeaturesInline is removed, null suppliers cannot exist. - * This method disappears and every node is handled by buildFullRecord() in callBatched(). - */ - private void collectNodeWrites(int newOrdinal, List pending) throws Exception { - var originalOrdinal = ordinalMapper.newToOld(newOrdinal); - long recordBase = baseOffset + (long) newOrdinal * recordSize; + for (int newOrdinal = startOrdinal; newOrdinal < endOrdinal; newOrdinal++) { + var originalOrdinal = ordinalMapper.newToOld(newOrdinal); + + if (originalOrdinal == OrdinalMapper.OMITTED) { + // 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) { + for (int i = 0; i < feature.featureSize(); i++) writer.writeByte(0); + } + writer.writeInt(0); // neighbor count + for (int n = 0; n < graph.getDegree(0); n++) writer.writeInt(-1); + continue; + } - if (originalOrdinal == OrdinalMapper.OMITTED) { - // OMITTED nodes are holes in the ordinal space. writeFeaturesInline() is never - // called for them, so the entire record is ours — write it in one shot. - ByteBuffer buf = allocRegion(recordSize); - buf.putInt(newOrdinal); - for (var feature : inlineFeatures) { - for (int i = 0; i < feature.featureSize(); i++) buf.put((byte) 0); + if (!graph.containsNode(originalOrdinal)) { + throw new IllegalStateException(String.format( + "Ordinal mapper mapped new ordinal %d to non-existing node %d", + newOrdinal, originalOrdinal)); } - buf.putInt(0); // neighbor count - for (int n = 0; n < graph.getDegree(0); n++) buf.putInt(-1); - buf.flip(); - pending.add(new PendingWrite(buf, recordBase)); - return; - } - if (!graph.containsNode(originalOrdinal)) { - throw new IllegalStateException(String.format( - "Ordinal mapper mapped new ordinal %d to non-existing node %d", - newOrdinal, originalOrdinal)); - } + // Ordinal: always owned. + writer.writeInt(newOrdinal); - // Accumulate contiguous owned bytes into a region buffer. When a null supplier - // (pre-written gap) is encountered, flush the current region and start a new one - // positioned past the gap. - // - // FUTURE IMPROVEMENT: this entire accumulation loop simplifies to a straight - // linear write through all features when null suppliers are impossible. - ByteBuffer region = allocRegion(recordSize); - long regionStart = recordBase; - - // Ordinal: always owned. - region.putInt(newOrdinal); - - for (var feature : inlineFeatures) { - var supplier = featureStateSuppliers.get(feature.id()); - if (supplier != null) { - // Owned: extend the current region directly. - // ByteBufferIndexWriter(region, false) picks up at region.position() - // without clearing, so it appends to whatever we have already written. - var fw = new ByteBufferIndexWriter(region, false); - feature.writeInline(fw, supplier.apply(originalOrdinal)); - } else { - // Pre-written gap: flush the current region (if any bytes are pending) - // then advance regionStart past both the flushed bytes and the gap. - // - // FUTURE IMPROVEMENT: this branch is only needed because writeFeaturesInline - // can place data at arbitrary sub-record offsets. Remove it with that API. - if (region.position() > 0) { - int owned = region.position(); - region.flip(); - pending.add(new PendingWrite(region, regionStart)); - regionStart += owned + feature.featureSize(); + 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 { - // Nothing accumulated yet in this region — just skip the gap. - regionStart += feature.featureSize(); + // 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)); + } + runFilePosition += (runEnd - runStart) + feature.featureSize(); + runStart = runEnd; } - region = allocRegion(recordSize); } - } - // Neighbor section: always owned — append to the current region. - 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))); - } - var nw = new ByteBufferIndexWriter(region, false); - nw.writeInt(neighbors.size()); - int n = 0; - for (; n < neighbors.size(); n++) { - int newNeighbor = ordinalMapper.oldToNew(neighbors.nextInt()); - if (newNeighbor < 0 || newNeighbor > ordinalMapper.maxOrdinal()) { + // 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( - "Neighbor ordinal out of bounds: %d/%d", - newNeighbor, ordinalMapper.maxOrdinal())); + "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); } - nw.writeInt(newNeighbor); + for (; n < graph.getDegree(0); n++) writer.writeInt(-1); } - for (; n < graph.getDegree(0); n++) nw.writeInt(-1); - if (region.position() > 0) { - region.flip(); - pending.add(new PendingWrite(region, regionStart)); + // Final trailing run. + int runEnd = rangeBuffer.position(); + if (runEnd > runStart) { + pending.add(new PendingWrite(sliceRange(rangeBuffer, runStart, runEnd), runFilePosition)); } + + 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(); } // ------------------------------------------------------------------------- @@ -321,15 +316,6 @@ private void buildFullRecord(ByteBufferIndexWriter writer, int newOrdinal) throw } } - /** Allocates a fresh BIG_ENDIAN ByteBuffer for a write region. */ - private ByteBuffer allocRegion(int capacity) { - var buf = useDirectBuffers - ? ByteBuffer.allocateDirect(capacity) - : ByteBuffer.allocate(capacity); - buf.order(java.nio.ByteOrder.BIG_ENDIAN); - return buf; - } - /** A not-yet-fully-written buffer destined for a fixed file offset. */ private static final class PendingWrite { final ByteBuffer buffer; @@ -342,33 +328,65 @@ private static final class PendingWrite { } /** - * Submits every write in {@code wave} without blocking, then joins them all. + * 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 wave. + * 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 writeAllFully(List wave) throws Exception { - while (!wave.isEmpty()) { - List> futures = new ArrayList<>(wave.size()); - for (var w : wave) { + 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 < wave.size(); i++) { + 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 = wave.get(i); + var w = chunk.get(i); if (w.buffer.hasRemaining()) { if (written == 0) { throw new IOException("Channel made no progress writing at position " + w.position); @@ -376,7 +394,7 @@ private void writeAllFully(List wave) throws Exception { retry.add(new PendingWrite(w.buffer, w.position + written)); } } - wave = retry; + chunk = retry; } } }