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 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/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: