From fe68c6b8be96f6a00db670a1ddb2256b7743292f Mon Sep 17 00:00:00 2001 From: jshook Date: Wed, 12 Aug 2026 17:43:18 +0000 Subject: [PATCH 1/4] graph: convert GraphIndexBuilder to the ParallelExecutor seam Route build and finalize through ParallelExecutor instead of a directly held ForkJoinPool, so an embedding host can supply its own pool or run caller-runs (no work escaping to ForkJoinPool.commonPool()). Existing ForkJoinPool constructors are preserved; they now wrap ParallelExecutor.forkJoin(pool), so current callers are unaffected. The internal parallel-for sites move from executor.submit(...) to the ParallelExecutor.forEachInt / forEach entry points. --- .../jvector/graph/GraphIndexBuilder.java | 258 ++++++++++++------ 1 file changed, 168 insertions(+), 90 deletions(-) diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndexBuilder.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndexBuilder.java index 4139a14b6..5adea127c 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndexBuilder.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndexBuilder.java @@ -74,8 +74,8 @@ public class GraphIndexBuilder implements Closeable, Accountable { private final BuildScoreProvider scoreProvider; - private final ForkJoinPool simdExecutor; - private final ForkJoinPool parallelExecutor; + private final ParallelExecutor simdExecutor; + private final ParallelExecutor parallelExecutor; private final ExplicitThreadLocal searchers; @@ -237,6 +237,39 @@ public GraphIndexBuilder(BuildScoreProvider scoreProvider, this(scoreProvider, dimension, List.of(M), beamWidth, neighborOverflow, alpha, addHierarchy, refineFinalGraph, simdExecutor, parallelExecutor); } + /** + * Reads all the vectors from vector values, builds a graph connecting them by their dense + * ordinals, using the given hyperparameter settings, and returns the resulting graph. + * + * @param scoreProvider describes how to determine the similarities between vectors + * @param M the maximum number of connections a node can have + * @param beamWidth the size of the beam search to use when finding nearest neighbors. + * @param neighborOverflow the ratio of extra neighbors to allow temporarily when inserting a + * node. larger values will build more efficiently, but use more memory. + * @param alpha how aggressive pruning diverse neighbors should be. Set alpha > 1.0 to + * allow longer edges. If alpha = 1.0 then the equivalent of the lowest level of + * an HNSW graph will be created, which is usually not what you want. + * @param addHierarchy whether we want to add an HNSW-style hierarchy on top of the Vamana index. + * @param refineFinalGraph whether we do a second pass over each node in the graph to refine its connections + * @param simdExecutor runs SIMD-heavy build iterations; use {@link ParallelExecutor#callerRuns()} + * to run them synchronously on the calling thread with no worker threads. + * @param parallelExecutor runs the parallel-stream build/cleanup iterations; use + * {@link ParallelExecutor#callerRuns()} for single-threaded, caller-runs finalize. + */ + public GraphIndexBuilder(BuildScoreProvider scoreProvider, + int dimension, + int M, + int beamWidth, + float neighborOverflow, + float alpha, + boolean addHierarchy, + boolean refineFinalGraph, + ParallelExecutor simdExecutor, + ParallelExecutor parallelExecutor) + { + this(scoreProvider, dimension, List.of(M), beamWidth, neighborOverflow, alpha, addHierarchy, refineFinalGraph, simdExecutor, parallelExecutor); + } + /** * Reads all the vectors from vector values, builds a graph connecting them by their dense * ordinals, using the given hyperparameter settings, and returns the resulting graph. @@ -295,6 +328,41 @@ public GraphIndexBuilder(BuildScoreProvider scoreProvider, boolean refineFinalGraph, ForkJoinPool simdExecutor, ForkJoinPool parallelExecutor) + { + this(scoreProvider, dimension, maxDegrees, beamWidth, neighborOverflow, alpha, addHierarchy, refineFinalGraph, + ParallelExecutor.forkJoin(simdExecutor), ParallelExecutor.forkJoin(parallelExecutor)); + } + + /** + * Reads all the vectors from vector values, builds a graph connecting them by their dense + * ordinals, using the given hyperparameter settings, and returns the resulting graph. + * + * @param scoreProvider describes how to determine the similarities between vectors + * @param maxDegrees the maximum number of connections a node can have in each layer; if fewer entries + * are specified than the number of layers, the last entry is used for all remaining layers. + * @param beamWidth the size of the beam search to use when finding nearest neighbors. + * @param neighborOverflow the ratio of extra neighbors to allow temporarily when inserting a + * node. larger values will build more efficiently, but use more memory. + * @param alpha how aggressive pruning diverse neighbors should be. Set alpha > 1.0 to + * allow longer edges. If alpha = 1.0 then the equivalent of the lowest level of + * an HNSW graph will be created, which is usually not what you want. + * @param addHierarchy whether we want to add an HNSW-style hierarchy on top of the Vamana index. + * @param refineFinalGraph whether we do a second pass over each node in the graph to refine its connections + * @param simdExecutor runs SIMD-heavy build iterations; use {@link ParallelExecutor#callerRuns()} + * to run them synchronously on the calling thread with no worker threads. + * @param parallelExecutor runs the parallel-stream build/cleanup iterations; use + * {@link ParallelExecutor#callerRuns()} for single-threaded, caller-runs finalize. + */ + public GraphIndexBuilder(BuildScoreProvider scoreProvider, + int dimension, + List maxDegrees, + int beamWidth, + float neighborOverflow, + float alpha, + boolean addHierarchy, + boolean refineFinalGraph, + ParallelExecutor simdExecutor, + ParallelExecutor parallelExecutor) { if (maxDegrees.stream().anyMatch(i -> i <= 0)) { throw new IllegalArgumentException("layer degrees must be positive"); @@ -352,6 +420,27 @@ public GraphIndexBuilder(BuildScoreProvider scoreProvider, */ @Experimental public GraphIndexBuilder(BuildScoreProvider buildScoreProvider, int dimension, MutableGraphIndex mutableGraphIndex, int beamWidth, float neighborOverflow, float alpha, boolean refineFinalGraph, ForkJoinPool simdExecutor, ForkJoinPool parallelExecutor) { + this(buildScoreProvider, dimension, mutableGraphIndex, beamWidth, neighborOverflow, alpha, refineFinalGraph, + ParallelExecutor.forkJoin(simdExecutor), ParallelExecutor.forkJoin(parallelExecutor)); + } + + /** + * Create this builder from an existing {@link io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex}, this is useful when we just loaded a graph from disk + * copy it into {@link OnHeapGraphIndex} and then start mutating it with minimal overhead of recreating the mutable {@link OnHeapGraphIndex} used in the new GraphIndexBuilder object + * + * @param buildScoreProvider the provider responsible for calculating build scores. + * @param mutableGraphIndex a mutable graph index. + * @param beamWidth the width of the beam used during the graph building process. + * @param neighborOverflow the factor determining how many additional neighbors are allowed beyond the configured limit. + * @param alpha the weight factor for balancing score computations. + * @param refineFinalGraph whether to perform a refinement step on the final graph structure. + * @param simdExecutor runs SIMD-heavy build iterations; use {@link ParallelExecutor#callerRuns()} + * to run them synchronously on the calling thread with no worker threads. + * @param parallelExecutor runs the parallel-stream build/cleanup iterations; use + * {@link ParallelExecutor#callerRuns()} for single-threaded, caller-runs finalize. + */ + @Experimental + public GraphIndexBuilder(BuildScoreProvider buildScoreProvider, int dimension, MutableGraphIndex mutableGraphIndex, int beamWidth, float neighborOverflow, float alpha, boolean refineFinalGraph, ParallelExecutor simdExecutor, ParallelExecutor parallelExecutor) { if (beamWidth <= 0) { throw new IllegalArgumentException("beamWidth must be positive"); } @@ -403,29 +492,27 @@ public static GraphIndexBuilder rescore(GraphIndexBuilder other, BuildScoreProvi var otherView = other.graph.getView(); // Copy each node and its neighbors from the old graph to the new one - other.parallelExecutor.submit(() -> { - IntStream.range(0, other.graph.getIdUpperBound()).parallel().forEach(i -> { - // Find the highest layer this node exists in - int maxLayer = other.graph.getMaxLevelForNode(i); - if (maxLayer < 0) { - return; - } + other.parallelExecutor.forEachInt(other.graph.getIdUpperBound(), i -> { + // Find the highest layer this node exists in + int maxLayer = other.graph.getMaxLevelForNode(i); + if (maxLayer < 0) { + return; + } - // Loop over 0..maxLayer, re-score neighbors for each layer - var sf = newProvider.searchProviderFor(i).scoreFunction(); - for (int lvl = 0; lvl <= maxLayer; lvl++) { - var oldNeighborsIt = otherView.getNeighborsIterator(lvl, i); - // Copy edges, compute new scores - var newNeighbors = new NodeArray(oldNeighborsIt.size()); - while (oldNeighborsIt.hasNext()) { - int neighbor = oldNeighborsIt.nextInt(); - // since we're using a different score provider, use insertSorted instead of addInOrder - newNeighbors.insertSorted(neighbor, sf.similarityTo(neighbor)); - } - newBuilder.graph.connectNode(lvl, i, newNeighbors); + // Loop over 0..maxLayer, re-score neighbors for each layer + var sf = newProvider.searchProviderFor(i).scoreFunction(); + for (int lvl = 0; lvl <= maxLayer; lvl++) { + var oldNeighborsIt = otherView.getNeighborsIterator(lvl, i); + // Copy edges, compute new scores + var newNeighbors = new NodeArray(oldNeighborsIt.size()); + while (oldNeighborsIt.hasNext()) { + int neighbor = oldNeighborsIt.nextInt(); + // since we're using a different score provider, use insertSorted instead of addInOrder + newNeighbors.insertSorted(neighbor, sf.similarityTo(neighbor)); } - }); - }).join(); + newBuilder.graph.connectNode(lvl, i, newNeighbors); + } + }); // Set the entry node newBuilder.graph.updateEntryNode(otherView.entryNode()); @@ -437,11 +524,9 @@ public ImmutableGraphIndex build(RandomAccessVectorValues ravv) { var vv = ravv.threadLocalSupplier(); int size = ravv.size(); - simdExecutor.submit(() -> { - IntStream.range(0, size).parallel().forEach(node -> { - addGraphNode(node, vv.get().getVector(node)); - }); - }).join(); + simdExecutor.forEachInt(size, node -> { + addGraphNode(node, vv.get().getVector(node)); + }); cleanup(); return graph; @@ -463,7 +548,8 @@ void validateEntryNode() { * Cleanup the graph by completing removal of marked-for-delete nodes, trimming * neighbor sets to the advertised degree, and updating the entry node. *

- * Uses default threadpool to process nodes in parallel. There is currently no way to restrict this to a single thread. + * Processes nodes in parallel via the builder's {@code parallelExecutor}. Construct the builder with + * {@link ParallelExecutor#callerRuns()} to run this cleanup synchronously on the calling thread instead. *

* Must be called before writing to disk. *

@@ -490,19 +576,15 @@ public void cleanup() { // It may be helpful for 2D use cases, but empirically it seems unnecessary for high-dimensional vectors. // It may bring a slight improvement in recall for small maximum degrees, // but it can be easily be compensated by using a slightly larger neighborOverflow. - parallelExecutor.submit(() -> { - graph.nodeStream(1).parallel().forEach(this::improveConnections); - }).join(); + parallelExecutor.forEach(graph.nodeStream(1), this::improveConnections); } // clean up overflowed neighbor lists - parallelExecutor.submit(() -> { - IntStream.range(0, graph.getIdUpperBound()).parallel().forEach(id -> { - for (int level = 0; level <= graph.getMaxLevel(); level++) { - graph.enforceDegree(id); - } - }); - }).join(); + parallelExecutor.forEachInt(graph.getIdUpperBound(), id -> { + for (int level = 0; level <= graph.getMaxLevel(); level++) { + graph.enforceDegree(id); + } + }); graph.setAllMutationsCompleted(); } @@ -701,67 +783,63 @@ public synchronized long removeDeletedNodes() { // strategy is proposed in "FreshDiskANN: A Fast and Accurate Graph-Based // ANN Index for Streaming Similarity Search" section 4.2. var newEdges = new ConcurrentHashMap>(); // new edges for key k are values v - parallelExecutor.submit(() -> { - IntStream.range(0, graph.getIdUpperBound()).parallel().forEach(i -> { - if (toDelete.get(i)) { - return; - } - for (var it = graph.getNeighborsIterator(level, i); it.hasNext(); ) { - var j = it.nextInt(); - if (toDelete.get(j)) { - var newEdgesForI = newEdges.computeIfAbsent(i, __ -> ConcurrentHashMap.newKeySet()); - for (var jt = graph.getNeighborsIterator(level, j); jt.hasNext(); ) { - int k = jt.nextInt(); - if (i != k && !toDelete.get(k)) { - newEdgesForI.add(k); - } + parallelExecutor.forEachInt(graph.getIdUpperBound(), i -> { + if (toDelete.get(i)) { + return; + } + for (var it = graph.getNeighborsIterator(level, i); it.hasNext(); ) { + var j = it.nextInt(); + if (toDelete.get(j)) { + var newEdgesForI = newEdges.computeIfAbsent(i, __ -> ConcurrentHashMap.newKeySet()); + for (var jt = graph.getNeighborsIterator(level, j); jt.hasNext(); ) { + int k = jt.nextInt(); + if (i != k && !toDelete.get(k)) { + newEdgesForI.add(k); } } } - }); - }).join(); + } + }); // Remove deleted nodes from neighbors lists; // Score the new edges, and connect the most diverse ones as neighbors - simdExecutor.submit(() -> { - newEdges.entrySet().stream().parallel().forEach(e -> { - // turn the new edges into a NodeArray - int node = e.getKey(); - // each deleted node has ALL of its neighbors added as candidates, so using approximate - // scoring and then re-scoring only the best options later makes sense here - var sf = scoreProvider.searchProviderFor(node).scoreFunction(); - var candidates = new NodeArray(graph.getDegree(level)); - for (var k : e.getValue()) { - candidates.insertSorted(k, sf.similarityTo(k)); - } + simdExecutor.forEach(newEdges.entrySet().stream(), e -> { + // turn the new edges into a NodeArray + int node = e.getKey(); + // each deleted node has ALL of its neighbors added as candidates, so using approximate + // scoring and then re-scoring only the best options later makes sense here + var sf = scoreProvider.searchProviderFor(node).scoreFunction(); + var candidates = new NodeArray(graph.getDegree(level)); + for (var k : e.getValue()) { + candidates.insertSorted(k, sf.similarityTo(k)); + } - // it's unlikely, but possible, that all the potential replacement edges were to nodes that have also - // been deleted. if that happens, keep the graph connected by adding random edges. - // (this is overly conservative -- really what we care about is that the end result of - // replaceDeletedNeighbors not be empty -- but we want to avoid having the node temporarily - // neighborless while concurrent searches run. empirically, this only results in a little extra work.) - if (candidates.size() == 0) { - var R = ThreadLocalRandom.current(); - // doing actual sampling-without-replacement is expensive so we'll loop a fixed number of times instead - for (int i = 0; i < 2 * graph.getDegree(level); i++) { - int randomNode = R.nextInt(graph.getIdUpperBound()); - while (toDelete.get(randomNode)) { - randomNode = R.nextInt(graph.getIdUpperBound()); - } - if (randomNode != node && !candidates.contains(randomNode) && graph.contains(level, randomNode)) { - float score = sf.similarityTo(randomNode); - candidates.insertSorted(randomNode, score); - } - if (candidates.size() == graph.getDegree(level)) { - break; - } + // it's unlikely, but possible, that all the potential replacement edges were to nodes that have also + // been deleted. if that happens, keep the graph connected by adding random edges. + // (this is overly conservative -- really what we care about is that the end result of + // replaceDeletedNeighbors not be empty -- but we want to avoid having the node temporarily + // neighborless while concurrent searches run. empirically, this only results in a little extra work.) + if (candidates.size() == 0) { + var R = ThreadLocalRandom.current(); + // doing actual sampling-without-replacement is expensive so we'll loop a fixed number of times instead + for (int i = 0; i < 2 * graph.getDegree(level); i++) { + int randomNode = R.nextInt(graph.getIdUpperBound()); + while (toDelete.get(randomNode)) { + randomNode = R.nextInt(graph.getIdUpperBound()); + } + if (randomNode != node && !candidates.contains(randomNode) && graph.contains(level, randomNode)) { + float score = sf.similarityTo(randomNode); + candidates.insertSorted(randomNode, score); + } + if (candidates.size() == graph.getDegree(level)) { + break; } } + } - // remove edges to deleted nodes and add the new connections, maintaining diversity - graph.replaceDeletedNeighbors(level, node, toDelete, candidates); - }); - }).join(); + // remove edges to deleted nodes and add the new connections, maintaining diversity + graph.replaceDeletedNeighbors(level, node, toDelete, candidates); + }); } // Generally we want to keep entryPoint update and node removal distinct, because both can be expensive, From bd2dfc55636d80e44223d5d883749f30a9744bb3 Mon Sep 17 00:00:00 2001 From: jshook Date: Wed, 12 Aug 2026 17:43:22 +0000 Subject: [PATCH 2/4] quantization: accept ParallelExecutor for PQ/NVQ/BQ encode Let PQ, NVQ, and BQ encode/refine on a caller-supplied ParallelExecutor so quantization participates in the host's execution budget instead of the common pool. Additive: the existing ForkJoinPool entry points are kept as delegating overloads, so no current call path changes. --- .../quantization/BinaryQuantization.java | 33 +-- .../jvector/quantization/NVQuantization.java | 27 +-- .../jvector/quantization/PQVectors.java | 40 ++-- .../quantization/ProductQuantization.java | 120 ++++++++--- .../quantization/VectorCompressor.java | 32 ++- .../QuantizationCallerRunsTest.java | 202 ++++++++++++++++++ 6 files changed, 376 insertions(+), 78 deletions(-) create mode 100644 jvector-tests/src/test/java/io/github/jbellis/jvector/quantization/QuantizationCallerRunsTest.java diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/BinaryQuantization.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/BinaryQuantization.java index f0d660301..24d955650 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/BinaryQuantization.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/BinaryQuantization.java @@ -18,6 +18,7 @@ import io.github.jbellis.jvector.disk.IndexWriter; import io.github.jbellis.jvector.disk.RandomAccessReader; +import io.github.jbellis.jvector.graph.ParallelExecutor; import io.github.jbellis.jvector.graph.RandomAccessVectorValues; import io.github.jbellis.jvector.vector.VectorizationProvider; import io.github.jbellis.jvector.vector.types.VectorFloat; @@ -26,7 +27,6 @@ import java.io.IOException; import java.util.Objects; import java.util.concurrent.ForkJoinPool; -import java.util.stream.IntStream; /** * Binary Quantization of float vectors: each float is compressed to a single bit, @@ -57,28 +57,35 @@ public static BinaryQuantization compute(RandomAccessVectorValues ravv, ForkJoin return new BinaryQuantization(ravv.dimension()); } + /** {@link ParallelExecutor} overload; the executor is unused (BQ needs no training). */ + @Deprecated + public static BinaryQuantization compute(RandomAccessVectorValues ravv, ParallelExecutor parallelExecutor) { + return new BinaryQuantization(ravv.dimension()); + } + @Override public CompressedVectors createCompressedVectors(Object[] compressedVectors) { return new ImmutableBQVectors(this, (long[][]) compressedVectors); } @Override - public CompressedVectors encodeAll(RandomAccessVectorValues ravv, ForkJoinPool simdExecutor) { + public CompressedVectors encodeAll(RandomAccessVectorValues ravv, ParallelExecutor simdExecutor) { var ravvCopy = ravv.threadLocalSupplier(); - var cv = simdExecutor.submit(() -> IntStream.range(0, ravv.size()) - .parallel() - .mapToObj(i -> { - var localRavv = ravvCopy.get(); - VectorFloat v = localRavv.getVector(i); - return v == null - ? new long[compressedVectorSize() / Long.BYTES] - : encode(v); - }) - .toArray(long[][]::new)) - .join(); + // Distinct-index array fill; the executor's completion barrier publishes the writes. + long[][] cv = new long[ravv.size()][]; + simdExecutor.forEachInt(ravv.size(), i -> { + var localRavv = ravvCopy.get(); + VectorFloat v = localRavv.getVector(i); + cv[i] = v == null ? new long[compressedVectorSize() / Long.BYTES] : encode(v); + }); return new ImmutableBQVectors(this, cv); } + @Override + public CompressedVectors encodeAll(RandomAccessVectorValues ravv, ForkJoinPool simdExecutor) { + return encodeAll(ravv, ParallelExecutor.forkJoin(simdExecutor)); + } + /** * Encodes the input vector * diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/NVQuantization.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/NVQuantization.java index aef0325b9..9adf3ec14 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/NVQuantization.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/NVQuantization.java @@ -19,6 +19,7 @@ import io.github.jbellis.jvector.annotations.VisibleForTesting; import io.github.jbellis.jvector.disk.IndexWriter; import io.github.jbellis.jvector.disk.RandomAccessReader; +import io.github.jbellis.jvector.graph.ParallelExecutor; import io.github.jbellis.jvector.graph.RandomAccessVectorValues; import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; import io.github.jbellis.jvector.util.Accountable; @@ -33,8 +34,6 @@ import java.util.List; import java.util.Objects; import java.util.concurrent.ForkJoinPool; -import java.util.stream.Collectors; -import java.util.stream.IntStream; import static io.github.jbellis.jvector.quantization.KMeansPlusPlusClusterer.UNWEIGHTED; import static io.github.jbellis.jvector.vector.VectorUtil.sub; @@ -179,18 +178,20 @@ public CompressedVectors createCompressedVectors(Object[] compressedVectors) { * Encodes the given vectors in parallel using NVQ. */ @Override - public NVQVectors encodeAll(RandomAccessVectorValues ravv, ForkJoinPool parallelExecutor) { + public NVQVectors encodeAll(RandomAccessVectorValues ravv, ParallelExecutor parallelExecutor) { var ravvCopy = ravv.threadLocalSupplier(); - return new NVQVectors(this, - parallelExecutor.submit(() -> IntStream.range(0, ravv.size()) - .parallel() - .mapToObj(i -> { - var localRavv = ravvCopy.get(); - VectorFloat v = localRavv.getVector(i); - return encode(v); - }) - .toArray(QuantizedVector[]::new)) - .join()); + // Distinct-index array fill; the executor's completion barrier publishes the writes. + QuantizedVector[] encoded = new QuantizedVector[ravv.size()]; + parallelExecutor.forEachInt(ravv.size(), i -> { + var localRavv = ravvCopy.get(); + encoded[i] = encode(localRavv.getVector(i)); + }); + return new NVQVectors(this, encoded); + } + + @Override + public NVQVectors encodeAll(RandomAccessVectorValues ravv, ForkJoinPool parallelExecutor) { + return encodeAll(ravv, ParallelExecutor.forkJoin(parallelExecutor)); } /** diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java index 760eb38ee..34d60bb11 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java @@ -18,6 +18,7 @@ import io.github.jbellis.jvector.disk.IndexWriter; import io.github.jbellis.jvector.disk.RandomAccessReader; +import io.github.jbellis.jvector.graph.ParallelExecutor; import io.github.jbellis.jvector.graph.RandomAccessVectorValues; import io.github.jbellis.jvector.graph.disk.CompactionContext; import io.github.jbellis.jvector.graph.disk.QuantizationCompactionStrategy; @@ -91,10 +92,15 @@ public static PQVectors load(RandomAccessReader in, long offset) throws IOExcept * @param simdExecutor the ForkJoinPool to use for SIMD operations * @return the PQVectors instance */ - public static ImmutablePQVectors encodeAndBuild(ProductQuantization pq, int vectorCount, RandomAccessVectorValues ravv, ForkJoinPool simdExecutor) { + public static ImmutablePQVectors encodeAndBuild(ProductQuantization pq, int vectorCount, RandomAccessVectorValues ravv, ParallelExecutor simdExecutor) { return encodeAndBuild(pq, vectorCount, IntUnaryOperator.identity(), ravv, simdExecutor); } + /** {@link ForkJoinPool} overload of {@link #encodeAndBuild(ProductQuantization, int, RandomAccessVectorValues, ParallelExecutor)}. */ + public static ImmutablePQVectors encodeAndBuild(ProductQuantization pq, int vectorCount, RandomAccessVectorValues ravv, ForkJoinPool simdExecutor) { + return encodeAndBuild(pq, vectorCount, ravv, ParallelExecutor.forkJoin(simdExecutor)); + } + /** * Build a PQVectors instance from the given RandomAccessVectorValues. The vectors are encoded in parallel * and split into chunks to avoid exceeding the maximum array size. @@ -107,6 +113,15 @@ public static ImmutablePQVectors encodeAndBuild(ProductQuantization pq, int vect * @return the PQVectors instance */ public static ImmutablePQVectors encodeAndBuild(ProductQuantization pq, int vectorCount, IntUnaryOperator ordinalsMapping, RandomAccessVectorValues ravv, ForkJoinPool simdExecutor) { + return encodeAndBuild(pq, vectorCount, ordinalsMapping, ravv, ParallelExecutor.forkJoin(simdExecutor)); + } + + /** + * As {@link #encodeAndBuild(ProductQuantization, int, IntUnaryOperator, RandomAccessVectorValues, ForkJoinPool)}, + * but {@code simdExecutor} may be {@link ParallelExecutor#callerRuns()} to encode synchronously on + * the calling thread. Encoding is per-vector independent, so the output is identical across paths. + */ + public static ImmutablePQVectors encodeAndBuild(ProductQuantization pq, int vectorCount, IntUnaryOperator ordinalsMapping, RandomAccessVectorValues ravv, ParallelExecutor simdExecutor) { // Verify that the mapped ordinals are packed, ie. the total range of ordinals does not exceed // the total vector count as they are mapped from 0 --> vector count - 1. // This assumes no duplicates are included. @@ -134,19 +149,16 @@ public static ImmutablePQVectors encodeAndBuild(ProductQuantization pq, int vect // The changes are concurrent, but because they are coordinated and do not overlap, we can use parallel streams // and then we are guaranteed safe publication because we join the thread after completion. var ravvCopy = ravv.threadLocalSupplier(); - simdExecutor.submit(() -> IntStream.range(0, vectorCount) - .parallel() - .forEach(ordinal -> { - // Retrieve the slice and mutate it. - var localRavv = ravvCopy.get(); - var slice = PQVectors.get(chunks, ordinal, layout.fullChunkVectors, pq.getSubspaceCount()); - var vector = localRavv.getVector(ordinalsMapping.applyAsInt(ordinal)); - if (vector != null) - pq.encodeTo(vector, slice); - else - slice.zero(); - })) - .join(); + simdExecutor.forEachInt(vectorCount, ordinal -> { + // Retrieve the slice and mutate it. + var localRavv = ravvCopy.get(); + var slice = PQVectors.get(chunks, ordinal, layout.fullChunkVectors, pq.getSubspaceCount()); + var vector = localRavv.getVector(ordinalsMapping.applyAsInt(ordinal)); + if (vector != null) + pq.encodeTo(vector, slice); + else + slice.zero(); + }); return new ImmutablePQVectors(pq, chunks, vectorCount, layout.fullChunkVectors); } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java index 37823db28..47bb98e96 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java @@ -19,6 +19,7 @@ import io.github.jbellis.jvector.annotations.VisibleForTesting; import io.github.jbellis.jvector.disk.IndexWriter; import io.github.jbellis.jvector.disk.RandomAccessReader; +import io.github.jbellis.jvector.graph.ParallelExecutor; import io.github.jbellis.jvector.graph.RandomAccessVectorValues; import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; import io.github.jbellis.jvector.util.Accountable; @@ -36,12 +37,9 @@ import java.util.List; import java.util.Objects; import java.util.SplittableRandom; -import java.util.concurrent.Callable; import java.util.concurrent.ForkJoinPool; import java.util.concurrent.atomic.AtomicReference; import java.util.logging.Logger; -import java.util.stream.Collectors; -import java.util.stream.IntStream; import static io.github.jbellis.jvector.quantization.KMeansPlusPlusClusterer.UNWEIGHTED; import static io.github.jbellis.jvector.util.MathUtil.square; @@ -113,6 +111,30 @@ public static ProductQuantization compute(RandomAccessVectorValues ravv, float anisotropicThreshold, ForkJoinPool simdExecutor, ForkJoinPool parallelExecutor) + { + return compute(ravv, M, clusterCount, globallyCenter, anisotropicThreshold, + ParallelExecutor.forkJoin(simdExecutor), ParallelExecutor.forkJoin(parallelExecutor)); + } + + /** + * As {@link #compute(RandomAccessVectorValues, int, int, boolean, float, ForkJoinPool, ForkJoinPool)}, + * but the executors may be {@link ParallelExecutor#callerRuns()} to train synchronously on the + * calling thread with no worker threads. + *

+ * Codebook training draws random k-means++ seeds from {@code ThreadLocalRandom}, so the trained + * codebooks are already non-deterministic run-to-run and are not byte-identical between the + * caller-runs and pool-backed paths; both produce validly-trained, recall-equivalent codebooks. + * + * @param simdExecutor executor for SIMD/codebook operations + * @param parallelExecutor executor for training-vector extraction/centering + */ + public static ProductQuantization compute(RandomAccessVectorValues ravv, + int M, + int clusterCount, + boolean globallyCenter, + float anisotropicThreshold, + ParallelExecutor simdExecutor, + ParallelExecutor parallelExecutor) { checkClusterCount(clusterCount); @@ -126,9 +148,12 @@ public static ProductQuantization compute(RandomAccessVectorValues ravv, VectorFloat globalCentroid; if (globallyCenter) { globalCentroid = KMeansPlusPlusClusterer.centroidOf(vectors); - // subtract the centroid from each vector + // subtract the centroid from each vector (distinct-index array fill; the executor's + // completion barrier safely publishes the writes back to this thread) List> finalVectors = vectors; - vectors = simdExecutor.submit(() -> finalVectors.stream().parallel().map(v -> VectorUtil.sub(v, globalCentroid)).collect(Collectors.>toList())).join(); + VectorFloat[] centered = new VectorFloat[finalVectors.size()]; + simdExecutor.forEachInt(finalVectors.size(), i -> centered[i] = VectorUtil.sub(finalVectors.get(i), globalCentroid)); + vectors = Arrays.asList(centered); } else { globalCentroid = null; } @@ -138,11 +163,14 @@ public static ProductQuantization compute(RandomAccessVectorValues ravv, return new ProductQuantization(codebooks, clusterCount, subvectorSizesAndOffsets, globalCentroid, anisotropicThreshold); } - static List> extractTrainingVectors(RandomAccessVectorValues ravv, ForkJoinPool parallelExecutor) { - final IntStream ordinalStream; + static List> extractTrainingVectors(RandomAccessVectorValues ravv, ParallelExecutor parallelExecutor) { + final int[] ords; if (ravv.size() <= MAX_PQ_TRAINING_SET_SIZE) { - ordinalStream = IntStream.range(0, ravv.size()); + ords = new int[ravv.size()]; + for (int i = 0; i < ords.length; i++) { + ords[i] = i; + } } else { // Uses Floyd’s sampling algorithm to select MAX_PQ_TRAINING_SET_SIZE random ordinals from 0 to ravv.size() // while only iterating MAX_PQ_TRAINING_SET_SIZE times. @@ -157,25 +185,25 @@ static List> extractTrainingVectors(RandomAccessVectorValues ravv ordinals.add(t); } } - int[] ordinalArray = new int[ordinals.size()]; + ords = new int[ordinals.size()]; IntHashSet.IntIterator it = ordinals.iterator(); - for (int i = 0; i < ordinals.size(); i++) { + for (int i = 0; i < ords.length; i++) { assert it.hasNext(); - ordinalArray[i] = it.next(); + ords[i] = it.next(); } assert !it.hasNext(); - ordinalStream = IntStream.of(ordinalArray); } + // Distinct-index array fill in the target order; the executor's completion barrier safely + // publishes the writes (identical result under forkJoin and callerRuns). var ravvCopy = ravv.threadLocalSupplier(); - return parallelExecutor.submit(() -> ordinalStream.parallel() - .mapToObj(targetOrd -> { - var localRavv = ravvCopy.get(); - VectorFloat v = localRavv.getVector(targetOrd); - return localRavv.isValueShared() ? v.copy() : v; - }) - .collect(Collectors.toList())) - .join(); + VectorFloat[] out = new VectorFloat[ords.length]; + parallelExecutor.forEachInt(ords.length, i -> { + var localRavv = ravvCopy.get(); + VectorFloat v = localRavv.getVector(ords[i]); + out[i] = localRavv.isValueShared() ? v.copy() : v; + }); + return Arrays.asList(out); } /** @@ -196,6 +224,22 @@ public ProductQuantization refine(RandomAccessVectorValues ravv, float anisotropicThreshold, ForkJoinPool simdExecutor, ForkJoinPool parallelExecutor) + { + return refine(ravv, lloydsRounds, anisotropicThreshold, + ParallelExecutor.forkJoin(simdExecutor), ParallelExecutor.forkJoin(parallelExecutor)); + } + + /** + * As {@link #refine(RandomAccessVectorValues, int, float, ForkJoinPool, ForkJoinPool)}, but the + * executors may be {@link ParallelExecutor#callerRuns()} to refine synchronously on the calling + * thread. See {@link #compute(RandomAccessVectorValues, int, int, boolean, float, ParallelExecutor, ParallelExecutor)} + * for the determinism note (codebooks are recall-equivalent, not byte-identical, across paths). + */ + public ProductQuantization refine(RandomAccessVectorValues ravv, + int lloydsRounds, + float anisotropicThreshold, + ParallelExecutor simdExecutor, + ParallelExecutor parallelExecutor) { if (lloydsRounds < 0) { throw new IllegalArgumentException("lloydsRounds must be non-negative"); @@ -204,18 +248,21 @@ public ProductQuantization refine(RandomAccessVectorValues ravv, var subvectorSizesAndOffsets = getSubvectorSizesAndOffsets(ravv.dimension(), M); var vectorsMutable = extractTrainingVectors(ravv, parallelExecutor); if (globalCentroid != null) { - var vectors = vectorsMutable; - vectorsMutable = simdExecutor.submit(() -> vectors.stream().parallel().map(v -> VectorUtil.sub(v, globalCentroid)).collect(Collectors.>toList())).join(); + List> src = vectorsMutable; + VectorFloat[] centered = new VectorFloat[src.size()]; + simdExecutor.forEachInt(src.size(), i -> centered[i] = VectorUtil.sub(src.get(i), globalCentroid)); + vectorsMutable = Arrays.asList(centered); } var vectors = vectorsMutable; // "effectively final" to make the closure happy - Callable[]> callable = () -> IntStream.range(0, M).parallel().mapToObj(m -> { + // Per-subquantizer (independent) codebook fill; completion barrier publishes the writes. + VectorFloat[] refinedCodebooks = new VectorFloat[M]; + simdExecutor.forEachInt(M, m -> { VectorFloat[] subvectors = extractSubvectors(vectors, m, subvectorSizesAndOffsets); var clusterer = new KMeansPlusPlusClusterer(subvectors, codebooks[m], anisotropicThreshold); - return clusterer.cluster(anisotropicThreshold == UNWEIGHTED ? lloydsRounds : 0, - anisotropicThreshold == UNWEIGHTED ? 0 : lloydsRounds); - }).toArray(VectorFloat[]::new); - var refinedCodebooks = simdExecutor.submit(callable).join(); + refinedCodebooks[m] = clusterer.cluster(anisotropicThreshold == UNWEIGHTED ? lloydsRounds : 0, + anisotropicThreshold == UNWEIGHTED ? 0 : lloydsRounds); + }); return new ProductQuantization(refinedCodebooks, clusterCount, subvectorSizesAndOffsets, globalCentroid, anisotropicThreshold); } @@ -258,10 +305,15 @@ public ImmutablePQVectors createCompressedVectors(Object[] compressedVectors) { * as a zero vector. */ @Override - public PQVectors encodeAll(RandomAccessVectorValues ravv, ForkJoinPool simdExecutor) { + public PQVectors encodeAll(RandomAccessVectorValues ravv, ParallelExecutor simdExecutor) { return PQVectors.encodeAndBuild(this, ravv.size(), ravv, simdExecutor); } + @Override + public PQVectors encodeAll(RandomAccessVectorValues ravv, ForkJoinPool simdExecutor) { + return encodeAll(ravv, ParallelExecutor.forkJoin(simdExecutor)); + } + /** * Encodes the input vector using the PQ codebooks, weighing parallel loss more than orthogonal loss, into * the given ByteSequence. @@ -484,14 +536,16 @@ public int getClusterCount() { return clusterCount; } - static VectorFloat[] createCodebooks(List> vectors, int[][] subvectorSizeAndOffset, int clusters, float anisotropicThreshold, ForkJoinPool simdExecutor) { + static VectorFloat[] createCodebooks(List> vectors, int[][] subvectorSizeAndOffset, int clusters, float anisotropicThreshold, ParallelExecutor simdExecutor) { int M = subvectorSizeAndOffset.length; - Callable[]> callable = () -> IntStream.range(0, M).parallel().mapToObj(m -> { + // Per-subquantizer (independent) codebook fill; completion barrier publishes the writes. + VectorFloat[] codebooks = new VectorFloat[M]; + simdExecutor.forEachInt(M, m -> { VectorFloat[] subvectors = extractSubvectors(vectors, m, subvectorSizeAndOffset); var clusterer = new KMeansPlusPlusClusterer(subvectors, clusters, anisotropicThreshold); - return clusterer.cluster(K_MEANS_ITERATIONS, anisotropicThreshold == UNWEIGHTED ? 0 : K_MEANS_ITERATIONS); - }).toArray(VectorFloat[]::new); - return simdExecutor.submit(callable).join(); + codebooks[m] = clusterer.cluster(K_MEANS_ITERATIONS, anisotropicThreshold == UNWEIGHTED ? 0 : K_MEANS_ITERATIONS); + }); + return codebooks; } /** diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/VectorCompressor.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/VectorCompressor.java index c5708716c..f253acb16 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/VectorCompressor.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/VectorCompressor.java @@ -18,6 +18,7 @@ import io.github.jbellis.jvector.disk.IndexWriter; import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues; +import io.github.jbellis.jvector.graph.ParallelExecutor; import io.github.jbellis.jvector.graph.RandomAccessVectorValues; import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; import io.github.jbellis.jvector.util.PhysicalCoreExecutor; @@ -44,11 +45,24 @@ default CompressedVectors encodeAll(RandomAccessVectorValues ravv) { /** * Encode all vectors in the RandomAccessVectorValues. If the RandomAccessVectorValues * has a missing vector for a given ordinal, the value will be encoded as a zero vector. + *

+ * Encoding is per-vector independent, so with {@link ParallelExecutor#callerRuns()} it runs + * synchronously on the calling thread and produces output identical to the pool-backed path; + * only wall-clock and thread usage differ. * @param ravv RandomAccessVectorValues to encode - * @param simdExecutor ForkJoinPool to use for SIMD operations + * @param simdExecutor executor hosting the parallel encode; {@link ParallelExecutor#callerRuns()} + * runs it on the calling thread * @return CompressedVectors containing the encoded vectors */ - CompressedVectors encodeAll(RandomAccessVectorValues ravv, ForkJoinPool simdExecutor); + CompressedVectors encodeAll(RandomAccessVectorValues ravv, ParallelExecutor simdExecutor); + + /** + * As {@link #encodeAll(RandomAccessVectorValues, ParallelExecutor)}, hosting the parallel encode + * on {@code simdExecutor}. Retained for callers holding a {@link ForkJoinPool}. + */ + default CompressedVectors encodeAll(RandomAccessVectorValues ravv, ForkJoinPool simdExecutor) { + return encodeAll(ravv, ParallelExecutor.forkJoin(simdExecutor)); + } T encode(VectorFloat v); @@ -108,8 +122,16 @@ default double[] reconstructionErrors(RandomAccessVectorValues ravv) { * @return the reconstruction error for each vector */ default double[] reconstructionErrors(RandomAccessVectorValues ravv, ForkJoinPool simdExecutor) { - return simdExecutor.submit(() -> - IntStream.range(0, ravv.size()).mapToDouble(i -> reconstructionError(ravv.getVector(i))).toArray() - ).join(); + return reconstructionErrors(ravv, ParallelExecutor.forkJoin(simdExecutor)); + } + + /** + * As {@link #reconstructionErrors(RandomAccessVectorValues, ForkJoinPool)}, hosting the parallel + * computation on the given {@link ParallelExecutor}. + */ + default double[] reconstructionErrors(RandomAccessVectorValues ravv, ParallelExecutor simdExecutor) { + double[] out = new double[ravv.size()]; + simdExecutor.forEachInt(ravv.size(), i -> out[i] = reconstructionError(ravv.getVector(i))); + return out; } } diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/quantization/QuantizationCallerRunsTest.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/quantization/QuantizationCallerRunsTest.java new file mode 100644 index 000000000..f109e6edd --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/quantization/QuantizationCallerRunsTest.java @@ -0,0 +1,202 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.quantization; + +import com.carrotsearch.randomizedtesting.RandomizedTest; +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; +import io.github.jbellis.jvector.TestUtil; +import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues; +import io.github.jbellis.jvector.graph.ParallelExecutor; +import io.github.jbellis.jvector.graph.RandomAccessVectorValues; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import org.junit.Test; + +import java.util.List; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ForkJoinPool; + +import static io.github.jbellis.jvector.quantization.KMeansPlusPlusClusterer.UNWEIGHTED; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +/// Tests that the quantization path accepts {@link ParallelExecutor} and works caller-runs. +/// +/// The contract is deliberately split (see {@link ProductQuantization#compute} javadoc): +/// - **Encoding** is per-vector independent and deterministic given a fixed codebook, so +/// {@code encodeAll} produces byte-identical output under {@code callerRuns()} and a +/// {@code ForkJoinPool}. +/// - **Training** ({@code compute}/{@code refine}) draws k-means++ seeds from +/// {@code ThreadLocalRandom} and is therefore already non-deterministic run-to-run; it is +/// not byte-identical across executors, only recall/quality-equivalent. +@ThreadLeakScope(ThreadLeakScope.Scope.NONE) +public class QuantizationCallerRunsTest extends RandomizedTest { + + private static final int DIM = 64; + private static final int SIZE = 2_000; + private static final int M = 8; + private static final int CLUSTERS = 256; + + private List> vectors; + private ListRandomAccessVectorValues ravv; + + private void makeData() { + vectors = TestUtil.createRandomVectors(SIZE, DIM); + ravv = new ListRandomAccessVectorValues(vectors, DIM); + } + + /// Encoding a fixed PQ codebook must be byte-identical whether run on a pool or caller-runs. + @Test + public void pqEncodeIsByteIdenticalAcrossExecutors() { + makeData(); + ForkJoinPool pool = new ForkJoinPool(4); + try { + ProductQuantization pq = ProductQuantization.compute(ravv, M, CLUSTERS, true, UNWEIGHTED, + ParallelExecutor.forkJoin(pool), ParallelExecutor.forkJoin(pool)); + + PQVectors onPool = pq.encodeAll(ravv, ParallelExecutor.forkJoin(pool)); + PQVectors callerRuns = pq.encodeAll(ravv, ParallelExecutor.callerRuns()); + + assertEquals("PQ encoding must be byte-identical across executors", onPool, callerRuns); + } finally { + pool.shutdown(); + } + } + + /// Binary quantization encoding must also be byte-identical across executors. + @Test + public void bqEncodeIsByteIdenticalAcrossExecutors() { + makeData(); + ForkJoinPool pool = new ForkJoinPool(4); + try { + BinaryQuantization bq = new BinaryQuantization(DIM); + CompressedVectors onPool = bq.encodeAll(ravv, ParallelExecutor.forkJoin(pool)); + CompressedVectors callerRuns = bq.encodeAll(ravv, ParallelExecutor.callerRuns()); + assertEquals("BQ encoding must be byte-identical across executors", onPool, callerRuns); + } finally { + pool.shutdown(); + } + } + + /// Training on a pool vs caller-runs is NOT byte-identical (ThreadLocalRandom seeds), but must + /// produce codebooks of equivalent quality — asserted via average reconstruction error. + @Test + public void pqTrainingIsQualityEquivalentAcrossExecutors() { + makeData(); + ForkJoinPool pool = new ForkJoinPool(4); + try { + ProductQuantization onPool = ProductQuantization.compute(ravv, M, CLUSTERS, true, UNWEIGHTED, + ParallelExecutor.forkJoin(pool), ParallelExecutor.forkJoin(pool)); + ProductQuantization callerRuns = ProductQuantization.compute(ravv, M, CLUSTERS, true, UNWEIGHTED, + ParallelExecutor.callerRuns(), ParallelExecutor.callerRuns()); + + double msePool = meanReconstructionError(onPool, ParallelExecutor.forkJoin(pool)); + double mseCaller = meanReconstructionError(callerRuns, ParallelExecutor.callerRuns()); + + assertTrue("pool-trained codebook should reconstruct well: " + msePool, msePool < 0.05); + assertTrue("caller-runs-trained codebook should reconstruct well: " + mseCaller, mseCaller < 0.05); + assertTrue("training quality must be equivalent across executors (" + msePool + " vs " + mseCaller + ")", + Math.abs(msePool - mseCaller) < 0.01); + } finally { + pool.shutdown(); + } + } + + /// refine() must also run caller-runs and produce a quality-equivalent codebook. + @Test + public void pqRefineRunsCallerRuns() { + makeData(); + ForkJoinPool pool = new ForkJoinPool(4); + try { + ProductQuantization base = ProductQuantization.compute(ravv, M, CLUSTERS, true, UNWEIGHTED, + ParallelExecutor.forkJoin(pool), ParallelExecutor.forkJoin(pool)); + ProductQuantization refined = base.refine(ravv, 1, UNWEIGHTED, + ParallelExecutor.callerRuns(), ParallelExecutor.callerRuns()); + assertTrue("refined codebook should reconstruct well", + meanReconstructionError(refined, ParallelExecutor.callerRuns()) < 0.05); + } finally { + pool.shutdown(); + } + } + + /// With callerRuns(), train + encode must execute only on the calling thread — no worker threads + /// and the common pool untouched. + @Test + public void callerRunsStaysOnCallingThread() { + makeData(); + Set observed = ConcurrentHashMap.newKeySet(); + RecordingRAVV recording = new RecordingRAVV(ravv, observed); + String mainName = Thread.currentThread().getName(); + + ProductQuantization pq = ProductQuantization.compute(recording, M, CLUSTERS, true, UNWEIGHTED, + ParallelExecutor.callerRuns(), ParallelExecutor.callerRuns()); + pq.encodeAll(recording, ParallelExecutor.callerRuns()); + + assertFalse("some vector reads must have occurred", observed.isEmpty()); + for (String name : observed) { + assertTrue("caller-runs must read only on the calling thread, but saw: " + name + " (all=" + observed + ")", + name.equals(mainName)); + } + } + + private double meanReconstructionError(ProductQuantization pq, ParallelExecutor ex) { + double[] errs = pq.reconstructionErrors(ravv, ex); + double s = 0; + for (double e : errs) { + s += e; + } + return s / errs.length; + } + + /// Records the thread of every getVector call, funnelling thread-local copies through one set. + private static final class RecordingRAVV implements RandomAccessVectorValues { + private final RandomAccessVectorValues delegate; + private final Set observed; + + RecordingRAVV(RandomAccessVectorValues delegate, Set observed) { + this.delegate = delegate; + this.observed = observed; + } + + @Override + public int size() { + return delegate.size(); + } + + @Override + public int dimension() { + return delegate.dimension(); + } + + @Override + public VectorFloat getVector(int nodeId) { + observed.add(Thread.currentThread().getName()); + return delegate.getVector(nodeId); + } + + @Override + public boolean isValueShared() { + return delegate.isValueShared(); + } + + @Override + public RandomAccessVectorValues copy() { + return new RecordingRAVV(delegate.copy(), observed); + } + } +} From 93f1f2a8462630a57153d9b2427c303fb5a8b2b2 Mon Sep 17 00:00:00 2001 From: jshook Date: Wed, 12 Aug 2026 18:25:07 +0000 Subject: [PATCH 3/4] graph: EmbeddedExecutionContext as the single execution carrier Carry the execution resources an embedder passes to jvector in one place -- a compute ParallelExecutor plus merge/io executors -- so the pool is supplied once and "no work escapes to the common pool" holds in a single spot. of(pool) bounds everything to a host pool; callerRuns() runs inline on the calling thread (the memtable-flush case). Wired here to the seams that exist on this branch: graph build (newBuilder) and PQ/NVQ train/refine/encode route through the compute executor. The compaction and parallel-writer wiring depends on the compactor's calling-convention conversion, which is illustrated in prose in doc/compaction-seam-conversion.md; the merge/io roles are carried so the note can reference them directly. --- .../graph/EmbeddedExecutionContext.java | 200 ++++++++++++++++++ 1 file changed, 200 insertions(+) create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/EmbeddedExecutionContext.java diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/EmbeddedExecutionContext.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/EmbeddedExecutionContext.java new file mode 100644 index 000000000..d8683af28 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/EmbeddedExecutionContext.java @@ -0,0 +1,200 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph; + +import io.github.jbellis.jvector.annotations.Experimental; +import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider; +import io.github.jbellis.jvector.quantization.NVQVectors; +import io.github.jbellis.jvector.quantization.NVQuantization; +import io.github.jbellis.jvector.quantization.PQVectors; +import io.github.jbellis.jvector.quantization.ProductQuantization; + +import java.util.Collections; +import java.util.List; +import java.util.Objects; +import java.util.concurrent.AbstractExecutorService; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.ForkJoinPool; +import java.util.concurrent.TimeUnit; + +import static io.github.jbellis.jvector.quantization.KMeansPlusPlusClusterer.UNWEIGHTED; + +/** + * A single carrier for the execution resources an embedder supplies to jvector, so the pool is + * passed once and the "no parallel work escapes to {@link io.github.jbellis.jvector.util.PhysicalCoreExecutor#pool()} + * or {@link ForkJoinPool#commonPool()}" guarantee lives in one place. Construct the context + * with a bounded pool (or {@link #callerRuns()}), then obtain a builder or run quantization through + * it; every operation routes only through the carried executors and never through a default + * overload. + * + *

Two modes. {@link #of(ForkJoinPool)} bounds all work to a supplied pool (the compaction + * use case). {@link #callerRuns()} runs everything synchronously on the calling thread — no worker + * threads, no pool — which lets a memtable flush run graph build and PQ/NVQ train/encode entirely on + * its own flush-writer thread. + * + *

Three executor roles. {@code compute} (a {@link ParallelExecutor}) drives graph build / + * cleanup and all PQ/NVQ quantization; {@code merge} (an {@link Executor}) drives the compaction + * merge; {@code io} (an {@link ExecutorService}) drives the optional IO-bound parallel graph writer. + * The factories derive all three consistently from one input, so "pass one thing" stays the simple, + * correct default; the general constructor accepts them separately for embedders that size an IO pool + * distinctly. + * + *

Scope of this branch. The compaction and parallel-writer wiring — {@code newCompactor()}, + * {@code newParallelWriter()}, and the PQ-retrain entry points — depends on the compactor's calling + * conventions being converted onto the seams (executor, {@link io.github.jbellis.jvector.util.work.ProgressLimiter}, + * {@link io.github.jbellis.jvector.graph.disk.CompactionDestination}). That conversion is illustrated + * in prose in {@code doc/compaction-seam-conversion.md} rather than carried as code here; the + * {@code merge} and {@code io} roles are still carried below so the note can reference + * {@link #mergeExecutor()} / {@link #ioExecutor()} directly. + * + *

Lifecycle. The context neither owns nor shuts down the embedder's pool(s); the embedder + * supplies and disposes them. + */ +@Experimental +public final class EmbeddedExecutionContext { + private final ParallelExecutor compute; + private final Executor merge; + private final ExecutorService io; + + /** + * @param compute drives graph build/cleanup and all PQ/NVQ work + * @param merge drives the compaction merge (a {@code ForkJoinPool} or {@code Runnable::run} for + * caller-runs) + * @param io drives the IO-bound parallel graph writer + */ + public EmbeddedExecutionContext(ParallelExecutor compute, Executor merge, ExecutorService io) { + this.compute = Objects.requireNonNull(compute, "compute"); + this.merge = Objects.requireNonNull(merge, "merge"); + this.io = Objects.requireNonNull(io, "io"); + } + + /** One pool for everything (compute, merge, and IO). */ + public static EmbeddedExecutionContext of(ForkJoinPool pool) { + return new EmbeddedExecutionContext(ParallelExecutor.forkJoin(pool), pool, pool); + } + + /** A compute pool plus a distinct IO pool; the compute pool also drives the merge. */ + public static EmbeddedExecutionContext of(ForkJoinPool computePool, ExecutorService io) { + return new EmbeddedExecutionContext(ParallelExecutor.forkJoin(computePool), computePool, io); + } + + /** + * Runs every operation synchronously on the calling thread — no worker threads, no pool, common + * pool untouched. Suitable for a memtable flush that wants graph build + PQ/NVQ encode on its own + * thread. Encoding is byte-identical to the pool-backed path; PQ training is quality-equivalent + * (its k-means seeds are drawn from {@code ThreadLocalRandom} either way). + */ + public static EmbeddedExecutionContext callerRuns() { + return new EmbeddedExecutionContext(ParallelExecutor.callerRuns(), Runnable::run, directExecutorService()); + } + + /** The compute executor (graph build + PQ/NVQ). */ + public ParallelExecutor parallelExecutor() { + return compute; + } + + /** The merge executor (compaction). See {@code doc/compaction-seam-conversion.md}. */ + public Executor mergeExecutor() { + return merge; + } + + /** The IO executor (parallel graph writer). See {@code doc/compaction-seam-conversion.md}. */ + public ExecutorService ioExecutor() { + return io; + } + + // ---- graph construction ---- + + /** + * A {@link GraphIndexBuilder} wired to the compute executor for both its executor roles. + */ + public GraphIndexBuilder newBuilder(BuildScoreProvider scoreProvider, + int dimension, + int M, + int beamWidth, + float neighborOverflow, + float alpha, + boolean addHierarchy, + boolean refineFinalGraph) { + return new GraphIndexBuilder(scoreProvider, dimension, M, beamWidth, neighborOverflow, alpha, + addHierarchy, refineFinalGraph, compute, compute); + } + + // ---- product quantization ---- + + /** Trains PQ on the compute executor (isotropic / unweighted). */ + public ProductQuantization trainPQ(RandomAccessVectorValues ravv, int M, int clusterCount, boolean globallyCenter) { + return trainPQ(ravv, M, clusterCount, globallyCenter, UNWEIGHTED); + } + + /** Trains PQ on the compute executor with an explicit anisotropic threshold. */ + public ProductQuantization trainPQ(RandomAccessVectorValues ravv, int M, int clusterCount, boolean globallyCenter, float anisotropicThreshold) { + return ProductQuantization.compute(ravv, M, clusterCount, globallyCenter, anisotropicThreshold, compute, compute); + } + + /** Refines an existing PQ codebook on the compute executor (one Lloyd's round, unweighted). */ + public ProductQuantization refinePQ(ProductQuantization base, RandomAccessVectorValues ravv) { + return base.refine(ravv, 1, UNWEIGHTED, compute, compute); + } + + /** Encodes all vectors with PQ on the compute executor. */ + public PQVectors encodePQ(ProductQuantization pq, RandomAccessVectorValues ravv) { + return pq.encodeAll(ravv, compute); + } + + // ---- non-uniform vector quantization ---- + + /** Encodes all vectors with NVQ on the compute executor. */ + public NVQVectors encodeNVQ(NVQuantization nvq, RandomAccessVectorValues ravv) { + return nvq.encodeAll(ravv, compute); + } + + /** An {@link ExecutorService} that runs every submitted task synchronously on the calling thread. */ + private static ExecutorService directExecutorService() { + return new AbstractExecutorService() { + @Override + public void execute(Runnable command) { + command.run(); + } + + @Override + public void shutdown() { + } + + @Override + public List shutdownNow() { + return Collections.emptyList(); + } + + @Override + public boolean isShutdown() { + return false; + } + + @Override + public boolean isTerminated() { + return false; + } + + @Override + public boolean awaitTermination(long timeout, TimeUnit unit) { + return true; + } + }; + } +} From d6a3f68e3b8ee6996e8fe2ec8814d8303e128e52 Mon Sep 17 00:00:00 2001 From: jshook Date: Wed, 12 Aug 2026 18:46:20 +0000 Subject: [PATCH 4/4] docs: illustrate the compactor's calling-convention conversion Add docs/compaction-seam-conversion.md: how OnDiskGraphIndexCompactor's calling conventions map onto the integration-robustness seams -- execution (Executor / EmbeddedExecutionContext), progress + throttle (ProgressLimiter), output (CompactionDestination / SeekableSink), PQ retrain and the parallel writer via the carried executors, plus the compactor-local memory-safety guards (drain-on-unwind, truncate-reused-outputs) that land with it. In the current lineage the compactor's seam-wiring is inseparable from the compaction-algorithm work, so it is described here rather than carried as code; this note is the target the clean compaction work should hit. --- docs/compaction-seam-conversion.md | 190 +++++++++++++++++++++++++++++ 1 file changed, 190 insertions(+) create mode 100644 docs/compaction-seam-conversion.md diff --git a/docs/compaction-seam-conversion.md b/docs/compaction-seam-conversion.md new file mode 100644 index 000000000..ff97a26e0 --- /dev/null +++ b/docs/compaction-seam-conversion.md @@ -0,0 +1,190 @@ +# Compaction seam conversion (illustration) + +`integration-robustness` introduced the embedding seams as standalone types, with no +consumer: + +| Seam | Types | Role | +| --- | --- | --- | +| Execution | `graph.ParallelExecutor`, `graph.EmbeddedExecutionContext` | host supplies the pool (or caller-runs); no work escapes to `ForkJoinPool.commonPool()` | +| Progress + throttle | `util.work.ProgressLimiter` (`ProgressTracker` + `WorkLimiter`), `WorkStage` | progress reported up, write bandwidth admitted down, cancellation checkpoints | +| Output | `graph.disk.CompactionDestination`, `disk.SeekableSink` | compaction output written through a caller-owned channel | +| Runtime mode | `util.RuntimeMode` | diagnostic-only work gated off by default | +| Reader safety | `OnDiskGraphIndex` bounded reads (already landed) | a stale/out-of-range offset fails diagnosably instead of faulting | + +This branch converts the **general** build/quantize calling conventions onto the execution +seam as code: + +- `graph: convert GraphIndexBuilder to the ParallelExecutor seam` +- `quantization: accept ParallelExecutor for PQ/NVQ/BQ encode` +- `graph: EmbeddedExecutionContext as the single execution carrier` + +The **compactor** conversion is described here rather than carried as code: in the current +lineage the compactor's seam-wiring is inseparable from the ten compaction-algorithm +commits (a single ~1,100-line diff), so converting it cleanly is the job of the compaction +work that will be rebuilt on top of `integration-robustness`. This note is the target that +work should hit — each section is `before` (today's calling convention on `main`) and +`after` (the seam-based convention), using the real signatures from the integration lineage. + +--- + +## 1. Execution — a host pool instead of the common pool + +**Before.** The compactor reaches for a process-wide pool internally, so compaction work +competes on `ForkJoinPool.commonPool()` / `PhysicalCoreExecutor.pool()` with everything else +in the host. + +**After.** The primary constructor accepts any `Executor` and only falls back to the shared +pool when none is supplied. `taskWindowSize` bounds in-flight batches (the compactor drives +its merge through an `ExecutorCompletionService` over this executor). + +```java +public OnDiskGraphIndexCompactor( + List sources, + List liveNodes, + List remappers, + VectorSimilarityFunction similarityFunction, + Executor executor, // host pool; null => PhysicalCoreExecutor.pool() + int taskWindowSize) // bound on concurrent in-flight batches +``` + +`EmbeddedExecutionContext` is where a host passes its pool once; `newCompactor(...)` wires +the compactor to the carried `mergeExecutor()`: + +```java +public OnDiskGraphIndexCompactor newCompactor(List sources, + List liveNodes, + List remappers, + VectorSimilarityFunction similarityFunction, + int taskWindowSize) { + return new OnDiskGraphIndexCompactor(sources, liveNodes, remappers, + similarityFunction, mergeExecutor(), taskWindowSize); +} +``` + +`callerRuns()` collapses all three executor roles onto the calling thread, so a memtable +flush can run the whole merge on its own flush-writer thread with no worker pool at all. + +--- + +## 2. Progress + throttle — `ProgressLimiter` + +**Before.** `compact()` runs opaque: no progress until it returns, and its (often dominant) +internal write bandwidth is outside any host throughput budget. + +**After.** A single opt-in limiter, defaulting to a no-op, is installed per operation: + +```java +private volatile ProgressLimiter limiter = ProgressLimiter.UNLIMITED; + +/** null restores ProgressLimiter.UNLIMITED. Returns this for chaining. */ +public OnDiskGraphIndexCompactor setProgressLimiter(ProgressLimiter limiter) { ... } +``` + +Inside `compact()`, at each phase boundary the compactor calls the limiter both ways, and +each call doubles as the cancellation checkpoint: + +- **up:** `limiter.onProgress(stage, completed, total)` — advances `nodetool compactionstats` + / `system_views.sstable_tasks` *while* the merge runs, instead of jumping 0% → 100%. +- **down:** `limiter.acquire(bytes)` before a write batch — admits the merge's write + bandwidth against the host's shared compaction rate limiter. +- **cancel:** both entry points throw if the host requested stop; jvector drains its in-flight + workers before `compact()` unwinds, so no source read survives the cancellation. + +Phases are scoped with `startPhase(WorkStage)` so each stage (extract, retrain, link, write, +footer) is separately observable. + +The limiter is deliberately **not** carried on `EmbeddedExecutionContext`: throttle/progress +is per host operation, so it is set on the returned compactor per call. + +--- + +## 3. Output virtualization — `CompactionDestination` + +**Before.** The compactor allocates and opens its own output file. + +**After.** `compact(CompactionDestination)` writes the graph body through a caller-owned +target, so the host can hand it a slot inside a larger container (e.g. an SAI component after +a reserved header) with no temp file and one write of the body: + +```java +public long compact(CompactionDestination destination) throws IOException { + try (CompactionDestination.Target target = destination.open()) { + // ...write the compacted graph into target.file() at target.startOffset()... + target.commit(bodyLength); // success: body durable; embedder finalizes its own footer + return bodyLength; + } // close() always runs; no commit() => aborted, partial output discarded +} +``` + +The embedder computes its own footer/checksum over the body via +`SeekableSink.over(channel, target.startOffset())`. The path-based `compact(Path)` / +`compact(Path, long)` entry points remain for callers that still want jvector to own the file. + +--- + +## 4. PQ retrain through the executor + +**Before.** Retrain during compaction submits to the common pool — the historical "retrain +leak" where work escaped the host's budget. + +**After.** Retrain takes the compute executor, closing the leak: + +```java +// PQRetrainer +public ProductQuantization retrain(VectorSimilarityFunction sim, + ParallelExecutor simd, ParallelExecutor parallel); +public ProductQuantization retrain(VectorSimilarityFunction sim, ProductQuantization basePQ, + ParallelExecutor simd, ParallelExecutor parallel); +``` + +`EmbeddedExecutionContext.retrainPQ(retrainer, sim[, basePQ])` forwards the carried compute +executor. (The general PQ train/refine/encode calls already route through the executor on +this branch — see `EmbeddedExecutionContext.trainPQ/refinePQ/encodePQ`.) + +--- + +## 5. Parallel (IO-bound) graph writer + +**After.** The parallel writer runs on the host's IO executor rather than a default pool: + +```java +new OnDiskParallelGraphIndexWriter.Builder(graph, outputPath).withExecutor(io); +``` + +`EmbeddedExecutionContext.newParallelWriter(graph, path)` supplies `ioExecutor()`. + +--- + +## 6. Compactor-side memory safety (lands with this conversion) + +The reader-side guard — bounded record reads in `OnDiskGraphIndex`, plus the mmap-lifecycle +`close()` contract on `ReaderSupplier` / `SimpleMappedReader` — is already on +`integration-robustness`. The remaining, compactor-local guards belong with this conversion +because they live inside `compact()`: + +- **drain on unwind** — on any exception or host cancel, in-flight merge workers are drained + before `compact()` returns/throws, so no worker read outlives the call (no use-after-unmap + when the host reclaims source mappings). +- **truncate reused outputs** — an output being reused is truncated before the new body is + written, so stale tail bytes past the new body can never be mistaken for graph data. + +--- + +## What `EmbeddedExecutionContext` regains after the conversion + +The trimmed context on this branch carries `mergeExecutor()` / `ioExecutor()` but omits the +compactor/writer/retrain factory methods, because they depend on the signatures above. Once +the compactor is converted, they return as thin wiring: + +```java +newCompactor(...) -> new OnDiskGraphIndexCompactor(..., mergeExecutor(), window) +newParallelWriter(...) -> new ...Writer.Builder(graph, path).withExecutor(ioExecutor()) +retrainPQ(...) -> retrainer.retrain(sim, parallelExecutor(), parallelExecutor()) +``` + +## What stays for the clean compaction work + +Everything above is the **calling-convention** surface. The body of `compact()` — layer +enumeration, cross-source linking, retain-largest, pre-encode caching, and the rest of the +ten algorithm commits — is the compaction work proper, to be rebuilt cleanly on top of +`integration-robustness` against exactly these seams.