Add query and search related cases for null and default value support (#1108)

* Add query and search related cases for null and default value support

Signed-off-by: qixuan <673771573@qq.com>

* Add vector operation related cases for null and default value support

Signed-off-by: qixuan <673771573@qq.com>

---------

Signed-off-by: qixuan <673771573@qq.com>

tests/milvustestv2/src/main/java/com/zilliz/milvustestv2/common/BaseTest.java

@@ -3,6 +3,7 @@ package com.zilliz.milvustestv2.common;
 import com.google.gson.JsonObject;
 import com.zilliz.milvustestv2.Milvustestv2Application;
 import com.zilliz.milvustestv2.config.ConnectInfoConfig;
+import com.zilliz.milvustestv2.params.FieldParam;
 import com.zilliz.milvustestv2.utils.PropertyFilesUtil;
 import io.milvus.param.MetricType;
 import io.milvus.v2.client.ConnectConfig;
@@ -25,6 +26,7 @@ import org.testng.annotations.BeforeSuite;
 import org.testng.annotations.Parameters;
 
 import java.net.URI;
+import java.util.ArrayList;
 import java.util.List;
 
 /**
@@ -107,5 +109,4 @@ public class BaseTest extends AbstractTestNGSpringContextTests {
         CommonFunction.createNewCollection(CommonData.dim,CommonData.defaultSparseFloatVectorCollection, DataType.SparseFloatVector);
         CommonFunction.createIndexAndInsertAndLoad(CommonData.defaultSparseFloatVectorCollection,DataType.SparseFloatVector,true,CommonData.numberEntities);
     }
-
 }

tests/milvustestv2/src/main/java/com/zilliz/milvustestv2/common/CommonFunction.java

@@ -617,6 +617,68 @@ public class CommonFunction {
         return jsonList;
     }
 
+    /**
+     * 为collection提供导入含有NULL的数据，目前只支持行式插入
+     *
+     * @param num 数据量
+     * @param dim 维度
+     * @return List<JsonObject>
+     */
+    public static List<JsonObject> generateSimpleNullData(long startId,long num, int dim, DataType vectorType) {
+        List<JsonObject> jsonList = new ArrayList<>();
+        Random ran = new Random();
+        Gson gson = new Gson();
+        for (long i = startId; i < (num+startId); i++) {
+            JsonObject row = new JsonObject();
+            row.addProperty(CommonData.fieldInt64, i);
+            row.addProperty(CommonData.fieldInt32, (String) null);
+            row.addProperty(CommonData.fieldInt16, (int) i % 32767);
+            row.addProperty(CommonData.fieldInt8, (short) i % 127);
+            row.addProperty(CommonData.fieldBool, i % 2 == 0);
+            if (i % 2 == 1) {
+                row.addProperty(CommonData.fieldDouble, (double) i);
+                row.addProperty(CommonData.fieldVarchar, "Str" + i);
+                row.addProperty(CommonData.fieldFloat, (float) i);
+                row.add(CommonData.fieldArray, gson.toJsonTree(Arrays.asList(i, i + 1, i + 2)));
+            }
+            // 判断vectorType
+            if (vectorType == DataType.FloatVector) {
+                List<Float> vector = new ArrayList<>();
+                for (int k = 0; k < dim; ++k) {
+                    vector.add(ran.nextFloat());
+                }
+                row.add(CommonData.fieldFloatVector, gson.toJsonTree(vector));
+            }
+            if (vectorType == DataType.BinaryVector) {
+                row.add(CommonData.fieldBinaryVector, gson.toJsonTree(generateBinaryVector(dim).array()));
+            }
+            if (vectorType == DataType.Float16Vector) {
+                row.add(CommonData.fieldFloat16Vector, gson.toJsonTree(generateFloat16Vector(dim).array()));
+            }
+            if (vectorType == DataType.BFloat16Vector) {
+                row.add(CommonData.fieldBF16Vector, gson.toJsonTree(generateBF16Vector(dim).array()));
+            }
+            if (vectorType == DataType.SparseFloatVector) {
+                row.add(CommonData.fieldSparseVector, gson.toJsonTree(generateSparseVector(dim)));
+            }
+
+            JsonObject json = new JsonObject();
+            json.addProperty(CommonData.fieldInt64, (int) i % 32767);
+            json.addProperty(CommonData.fieldInt32, (int) i % 32767);
+            json.addProperty(CommonData.fieldDouble, (double) i);
+            json.add(CommonData.fieldArray, gson.toJsonTree(Arrays.asList(i, i + 1, i + 2)));
+            json.addProperty(CommonData.fieldBool, i % 2 == 0);
+            if ( i % 2 == 1) {
+                json.addProperty(CommonData.fieldVarchar, "Str" + i);
+            }
+            json.addProperty(CommonData.fieldFloat, (float) i);
+            row.add(CommonData.fieldJson, json);
+            jsonList.add(row);
+        }
+        return jsonList;
+    }
+
+
     /**
      * 快速创建一个collection，只有主键和向量字段
      *

tests/milvustestv2/src/test/java/com/zilliz/milvustestv2/vectorOperation/InsertTest.java

@@ -33,17 +33,19 @@ import java.util.List;
 public class InsertTest extends BaseTest {
     private String newCollectionName;
     private String simpleCollectionName;
+    private String nullableDefaultCollectionName;
 
     @BeforeClass(alwaysRun = true)
     public void providerCollection() {
         newCollectionName = CommonFunction.createNewCollection(CommonData.dim, null, DataType.FloatVector);
         simpleCollectionName = CommonFunction.createSimpleCollection(CommonData.dim, null);
+        nullableDefaultCollectionName = CommonFunction.createNewNullableDefaultValueCollection(CommonData.dim, null, DataType.FloatVector);
     }
-
     @AfterClass(alwaysRun = true)
     public void cleanTestData() {
         milvusClientV2.dropCollection(DropCollectionReq.builder().collectionName(newCollectionName).build());
         milvusClientV2.dropCollection(DropCollectionReq.builder().collectionName(simpleCollectionName).build());
+        milvusClientV2.dropCollection(DropCollectionReq.builder().collectionName(nullableDefaultCollectionName).build());
     }
 
     @DataProvider(name = "VectorTypeList")
@@ -88,4 +90,14 @@ public class InsertTest extends BaseTest {
         milvusClientV2.dropCollection(DropCollectionReq.builder().collectionName(newCollection).build());
     }
 
+    @Test(description = "insert nullable field collection test", groups = {"Smoke"}, dataProvider = "VectorTypeList")
+    public void insertNullableCollection(DataType dataType) {
+        String newCollection = CommonFunction.createNewNullableDefaultValueCollection(CommonData.dim, null, dataType);
+        List<JsonObject> jsonObjects = CommonFunction.generateSimpleNullData(0,CommonData.numberEntities, CommonData.dim, dataType);
+        InsertResp insert = milvusClientV2.insert(InsertReq.builder()
+                .collectionName(newCollection)
+                .data(jsonObjects).build());
+        Assert.assertEquals(insert.getInsertCnt(), CommonData.numberEntities);
+        milvusClientV2.dropCollection(DropCollectionReq.builder().collectionName(newCollection).build());
+    }
 }

tests/milvustestv2/src/test/java/com/zilliz/milvustestv2/vectorOperation/QueryTest.java

@@ -1,19 +1,30 @@
 package com.zilliz.milvustestv2.vectorOperation;
 
 import com.google.common.collect.Lists;
+import com.google.gson.JsonObject;
 import com.zilliz.milvustestv2.common.BaseTest;
 import com.zilliz.milvustestv2.common.CommonData;
+import com.zilliz.milvustestv2.common.CommonFunction;
+import com.zilliz.milvustestv2.params.FieldParam;
 import com.zilliz.milvustestv2.utils.DataProviderUtils;
 import io.milvus.v2.common.ConsistencyLevel;
 import io.milvus.v2.common.DataType;
+import io.milvus.v2.common.IndexParam;
+import io.milvus.v2.service.collection.request.DropCollectionReq;
 import io.milvus.v2.service.collection.request.GetCollectionStatsReq;
+import io.milvus.v2.service.collection.request.LoadCollectionReq;
 import io.milvus.v2.service.collection.response.GetCollectionStatsResp;
+import io.milvus.v2.service.vector.request.InsertReq;
 import io.milvus.v2.service.vector.request.QueryReq;
+import io.milvus.v2.service.vector.response.InsertResp;
 import io.milvus.v2.service.vector.response.QueryResp;
 import org.testng.Assert;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 
@@ -22,6 +33,9 @@ import java.util.List;
  * @Date 2024/2/19 17:03
  */
 public class QueryTest extends BaseTest {
+
+    String nullableDefaultCollectionName;
+    
     @DataProvider(name = "filterAndExcept")
     public Object[][] providerData() {
         return new Object[][]{
@@ -88,6 +102,47 @@ public class QueryTest extends BaseTest {
         return objects;
     }
 
+    @BeforeClass(alwaysRun = true)
+    public void providerCollection() {
+        nullableDefaultCollectionName = CommonFunction.createNewNullableDefaultValueCollection(CommonData.dim, null, DataType.FloatVector);
+        // insert data
+        List<JsonObject> jsonObjects = CommonFunction.generateSimpleNullData(0,CommonData.numberEntities, CommonData.dim,DataType.FloatVector);
+        InsertResp insert = milvusClientV2.insert(InsertReq.builder().collectionName(nullableDefaultCollectionName).data(jsonObjects).build());
+        CommonFunction.createVectorIndex(nullableDefaultCollectionName,CommonData.fieldFloatVector, IndexParam.IndexType.AUTOINDEX, IndexParam.MetricType.L2);
+        // Build Scalar Index
+        List<FieldParam> FieldParamList = new ArrayList<FieldParam>() {{
+            add(FieldParam.builder().fieldName(CommonData.fieldVarchar).indextype(IndexParam.IndexType.BITMAP).build());
+            add(FieldParam.builder().fieldName(CommonData.fieldInt8).indextype(IndexParam.IndexType.BITMAP).build());
+            add(FieldParam.builder().fieldName(CommonData.fieldInt16).indextype(IndexParam.IndexType.BITMAP).build());
+            add(FieldParam.builder().fieldName(CommonData.fieldInt32).indextype(IndexParam.IndexType.BITMAP).build());
+            add(FieldParam.builder().fieldName(CommonData.fieldInt64).indextype(IndexParam.IndexType.STL_SORT).build());
+            add(FieldParam.builder().fieldName(CommonData.fieldBool).indextype(IndexParam.IndexType.BITMAP).build());
+            add(FieldParam.builder().fieldName(CommonData.fieldArray).indextype(IndexParam.IndexType.BITMAP).build());
+        }};
+        CommonFunction.createScalarCommonIndex(nullableDefaultCollectionName, FieldParamList);
+        milvusClientV2.loadCollection(LoadCollectionReq.builder().collectionName(nullableDefaultCollectionName).build());
+        // create partition
+        CommonFunction.createPartition(nullableDefaultCollectionName,CommonData.partitionNameA);
+        List<JsonObject> jsonObjectsA = CommonFunction.generateSimpleNullData(0,CommonData.numberEntities*3, CommonData.dim,DataType.FloatVector);
+        milvusClientV2.insert(InsertReq.builder().collectionName(nullableDefaultCollectionName).partitionName(CommonData.partitionNameA).data(jsonObjectsA).build());
+    }
+
+    @AfterClass(alwaysRun = true)
+    public void cleanTestData() {
+        milvusClientV2.dropCollection(DropCollectionReq.builder().collectionName(nullableDefaultCollectionName).build());
+    }
+
+    @DataProvider(name = "queryNullableField")
+    private Object[][] provideNullableFieldQueryParams() {
+        return new Object[][]{
+                {CommonData.fieldInt32 + " == 1 ", CommonData.numberEntities*3},
+                {CommonData.fieldDouble + " > 1 ", CommonData.numberEntities*3/2 - 1},
+                {CommonData.fieldVarchar + " == \"1.0\" ", CommonData.numberEntities*3/2},
+                {CommonData.fieldFloat + " == 1.0 ", CommonData.numberEntities*3/2 + 1},
+                {"fieldJson[\"" + CommonData.fieldVarchar + "\"] in [\"Str1\", \"Str3\"]", 2},
+                {"ARRAY_CONTAINS(" + CommonData.fieldArray + ", 1)", 1},
+        };
+    }
 
     @Test(description = "query", groups = {"Smoke"}, dataProvider = "filterAndExcept")
     public void query(String filter, long expect) {
@@ -164,4 +219,14 @@ public class QueryTest extends BaseTest {
         Assert.assertEquals(query.getQueryResults().size(), expect);
     }
 
+    @Test(description = "query with nullable field", groups = {"Smoke"}, dataProvider = "queryNullableField")
+    public void queryByNullFilter(String filter, long expect) {
+        QueryResp query = milvusClientV2.query(QueryReq.builder()
+                .collectionName(nullableDefaultCollectionName)
+                .consistencyLevel(ConsistencyLevel.BOUNDED)
+                .outputFields(Lists.newArrayList("*"))
+                .filter(filter)
+                .build());
+        Assert.assertEquals(query.getQueryResults().size(), expect);
+    }
 }

tests/milvustestv2/src/test/java/com/zilliz/milvustestv2/vectorOperation/SearchTest.java

@@ -55,6 +55,7 @@ public class SearchTest extends BaseTest {
     int topK = 10;
     private MilvusServiceClient milvusServiceClient;
     String newCollectionName;
+    String nullableDefaultCollectionName;
 
     @DataProvider(name = "filterAndExcept")
     public Object[][] providerData() {
@@ -137,20 +138,49 @@ public class SearchTest extends BaseTest {
         };
     }
 
+    @DataProvider(name = "searchNullableField")
+    private Object[][] provideNullableFieldSearchParams() {
+        return new Object[][]{
+                {CommonData.fieldInt32 + " == 1 ", topK},
+                {CommonData.fieldDouble + " > 1 ", topK},
+                {CommonData.fieldVarchar + " == \"1.0\" ", topK},
+                {CommonData.fieldFloat + " == 1.0 ", topK},
+                {"fieldJson[\"" + CommonData.fieldVarchar + "\"] in [\"Str1\", \"Str3\"]", 2},
+                {"ARRAY_CONTAINS(" + CommonData.fieldArray + ", 1)", 1},
+        };
+    }
+
+    @DataProvider(name = "collectionNameList")
+    public Object[][] providerCollectionName() {
+        return new Object[][]{
+                {newCollectionName},
+                {nullableDefaultCollectionName},
+        };
+    }
+
     @BeforeClass(alwaysRun = true)
     public void providerCollection() {
         newCollectionName = CommonFunction.createNewCollection(CommonData.dim, null, DataType.FloatVector);
         List<JsonObject> jsonObjects = CommonFunction.generateDefaultData(0,CommonData.numberEntities * 10, CommonData.dim, DataType.FloatVector);
         milvusClientV2.insert(InsertReq.builder().collectionName(newCollectionName).data(jsonObjects).build());
+        nullableDefaultCollectionName = CommonFunction.createNewNullableDefaultValueCollection(CommonData.dim, null, DataType.FloatVector);
+        // insert data
+        List<JsonObject> jsonNullableObjects = CommonFunction.generateSimpleNullData(0,CommonData.numberEntities, CommonData.dim,DataType.FloatVector);
+        milvusClientV2.insert(InsertReq.builder().collectionName(nullableDefaultCollectionName).data(jsonNullableObjects).build());
+        // create partition
+        CommonFunction.createPartition(nullableDefaultCollectionName,CommonData.partitionNameA);
+        List<JsonObject> jsonObjectsA = CommonFunction.generateSimpleNullData(0,CommonData.numberEntities, CommonData.dim,DataType.FloatVector);
+        milvusClientV2.insert(InsertReq.builder().collectionName(nullableDefaultCollectionName).partitionName(CommonData.partitionNameA).data(jsonObjectsA).build());
     }
 
     @AfterClass(alwaysRun = true)
     public void cleanTestData() {
         milvusClientV2.dropCollection(DropCollectionReq.builder().collectionName(newCollectionName).build());
+        milvusClientV2.dropCollection(DropCollectionReq.builder().collectionName(nullableDefaultCollectionName).build());
     }
 
-    @Test(description = "Create vector and scalar index", groups = {"Smoke"})
-    public void createVectorAndScalarIndex() {
+    @Test(description = "Create vector and scalar index", groups = {"Smoke"}, dataProvider = "collectionNameList")
+    public void createVectorAndScalarIndex(String newCollectionName) {
         // Build Vector index
         IndexParam indexParam = IndexParam.builder()
                 .fieldName(CommonData.fieldFloatVector)
@@ -346,4 +376,23 @@ public class SearchTest extends BaseTest {
         Assert.assertEquals(search.getSearchResults().size(), CommonData.nq);
         Assert.assertEquals(search.getSearchResults().get(0).size(), expect);
     }
+
+    @Test(description = "search collection with nullable field", groups = {"Smoke"}, dependsOnMethods = {"createVectorAndScalarIndex"}, dataProvider = "searchNullableField")
+    public void searchNullableCollection(String filter, int expect) {
+        List<BaseVector> data = CommonFunction.providerBaseVector(CommonData.nq, CommonData.dim, DataType.FloatVector);
+        SearchReq searchParams = SearchReq.builder()
+                .collectionName(nullableDefaultCollectionName)
+                .filter(filter)
+                .outputFields(Lists.newArrayList("*"))
+                .consistencyLevel(ConsistencyLevel.STRONG)
+                .annsField(CommonData.fieldFloatVector)
+                .data(data)
+                .topK(topK)
+                .build();
+        System.out.println(searchParams);
+        SearchResp search = milvusClientV2.search(searchParams);
+        System.out.println(search);
+        Assert.assertEquals(search.getSearchResults().size(), CommonData.nq);
+        Assert.assertEquals(search.getSearchResults().get(0).size(), expect);
+    }
 }

tests/milvustestv2/src/test/java/com/zilliz/milvustestv2/vectorOperation/UpsertTest.java

@@ -34,17 +34,20 @@ import java.util.Map;
  */
 public class UpsertTest extends BaseTest {
     String newCollectionName;
+    String nullableDefaultCollectionName;
 
     @BeforeClass(alwaysRun = true)
     public void providerCollection() {
         newCollectionName = CommonFunction.createNewCollection(CommonData.dim, null, DataType.FloatVector);
         List<JsonObject> jsonObjects = CommonFunction.generateDefaultData(0,CommonData.numberEntities, CommonData.dim, DataType.FloatVector);
         milvusClientV2.insert(InsertReq.builder().collectionName(newCollectionName).data(jsonObjects).build());
+        nullableDefaultCollectionName = CommonFunction.createNewNullableDefaultValueCollection(CommonData.dim, null, DataType.FloatVector);
     }
 
     @AfterClass(alwaysRun = true)
     public void cleanTestData() {
         milvusClientV2.dropCollection(DropCollectionReq.builder().collectionName(newCollectionName).build());
+        milvusClientV2.dropCollection(DropCollectionReq.builder().collectionName(nullableDefaultCollectionName).build());
     }
 
     @DataProvider(name = "DifferentCollection")
@@ -114,4 +117,33 @@ public class UpsertTest extends BaseTest {
         System.out.println(upsert);
     }
 
+    @Test(description = "upsert nullable collection", groups = {"Smoke"}, dataProvider = "DifferentCollection")
+    public void nullableCollectionUpsert( DataType vectorType) {
+        String collectionName = CommonFunction.createNewNullableDefaultValueCollection(CommonData.dim, null, vectorType);
+        CommonFunction.createIndexAndInsertAndLoad(collectionName,vectorType,true,CommonData.numberEntities);
+
+        List<JsonObject> jsonObjects = CommonFunction.generateSimpleNullData(0,1, CommonData.dim, vectorType);
+        for (int i = 1; i < 10; i++) {
+            JsonObject jsonObject0 = jsonObjects.get(0).deepCopy();
+            jsonObject0.addProperty(CommonData.fieldInt64, i);
+            jsonObjects.add(jsonObject0);
+        }
+        UpsertResp upsert = milvusClientV2.upsert(UpsertReq.builder()
+                .collectionName(collectionName)
+                .data(jsonObjects)
+                .partitionName("_default")
+                .build());
+        System.out.println(upsert);
+        Assert.assertEquals(upsert.getUpsertCnt(), 10);
+
+        // query
+        QueryResp query = milvusClientV2.query(QueryReq.builder()
+                .collectionName(collectionName)
+                .filter(CommonData.fieldInt32 + " == 1")
+                .partitionNames(Lists.newArrayList(CommonData.defaultPartitionName))
+                .outputFields(Lists.newArrayList(CommonData.fieldInt64, CommonData.fieldInt32))
+                .consistencyLevel(ConsistencyLevel.STRONG).build());
+        Assert.assertEquals(query.getQueryResults().size(),10);
+        milvusClientV2.dropCollection(DropCollectionReq.builder().collectionName(collectionName).build());
+    }
 }