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

…va to jni layer Signed-off-by: Navneet Verma <[email protected]>
navneet1v · Jul 9, 2024 · 84c9b4c · 84c9b4c
1 parent 5139b16
commit 84c9b4c
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -18,6 +18,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 * Adds dynamic query parameter ef_search in radial search faiss engine [#1790](https://github.com/opensearch-project/k-NN/pull/1790)
 ### Enhancements
 ### Bug Fixes
+* Fixing the arithmetic to find the number of vectors to stream from java to jni layer.[#1804](https://github.com/opensearch-project/k-NN/pull/1804)
 ### Infrastructure
 ### Documentation
 * Update dev guide to fix clang linking issue on arm [#1746](https://github.com/opensearch-project/k-NN/pull/1746)

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) {