Parallel Write performance improvements

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<ImmutableGraphIndex.View> viewPerThread;
-    private final ThreadLocal<ByteBuffer> bufferPerThread;
     private final CopyOnWriteArrayList<ImmutableGraphIndex.View> 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<Void> 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();
                 });
 

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<BenchResult> 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);
 

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,

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<Integer> mGrid,
                        List<Integer> efConstructionGrid,
                        List<Float> 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<Boolean> 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<Set<FeatureId>, ImmutableGraphIndex> buildOnDisk(List<? exten
                                                                         Path outputDir,
                                                                         VectorCompressor<?> buildCompressor,
                                                                         Map<Set<FeatureId>, OnDiskGraphIndexCache.WriteHandle> handles,
-                                                                        ConstructionMetrics constructionMetrics) throws IOException
+                                                                        ConstructionMetrics constructionMetrics,
+                                                                        boolean useParallelConstruction) throws IOException
     {
         Files.createDirectories(outputDir);
 
@@ -388,9 +392,9 @@ private static Map<Set<FeatureId>, ImmutableGraphIndex> buildOnDisk(List<? exten
         GraphIndexBuilder builder = new GraphIndexBuilder(bsp, floatVectors.dimension(), M, efConstruction, neighborOverflow, 1.2f, addHierarchy, refineFinalGraph);
 
         // use the inline vectors index as the score provider for graph construction
-        Map<Set<FeatureId>, OnDiskGraphIndexWriter> writers = new HashMap<>();
+        Map<Set<FeatureId>, RandomAccessOnDiskGraphIndexWriter> writers = new HashMap<>();
         Map<Set<FeatureId>, Map<FeatureId, IntFunction<Feature.State>>> 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<Set<FeatureId>, ImmutableGraphIndex> buildOnDisk(List<? exten
             } else {
                 graphPath = outputDir.resolve("graph" + n++);
             }
-            var bws = builderWithSuppliers(features, builder.getGraph(), graphPath, floatVectors, pq.getCompressor(), constructionMetrics);
-            var writer = bws.builder.build();
+            var bws = builderWithSuppliers(features, builder.getGraph(), graphPath, floatVectors, pq.getCompressor(), constructionMetrics, useParallelConstruction);
+            var writer = bws.writer;
             writers.put(features, writer);
             suppliers.put(features, bws.suppliers);
             if (features.contains(FeatureId.INLINE_VECTORS) || features.contains(FeatureId.NVQ_VECTORS)) {
@@ -483,14 +487,35 @@ private static BuilderWithSuppliers builderWithSuppliers(Set<FeatureId> 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<FeatureId, IntFunction<Feature.State>> 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<FeatureId> features,
+                                             ImmutableGraphIndex onHeapGraph,
+                                             RandomAccessVectorValues floatVectors,
+                                             ProductQuantization pq,
+                                             ConstructionMetrics constructionMetrics,
+                                             Map<FeatureId, IntFunction<Feature.State>> suppliers) {
         for (var featureId : features) {
             switch (featureId) {
                 case INLINE_VECTORS:
@@ -513,10 +538,8 @@ private static BuilderWithSuppliers builderWithSuppliers(Set<FeatureId> 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<FeatureId, IntFunction<Feature.State>> suppliers;
 
-        public BuilderWithSuppliers(OnDiskGraphIndexWriter.Builder builder, Map<FeatureId, IntFunction<Feature.State>> suppliers) {
-            this.builder = builder;
+        public BuilderWithSuppliers(RandomAccessOnDiskGraphIndexWriter writer, Map<FeatureId, IntFunction<Feature.State>> suppliers) {
+            this.writer = writer;
             this.suppliers = suppliers;
         }
     }
@@ -594,8 +617,8 @@ private static Map<Set<FeatureId>, ImmutableGraphIndex> buildInMemory(List<? ext
                 continue;
             }
             var graphPath = testDirectory.resolve("graph" + n++);
-            var bws = builderWithSuppliers(features, onHeapGraph, graphPath, floatVectors, null, null);
-            try (var writer = bws.builder.build()) {
+            var bws = builderWithSuppliers(features, onHeapGraph, graphPath, floatVectors, null, null, false);
+            try (var writer = bws.writer) {
                 start = System.nanoTime();
                 writer.write(bws.suppliers);
                 System.out.format("Wrote %s in %.2fs%n", features, (System.nanoTime() - start) / 1_000_000_000.0);
@@ -812,6 +835,7 @@ private static Map<String, Object> ordered(Object... kv) {
     public static List<BenchResult> runAllAndCollectResults(
             DataSet ds,
             boolean enableIndexCache,
+            boolean useParallelConstruction,
             List<Integer> mGrid,
             List<Integer> efConstructionGrid,
             List<Float> neighborOverflowGrid,
@@ -904,7 +928,7 @@ public static List<BenchResult> 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);
                                             }
 

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,

jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/ConstructionParameters.java

@@ -32,6 +32,11 @@ public class ConstructionParameters extends CommonParameters {
     public List<String> reranking;
     public List<Boolean> fusedGraph;
     public Boolean useSavedIndexIfExists;
+    public Boolean parallelGraphConstruction;
+
+    public boolean isParallelGraphConstruction() {
+        return Boolean.TRUE.equals(parallelGraphConstruction);
+    }
 
     public List<EnumSet<FeatureId>> getFeatureSets() {
         List<EnumSet<FeatureId>> featureSets = null;

jvector-examples/yaml-configs/index-parameters/cap-10M.yml → 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]

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:

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]

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: