Small bug fixing in L2 quantizer.

Additional logging in benchmark.
JetBrains · Nov 6, 2023 · 1728330 · 1728330
1 parent b68b0cf
commit 1728330
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 9 deletions.
diff --git a/vectoriadb-bench/src/main/java/jetbrains/vectoriadb/bench/Sift1MBench.java b/vectoriadb-bench/src/main/java/jetbrains/vectoriadb/bench/Sift1MBench.java
@@ -100,16 +100,14 @@ public static void main(String[] args) {
 
             while (true) {
                 var indexState = client.indexState(indexName);
-                if (indexState != IndexState.BUILDING && indexState != IndexState.BUILT &&
-                        indexState != IndexState.IN_BUILD_QUEUE) {
-                    throw new IllegalStateException("Unexpected index state: " + indexState);
-                }
-
                 if (indexState == IndexState.BUILT) {
                     break;
                 }
+                Thread.sleep(2_000);
             }
 
+            stopPrintStatus.set(true);
+
             ts2 = System.currentTimeMillis();
             System.out.printf("Index %s built in %d ms%n", indexName, ts2 - ts1);
             client.switchToSearchMode();
@@ -130,9 +128,16 @@ public static void main(String[] args) {
             System.out.println("Ground truth is read, searching...");
             System.out.println("Warming up ...");
 
-            for (int i = 0; i < 10; i++) {
-                for (float[] vector : queryVectors) {
+            for (int i = 0; i < 5; i++) {
+                System.out.printf("Iteration %d out of 5 %n", (i + 1));
+
+                for (int j = 0; j < vectors.length; j++) {
+                    var vector =  vectors[j];
                     client.findNearestNeighbours(indexName, vector, 1);
+
+                    if ((j + 1) % 1_000 == 0) {
+                        System.out.printf("%d vectors are processed out of %d%n", j, vectors.length);
+                    }
                 }
             }
 

diff --git a/vectoriadb-index/src/main/java/jetbrains/vectoriadb/index/L2PQQuantizer.java b/vectoriadb-index/src/main/java/jetbrains/vectoriadb/index/L2PQQuantizer.java
@@ -269,7 +269,7 @@ private void calculateClusters(int numClusters, int iterations, long numVectors,
                 "vectors count", String.valueOf(numVectors),
                 "cores", String.valueOf(cores));
         try (var memorySession = Arena.ofShared()) {
-            var centroidIndexes = memorySession.allocate(numVectors * Integer.SIZE,
+            var centroidIndexes = memorySession.allocate(numVectors * Integer.BYTES,
                     ValueLayout.JAVA_INT.byteAlignment());
             var rng = RandomSource.XO_RO_SHI_RO_128_PP.create();
 

diff --git a/vectoriadb-java-client/src/main/java/jetbrains/vectoriadb/client/VectoriaDBClient.java b/vectoriadb-java-client/src/main/java/jetbrains/vectoriadb/client/VectoriaDBClient.java
@@ -246,7 +246,7 @@ private static void uploadVectorsArray(String indexName, float[][] vectors,
                 requestObserver.onNext(request);
                 if (progressIndicator != null) {
                     counter[0]++;
-                    progressIndicator.accept(counter[0], vector.length);
+                    progressIndicator.accept(counter[0], vectors.length);
                 }
             });