Make ravv usage thread-safe (#381)

It solves the issue that @jkni caught when reviewing #374. It renders the parallel vector accesses during "encodeAll" thread safe.
jbellis · Jan 3, 2025 · 5b0df95 · 5b0df95
1 parent e8d5c3c
commit 5b0df95
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
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
@@ -64,14 +64,16 @@ public CompressedVectors createCompressedVectors(Object[] compressedVectors) {
 
     @Override
     public CompressedVectors encodeAll(RandomAccessVectorValues ravv, ForkJoinPool simdExecutor) {
+        var ravvCopy = ravv.threadLocalSupplier();
         var cv = simdExecutor.submit(() -> IntStream.range(0, ravv.size())
                 .parallel()
                 .mapToObj(i -> {
-                    var vector = ravv.getVector(i);
-                    return vector == null
+                    var localRavv = ravvCopy.get();
+                    VectorFloat<?> v = localRavv.getVector(i);
+                    return v == null
                             ? new long[compressedVectorSize() / Long.BYTES]
-                            : encode(vector);
-                })
+                            : encode(v);
+                        })
                 .toArray(long[][]::new))
                 .join();
         return new ImmutableBQVectors(this, cv);

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
@@ -137,12 +137,14 @@ public static ImmutablePQVectors encodeAndBuild(ProductQuantization pq, int vect
         // Encode the vectors in parallel into the compressed data chunks
         // 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, ravv.size())
                         .parallel()
                         .forEach(ordinal -> {
                             // Retrieve the slice and mutate it.
+                            var localRavv = ravvCopy.get();
                             var slice = PQVectors.get(chunks, ordinal, vectorsPerChunk, pq.getSubspaceCount());
-                            var vector = ravv.getVector(ordinal);
+                            var vector = localRavv.getVector(ordinal);
                             if (vector != null)
                                 pq.encodeTo(vector, slice);
                             else