From 77cf2b77f9a13bfbc84cf19ff7a9b63daca2c6a7 Mon Sep 17 00:00:00 2001 From: Navneet Verma Date: Sat, 7 Sep 2024 19:16:05 -0700 Subject: [PATCH] [Bug-Fix] Fixing the Quantization to clone the quantized vector, prepare the quantized vector correctly. Signed-off-by: Navneet Verma --- .../index/codec/nativeindex/QuantizationIndexUtils.java | 3 ++- .../quantizationOutput/BinaryQuantizationOutput.java | 7 +++++++ .../models/quantizationOutput/QuantizationOutput.java | 2 ++ .../knn/quantization/quantizer/OneBitScalarQuantizer.java | 3 ++- 4 files changed, 13 insertions(+), 2 deletions(-) diff --git a/src/main/java/org/opensearch/knn/index/codec/nativeindex/QuantizationIndexUtils.java b/src/main/java/org/opensearch/knn/index/codec/nativeindex/QuantizationIndexUtils.java index bebe9e8b0b..fa951a4cca 100644 --- a/src/main/java/org/opensearch/knn/index/codec/nativeindex/QuantizationIndexUtils.java +++ b/src/main/java/org/opensearch/knn/index/codec/nativeindex/QuantizationIndexUtils.java @@ -33,7 +33,8 @@ static Object processAndReturnVector(KNNVectorValues knnVectorValues, IndexBu knnVectorValues.getVector(), indexBuildSetup.getQuantizationOutput() ); - return indexBuildSetup.getQuantizationOutput().getQuantizedVector(); + // We should get the cloned Vector because Quantization output always have the same reference + return indexBuildSetup.getQuantizationOutput().cloneQuantizedVector(); } else { return knnVectorValues.conditionalCloneVector(); } diff --git a/src/main/java/org/opensearch/knn/quantization/models/quantizationOutput/BinaryQuantizationOutput.java b/src/main/java/org/opensearch/knn/quantization/models/quantizationOutput/BinaryQuantizationOutput.java index 388fd9e94f..72f65eb4ce 100644 --- a/src/main/java/org/opensearch/knn/quantization/models/quantizationOutput/BinaryQuantizationOutput.java +++ b/src/main/java/org/opensearch/knn/quantization/models/quantizationOutput/BinaryQuantizationOutput.java @@ -54,6 +54,13 @@ public boolean isPrepared(int vectorLength) { return vectorLength == currentVectorLength && quantizedVector != null; } + @Override + public byte[] cloneQuantizedVector() { + byte[] clonedByteArray = new byte[quantizedVector.length]; + System.arraycopy(quantizedVector, 0, clonedByteArray, 0, quantizedVector.length); + return clonedByteArray; + } + /** * Returns the quantized vector. * diff --git a/src/main/java/org/opensearch/knn/quantization/models/quantizationOutput/QuantizationOutput.java b/src/main/java/org/opensearch/knn/quantization/models/quantizationOutput/QuantizationOutput.java index 4d088f91fd..31f6465438 100644 --- a/src/main/java/org/opensearch/knn/quantization/models/quantizationOutput/QuantizationOutput.java +++ b/src/main/java/org/opensearch/knn/quantization/models/quantizationOutput/QuantizationOutput.java @@ -33,4 +33,6 @@ public interface QuantizationOutput { * @return true if the quantized vector is already prepared, false otherwise. */ boolean isPrepared(int vectorLength); + + T cloneQuantizedVector(); } diff --git a/src/main/java/org/opensearch/knn/quantization/quantizer/OneBitScalarQuantizer.java b/src/main/java/org/opensearch/knn/quantization/quantizer/OneBitScalarQuantizer.java index ac48a95235..fee0349e93 100644 --- a/src/main/java/org/opensearch/knn/quantization/quantizer/OneBitScalarQuantizer.java +++ b/src/main/java/org/opensearch/knn/quantization/quantizer/OneBitScalarQuantizer.java @@ -84,7 +84,8 @@ public void quantize(final float[] vector, final QuantizationState state, final if (thresholds == null || thresholds.length != vectorLength) { throw new IllegalArgumentException("Thresholds must not be null and must match the dimension of the vector."); } - if (!output.isPrepared(vectorLength)) output.prepareQuantizedVector(vectorLength); + // This will ensure that we are preparing the Quantized Vector always. aka setting all the bits to 0. + output.prepareQuantizedVector(vectorLength); BitPacker.quantizeAndPackBits(vector, thresholds, output.getQuantizedVector()); }