improving tests with null checks and added diskbbq format to densevectorfieldtype test

Author: john-wagster

server/src/internalClusterTest/java/org/elasticsearch/search/query/RescoreKnnVectorQueryIT.java

@@ -116,7 +116,7 @@ private record TestParams(
         float[] queryVector,
         int k,
         int numCands,
-        float visitPercentage,
+        Float visitPercentage,
         RescoreVectorBuilder rescoreVectorBuilder
     ) {
         public static TestParams generate() {
@@ -129,7 +129,7 @@ public static TestParams generate() {
                 randomVector(numDims),
                 k,
                 (int) (k * randomFloatBetween(1.0f, 10.0f, true)),
-                randomFloatBetween(0.0f, 100.0f, true),
+                randomBoolean() ? null : randomFloatBetween(0.0f, 100.0f, true),
                 new RescoreVectorBuilder(randomFloatBetween(1.0f, 100f, true))
             );
         }

server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldTypeTests.java

@@ -25,22 +25,29 @@
 import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.VectorSimilarity;
 import org.elasticsearch.search.DocValueFormat;
 import org.elasticsearch.search.vectors.DenseVectorQuery;
+import org.elasticsearch.search.vectors.DiversifyingChildrenIVFKnnFloatVectorQuery;
 import org.elasticsearch.search.vectors.DiversifyingParentBlockQuery;
 import org.elasticsearch.search.vectors.ESKnnByteVectorQuery;
 import org.elasticsearch.search.vectors.ESKnnFloatVectorQuery;
+import org.elasticsearch.search.vectors.IVFKnnFloatVectorQuery;
 import org.elasticsearch.search.vectors.RescoreKnnVectorQuery;
 import org.elasticsearch.search.vectors.VectorData;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Set;
 import java.util.function.Function;
 
+import static org.elasticsearch.index.codec.vectors.IVFVectorsFormat.MAX_VECTORS_PER_CLUSTER;
+import static org.elasticsearch.index.codec.vectors.IVFVectorsFormat.MIN_VECTORS_PER_CLUSTER;
 import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.BBQ_MIN_DIMS;
 import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.ElementType.BIT;
 import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.ElementType.BYTE;
 import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.ElementType.FLOAT;
+import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.IVF_FORMAT;
 import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.OVERSAMPLE_LIMIT;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
@@ -66,7 +73,7 @@ private DenseVectorFieldMapper.DenseVectorIndexOptions randomIndexOptionsNonQuan
     }
 
     public static DenseVectorFieldMapper.DenseVectorIndexOptions randomIndexOptionsAll() {
-        return randomFrom(
+        List<DenseVectorFieldMapper.DenseVectorIndexOptions> options = new ArrayList<>(Arrays.asList(
             new DenseVectorFieldMapper.HnswIndexOptions(randomIntBetween(1, 100), randomIntBetween(1, 10_000)),
             new DenseVectorFieldMapper.Int8HnswIndexOptions(
                 randomIntBetween(1, 100),
@@ -95,7 +102,18 @@ public static DenseVectorFieldMapper.DenseVectorIndexOptions randomIndexOptionsA
                 randomFrom((DenseVectorFieldMapper.RescoreVector) null, randomRescoreVector())
             ),
             new DenseVectorFieldMapper.BBQFlatIndexOptions(randomFrom((DenseVectorFieldMapper.RescoreVector) null, randomRescoreVector()))
-        );
+        ));
+
+        if(IVF_FORMAT.isEnabled()) {
+            options.add(
+                new DenseVectorFieldMapper.BBQIVFIndexOptions(
+                    randomIntBetween(MIN_VECTORS_PER_CLUSTER, MAX_VECTORS_PER_CLUSTER),
+                    randomFloatBetween(0.0f, 100.0f, true),
+                    randomFrom((DenseVectorFieldMapper.RescoreVector) null, randomRescoreVector()))
+            );
+        }
+
+        return randomFrom(options);
     }
 
     private DenseVectorFieldMapper.DenseVectorIndexOptions randomIndexOptionsHnswQuantized() {
@@ -244,7 +262,9 @@ public void testCreateNestedKnnQuery() {
             if (field.getIndexOptions().isFlat()) {
                 assertThat(query, instanceOf(DiversifyingParentBlockQuery.class));
             } else {
-                assertTrue(query instanceof DiversifyingChildrenFloatKnnVectorQuery || query instanceof PatienceKnnVectorQuery);
+                assertTrue(query instanceof DiversifyingChildrenFloatKnnVectorQuery
+                    || query instanceof PatienceKnnVectorQuery
+                    || query instanceof DiversifyingChildrenIVFKnnFloatVectorQuery);
             }
         }
         {
@@ -475,7 +495,9 @@ public void testCreateKnnQueryMaxDims() {
             if (fieldWith4096dims.getIndexOptions().isFlat()) {
                 assertThat(query, instanceOf(DenseVectorQuery.Floats.class));
             } else {
-                assertTrue(query instanceof KnnFloatVectorQuery || query instanceof PatienceKnnVectorQuery);
+                assertTrue(query instanceof KnnFloatVectorQuery
+                    || query instanceof PatienceKnnVectorQuery
+                    || query instanceof IVFKnnFloatVectorQuery);
             }
         }
 
@@ -792,7 +814,7 @@ private static void checkRescoreQueryParameters(
         DenseVectorFieldType fieldType,
         int k,
         int candidates,
-        float visitPercentage,
+        Float visitPercentage,
         float oversample,
         int expectedK,
         int expectedCandidates,

server/src/test/java/org/elasticsearch/search/retriever/KnnRetrieverBuilderParsingTests.java

@@ -52,7 +52,7 @@ public static KnnRetrieverBuilder createRandomKnnRetrieverBuilder() {
         float[] vector = randomVector(dim);
         int k = randomIntBetween(1, 100);
         int numCands = randomIntBetween(k + 20, 1000);
-        float visitPercentage = randomFloatBetween(0.0f, 100.0f, true);
+        Float visitPercentage = randomBoolean() ? null : randomFloatBetween(0.0f, 100.0f, true);
         Float similarity = randomBoolean() ? null : randomFloat();
         RescoreVectorBuilder rescoreVectorBuilder = randomBoolean()
             ? null

server/src/test/java/org/elasticsearch/search/retriever/RankDocsRetrieverBuilderTests.java

@@ -70,7 +70,7 @@ private List<RetrieverBuilder> innerRetrievers(QueryRewriteContext queryRewriteC
                     null,
                     randomInt(10),
                     randomIntBetween(10, 100),
-                    randomFloatBetween(0.0f, 100.0f, true),
+                    randomBoolean() ? null : randomFloatBetween(0.0f, 100.0f, true),
                     randomBoolean() ? null : new RescoreVectorBuilder(randomFloatBetween(1.0f, 10.0f, false)),
                     randomFloat()
                 );

server/src/test/java/org/elasticsearch/search/vectors/AbstractKnnVectorQueryBuilderTestCase.java

@@ -146,7 +146,7 @@ protected KnnVectorQueryBuilder doCreateTestQueryBuilder() {
         String fieldName = randomBoolean() ? VECTOR_FIELD : VECTOR_ALIAS_FIELD;
         int k = randomIntBetween(1, 100);
         int numCands = randomIntBetween(k + 20, 1000);
-        float visitPercentage = randomFloatBetween(0.0f, 100.0f, true);
+        Float visitPercentage = randomBoolean() ? null : randomFloatBetween(0.0f, 100.0f, true);
         KnnVectorQueryBuilder queryBuilder = createKnnVectorQueryBuilder(
             fieldName,
             k,

server/src/test/java/org/elasticsearch/search/vectors/KnnSearchBuilderTests.java

@@ -52,7 +52,7 @@ public static KnnSearchBuilder randomTestInstance() {
         float[] vector = randomVector(dim);
         int k = randomIntBetween(1, 100);
         int numCands = randomIntBetween(k + 20, 1000);
-        float visitPercentage = randomFloatBetween(0.0f, 100.0f, true);
+        Float visitPercentage = randomBoolean() ? null : randomFloatBetween(0.0f, 100.0f, true);
         RescoreVectorBuilder rescoreVectorBuilder = randomBoolean()
             ? null
             : new RescoreVectorBuilder(randomFloatBetween(1.0f, 10.0f, false));
@@ -230,7 +230,7 @@ public void testToQueryBuilder() {
         float[] vector = randomVector(randomIntBetween(2, 30));
         int k = randomIntBetween(1, 100);
         int numCands = randomIntBetween(k, 1000);
-        float visitPercentage = randomFloatBetween(0.0f, 100.0f, true);
+        Float visitPercentage = randomBoolean() ? null : randomFloatBetween(0.0f, 100.0f, true);
         Float similarity = randomBoolean() ? null : randomFloat();
         RescoreVectorBuilder rescoreVectorBuilder = randomBoolean()
             ? null

server/src/test/java/org/elasticsearch/search/vectors/KnnSearchRequestParserTests.java

@@ -277,7 +277,7 @@ private KnnSearch randomKnnSearch() {
 
         int k = randomIntBetween(1, 100);
         int numCands = randomIntBetween(k, 1000);
-        float visitPercentage = randomFloatBetween(0.0f, 100.0f, true);
+        Float visitPercentage = randomBoolean() ? null : randomFloatBetween(0.0f, 100.0f, true);
         return new KnnSearch(field, vector, k, numCands, visitPercentage);
     }
 

test/framework/src/main/java/org/elasticsearch/search/RandomSearchRequestGenerator.java

@@ -266,7 +266,7 @@ public static SearchSourceBuilder randomSearchSourceBuilder(
                 }
                 int k = randomIntBetween(1, 100);
                 int numCands = randomIntBetween(k, 1000);
-                float visitPercentage = randomFloatBetween(0.0f, 100.0f, true);
+                Float visitPercentage = randomBoolean() ? null : randomFloatBetween(0.0f, 100.0f, true);
                 RescoreVectorBuilder rescoreVectorBuilder = randomBoolean()
                     ? null
                     : new RescoreVectorBuilder(randomFloatBetween(1.0f, 10.0f, false));