Fixing the arithmetic to find the number of vectors to stream from ja…

…va to jni layer Signed-off-by: Navneet Verma <navneev@amazon.com>
navneet1v · Jul 9, 2024 · 7d78115 · 7d78115
1 parent 5139b16
commit 7d78115
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/src/main/java/org/opensearch/knn/index/codec/util/KNNCodecUtil.java b/src/main/java/org/opensearch/knn/index/codec/util/KNNCodecUtil.java
@@ -63,11 +63,13 @@ public static KNNCodecUtil.Pair getFloats(BinaryDocValues values) throws IOExcep
                 dimension = vector.length;
 
                 if (vectorsPerTransfer == Integer.MIN_VALUE) {
-                    vectorsPerTransfer = (dimension * Float.BYTES * totalLiveDocs) / vectorsStreamingMemoryLimit;
-                    // This condition comes if vectorsStreamingMemoryLimit is higher than total number floats to transfer
-                    // Doing this will reduce 1 extra trip to JNI layer.
+                    // if vectorsStreamingMemoryLimit is 100 bytes and we have floats with 5 dimension, then per transfer
+                    // we have to send 100/(5 * 4) => 20 vectors.
+                    vectorsPerTransfer = vectorsStreamingMemoryLimit / ((long) dimension * Float.BYTES);
+                    // If vectorsPerTransfer comes out to be 0, then we set number of vectors per transfer to 1, to ensure that
+                    // we are sending minimum number of vectors.
                     if (vectorsPerTransfer == 0) {
-                        vectorsPerTransfer = totalLiveDocs;
+                        vectorsPerTransfer = 1;
                     }
                 }
                 if (vectorList.size() == vectorsPerTransfer) {