Re-Call Issue Fix with Binary Quantized Vectors

src/main/java/org/opensearch/knn/index/codec/nativeindex/QuantizationIndexUtils.java

@@ -33,7 +33,7 @@ static Object processAndReturnVector(KNNVectorValues<?> knnVectorValues, IndexBu
                 knnVectorValues.getVector(),
                 indexBuildSetup.getQuantizationOutput()
             );
-            return indexBuildSetup.getQuantizationOutput().getQuantizedVector();
+            return indexBuildSetup.getQuantizationOutput().getQuantizedVectorCopy();
         } else {
             return knnVectorValues.conditionalCloneVector();
         }

src/main/java/org/opensearch/knn/quantization/models/quantizationOutput/BinaryQuantizationOutput.java

@@ -63,4 +63,16 @@ public boolean isPrepared(int vectorLength) {
     public byte[] getQuantizedVector() {
         return quantizedVector;
     }
+
+    /**
+     * Returns a copy of the quantized vector.
+     *
+     * @return a copy of the quantized vector byte array.
+     */
+    @Override
+    public byte[] getQuantizedVectorCopy() {
+        byte[] clonedByteArray = new byte[quantizedVector.length];
+        System.arraycopy(quantizedVector, 0, clonedByteArray, 0, quantizedVector.length);
+        return clonedByteArray;
+    }
 }

src/main/java/org/opensearch/knn/quantization/models/quantizationOutput/QuantizationOutput.java

@@ -33,4 +33,11 @@ public interface QuantizationOutput<T> {
      * @return true if the quantized vector is already prepared, false otherwise.
      */
     boolean isPrepared(int vectorLength);
+
+    /**
+     * Returns a copy of the quantized vector.
+     *
+     * @return a copy of the quantized data.
+     */
+    T getQuantizedVectorCopy();
 }

src/main/java/org/opensearch/knn/quantization/quantizer/MultiBitScalarQuantizer.java

@@ -139,7 +139,7 @@ public void quantize(final float[] vector, final QuantizationState state, final
         if (thresholds == null || thresholds[0].length != vector.length) {
             throw new IllegalArgumentException("Thresholds must not be null and must match the dimension of the vector.");
         }
-        if (!output.isPrepared(vectorLength)) output.prepareQuantizedVector(vectorLength);
+        output.prepareQuantizedVector(vectorLength);
         BitPacker.quantizeAndPackBits(vector, thresholds, bitsPerCoordinate, output.getQuantizedVector());
     }
 

src/main/java/org/opensearch/knn/quantization/quantizer/OneBitScalarQuantizer.java

@@ -84,7 +84,7 @@ 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);
+        output.prepareQuantizedVector(vectorLength);
         BitPacker.quantizeAndPackBits(vector, thresholds, output.getQuantizedVector());
     }
 

src/test/java/org/opensearch/knn/index/codec/nativeindex/DefaultIndexBuildStrategyTests.java

@@ -141,16 +141,16 @@ public void testBuildAndWrite_withQuantization() {
             ArgumentCaptor<float[]> vectorCaptor = ArgumentCaptor.forClass(float[].class);
             // New: Create QuantizationOutput and mock the quantization process
             QuantizationOutput<byte[]> quantizationOutput = mock(QuantizationOutput.class);
-            when(quantizationOutput.getQuantizedVector()).thenReturn(new byte[] { 1, 2 });
+            when(quantizationOutput.getQuantizedVectorCopy()).thenReturn(new byte[] { 1, 2 });
             when(quantizationService.createQuantizationOutput(eq(quantizationState.getQuantizationParams()))).thenReturn(
                 quantizationOutput
             );
 
             // Quantize the vector with the quantization output
             when(quantizationService.quantize(eq(quantizationState), vectorCaptor.capture(), eq(quantizationOutput))).thenAnswer(
                 invocation -> {
-                    quantizationOutput.getQuantizedVector();
-                    return quantizationOutput.getQuantizedVector();
+                    quantizationOutput.getQuantizedVectorCopy();
+                    return quantizationOutput.getQuantizedVectorCopy();
                 }
             );
             when(quantizationState.getDimensions()).thenReturn(2);

src/test/java/org/opensearch/knn/index/codec/nativeindex/MemOptimizedNativeIndexBuildStrategyTests.java

@@ -168,16 +168,16 @@ public void testBuildAndWrite_withQuantization() {
             ArgumentCaptor<float[]> vectorCaptor = ArgumentCaptor.forClass(float[].class);
             // New: Create QuantizationOutput and mock the quantization process
             QuantizationOutput<byte[]> quantizationOutput = mock(QuantizationOutput.class);
-            when(quantizationOutput.getQuantizedVector()).thenReturn(new byte[] { 1, 2 });
+            when(quantizationOutput.getQuantizedVectorCopy()).thenReturn(new byte[] { 1, 2 });
             when(quantizationService.createQuantizationOutput(eq(quantizationState.getQuantizationParams()))).thenReturn(
                 quantizationOutput
             );
 
             // Quantize the vector with the quantization output
             when(quantizationService.quantize(eq(quantizationState), vectorCaptor.capture(), eq(quantizationOutput))).thenAnswer(
                 invocation -> {
-                    quantizationOutput.getQuantizedVector();
-                    return quantizationOutput.getQuantizedVector();
+                    quantizationOutput.getQuantizedVectorCopy();
+                    return quantizationOutput.getQuantizedVectorCopy();
                 }
             );
             when(quantizationState.getDimensions()).thenReturn(2);