Test Improvements

Signed-off-by: Vikasht34 <[email protected]>
Vikasht34 · Feb 12, 2025 · 7ee5c9b · 7ee5c9b
1 parent 2209680
commit 7ee5c9b
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 3 deletions.
diff --git a/src/main/java/org/opensearch/knn/quantization/models/requests/TrainingRequest.java b/src/main/java/org/opensearch/knn/quantization/models/requests/TrainingRequest.java
@@ -31,7 +31,5 @@ public abstract class TrainingRequest<T> {
      */
     public abstract T getVectorAtThePosition(int position) throws IOException;
 
-    public void resetVectorValues() {
-        // Default no-op implementation
-    }
+    public abstract void resetVectorValues();
 }
diff --git a/src/main/java/org/opensearch/knn/training/FloatTrainingDataConsumer.java b/src/main/java/org/opensearch/knn/training/FloatTrainingDataConsumer.java
@@ -98,6 +98,12 @@ private List<byte[]> quantizeVectors(List<?> vectors) throws IOException {
             public float[] getVectorAtThePosition(int position) {
                 return ArrayUtils.toPrimitive((Float[]) vectors.get(position));
             }
+
+            @Override
+            public void resetVectorValues() {
+                // No-op
+            }
+
         };
         QuantizationState quantizationState = quantizer.train(trainingRequest);
         BinaryQuantizationOutput binaryQuantizationOutput = new BinaryQuantizationOutput(quantizationConfig.getQuantizationType().getId());

diff --git a/src/test/java/org/opensearch/knn/quantization/quantizer/MultiBitScalarQuantizerTests.java b/src/test/java/org/opensearch/knn/quantization/quantizer/MultiBitScalarQuantizerTests.java
@@ -223,5 +223,10 @@ public MockTrainingRequest(ScalarQuantizationParams params, float[][] vectors) {
         public float[] getVectorAtThePosition(int position) {
             return vectors[position];
         }
+
+        @Override
+        public void resetVectorValues() {
+            // No-op
+        }
     }
 }
diff --git a/src/test/java/org/opensearch/knn/quantization/quantizer/OneBitScalarQuantizerTests.java b/src/test/java/org/opensearch/knn/quantization/quantizer/OneBitScalarQuantizerTests.java
@@ -31,6 +31,11 @@ public void testTrain_withTrainingRequired() throws IOException {
             public float[] getVectorAtThePosition(int position) {
                 return vectors[position];
             }
+
+            @Override
+            public void resetVectorValues() {
+                // No-op
+            }
         };
         OneBitScalarQuantizer quantizer = new OneBitScalarQuantizer();
         QuantizationState state = quantizer.train(originalRequest);
@@ -47,6 +52,11 @@ public void testTrain_withBelowAboveThresholdMeans() throws IOException {
             public float[] getVectorAtThePosition(int position) {
                 return vectors[position];
             }
+
+            @Override
+            public void resetVectorValues() {
+                // No-op
+            }
         };
         OneBitScalarQuantizer quantizer = new OneBitScalarQuantizer();
         QuantizationState state = quantizer.train(trainingRequest);
@@ -150,6 +160,11 @@ public void testTrain_withRotationApplied() throws IOException {
             public float[] getVectorAtThePosition(int position) {
                 return vectors[position];
             }
+
+            @Override
+            public void resetVectorValues() {
+                // No-op
+            }
         };
 
         OneBitScalarQuantizer quantizer = new OneBitScalarQuantizer();
@@ -168,6 +183,11 @@ public void testTrain_withoutRotationMatrix() throws IOException {
             public float[] getVectorAtThePosition(int position) {
                 return vectors[position];
             }
+
+            @Override
+            public void resetVectorValues() {
+                // No-op
+            }
         };
 
         OneBitScalarQuantizer quantizer = new OneBitScalarQuantizer();
@@ -263,6 +283,11 @@ public void testCalculateMean() throws IOException {
             public float[] getVectorAtThePosition(int position) {
                 return vectors[position];
             }
+
+            @Override
+            public void resetVectorValues() {
+                // No-op
+            }
         };
 
         Sampler sampler = SamplingFactory.getSampler(SamplerType.RESERVOIR);
@@ -284,6 +309,11 @@ public void testCalculateMean_withNullVector() {
             public float[] getVectorAtThePosition(int position) {
                 return vectors[position];
             }
+
+            @Override
+            public void resetVectorValues() {
+                // No-op
+            }
         };
 
         Sampler sampler = SamplingFactory.getSampler(SamplerType.RESERVOIR);