moving updateAndGet last is not correct

jbellis · Jan 9, 2025 · 0017908 · 0017908
1 parent 9143fd3
commit 0017908
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/MutablePQVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/MutablePQVectors.java
@@ -48,15 +48,17 @@ public MutablePQVectors(ProductQuantization pq) {
     @Override
     public void encodeAndSet(int ordinal, VectorFloat<?> vector) {
         ensureChunkCapacity(ordinal);
-        pq.encodeTo(vector, get(ordinal));
+        // increase count first so get() works
         vectorCount.updateAndGet(current -> max(current, ordinal + 1));
+        pq.encodeTo(vector, get(ordinal));
     }
 
     @Override
     public void setZero(int ordinal) {
         ensureChunkCapacity(ordinal);
-        get(ordinal).zero();
+        // increase count first so get() works
         vectorCount.updateAndGet(current -> max(current, ordinal + 1));
+        get(ordinal).zero();
     }
 
     private void ensureChunkCapacity(int ordinal) {