Fix test, add coverage for byte element types

carlosdelest · Dec 5, 2024 · 7d083b7 · 7d083b7
1 parent 9c9773f
commit 7d083b7
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/server/src/main/java/org/elasticsearch/search/vectors/ESKnnByteVectorQuery.java b/server/src/main/java/org/elasticsearch/search/vectors/ESKnnByteVectorQuery.java
@@ -35,4 +35,8 @@ protected TopDocs mergeLeafResults(TopDocs[] perLeafResults) {
     public void profile(QueryProfiler queryProfiler) {
         queryProfiler.addVectorOpsCount(vectorOpsCount);
     }
+
+    public Integer kParam() {
+        return kParam;
+    }
 }
diff --git a/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldTypeTests.java b/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldTypeTests.java
@@ -23,6 +23,7 @@
 import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.VectorSimilarity;
 import org.elasticsearch.search.DocValueFormat;
 import org.elasticsearch.search.vectors.DenseVectorQuery;
+import org.elasticsearch.search.vectors.ESKnnByteVectorQuery;
 import org.elasticsearch.search.vectors.ESKnnFloatVectorQuery;
 import org.elasticsearch.search.vectors.RescoreKnnVectorQuery;
 import org.elasticsearch.search.vectors.VectorData;
@@ -409,10 +410,11 @@ public void testByteCreateKnnQuery() {
     }
 
     public void testRescoreOversampleUsedWithoutQuantization() {
+        DenseVectorFieldMapper.ElementType elementType = randomFrom(FLOAT, BYTE);
         DenseVectorFieldType nonQuantizedField = new DenseVectorFieldType(
             "f",
             IndexVersion.current(),
-            randomFrom(FLOAT, BYTE),
+            elementType,
             3,
             true,
             VectorSimilarity.COSINE,
@@ -430,9 +432,15 @@ public void testRescoreOversampleUsedWithoutQuantization() {
             null
         );
 
-        ESKnnFloatVectorQuery esKnnQuery = (ESKnnFloatVectorQuery) knnQuery;
-        assertThat(esKnnQuery.getK(), is(100));
-        assertThat(esKnnQuery.kParam(), is(10));
+        if (elementType == BYTE) {
+            ESKnnByteVectorQuery esKnnQuery = (ESKnnByteVectorQuery) knnQuery;
+            assertThat(esKnnQuery.getK(), is(100));
+            assertThat(esKnnQuery.kParam(), is(10));
+        } else {
+            ESKnnFloatVectorQuery esKnnQuery = (ESKnnFloatVectorQuery) knnQuery;
+            assertThat(esKnnQuery.getK(), is(100));
+            assertThat(esKnnQuery.kParam(), is(10));
+        }
     }
 
     public void testRescoreOversampleModifiesKnnParams() {