From 7d78115f0532b3a2fb75c745476ab7334ef9816c Mon Sep 17 00:00:00 2001 From: Navneet Verma Date: Tue, 9 Jul 2024 10:36:02 -0700 Subject: [PATCH] Fixing the arithmetic to find the number of vectors to stream from java to jni layer Signed-off-by: Navneet Verma --- .../opensearch/knn/index/codec/util/KNNCodecUtil.java | 10 ++++++---- 1 file changed, 6 insertions(+), 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 index e059626081..c684424dc8 100644 --- 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) {