Browse Source

add BITMAP index cases (#971)

Signed-off-by: wangting0128 <ting.wang@zilliz.com>
wt 1 year ago
parent
commit
63553c503e

+ 1 - 0
src/main/java/io/milvus/v2/common/IndexParam.java

@@ -86,6 +86,7 @@ public class IndexParam {
         // Only for scalar type field
         STL_SORT(200), // only for numeric type field
         INVERTED(201), // works for all scalar fields except JSON type field
+        BITMAP(202), // works for all scalar fields except JSON, FLOAT and DOUBLE type fields
 
         // Only for sparse vectors
         SPARSE_INVERTED_INDEX(300),

+ 56 - 8
tests/milvustestv2/src/main/java/com/zilliz/milvustestv2/common/CommonFunction.java

@@ -16,6 +16,7 @@ import io.milvus.v2.service.collection.request.DescribeCollectionReq;
 import io.milvus.v2.service.collection.request.LoadCollectionReq;
 import io.milvus.v2.service.collection.response.DescribeCollectionResp;
 import io.milvus.v2.service.index.request.CreateIndexReq;
+import io.milvus.v2.service.index.request.DropIndexReq;
 import io.milvus.v2.service.partition.request.CreatePartitionReq;
 import io.milvus.v2.service.vector.request.AnnSearchReq;
 import io.milvus.v2.service.vector.request.InsertReq;
@@ -29,6 +30,7 @@ import lombok.extern.slf4j.Slf4j;
 import javax.annotation.Nullable;
 import java.nio.ByteBuffer;
 import java.util.*;
+import java.util.stream.Collectors;
 
 /**
  * @Author yongpeng.li
@@ -647,6 +649,7 @@ public class CommonFunction {
 
     /**
      * 创建通用的collection方法,支持多个filed,多个向量
+     *
      * @param collectionName collection 可不传
      * @param pkDataType     主键类型
      * @param enableDynamic  是否开启动态列
@@ -683,6 +686,7 @@ public class CommonFunction {
 
     /**
      * 遍历fieldParamList生成对应的schema
+     *
      * @param fieldParamList field字段集合
      * @return List<CreateCollectionReq.FieldSchema> 给创建collection提供
      */
@@ -838,7 +842,8 @@ public class CommonFunction {
 
     /**
      * 创建通用索引
-     * @param collection   collection name
+     *
+     * @param collection     collection name
      * @param fieldParamList field集合
      */
     public static void createCommonIndex(String collection, List<FieldParam> fieldParamList) {
@@ -863,13 +868,55 @@ public class CommonFunction {
                 .build());
     }
 
+    /**
+     * Create Scalar Indexes
+     *
+     * @param collection     collection name
+     * @param fieldParamList scalar fields
+     */
+    public static void createScalarCommonIndex(String collection, List<FieldParam> FieldParamList) {
+        List<IndexParam> indexParamList = new ArrayList<>();
+        for (FieldParam FieldParam : FieldParamList) {
+            IndexParam.IndexType indexType = FieldParam.getIndextype();
+            String fieldName = FieldParam.getFieldName();
+
+            IndexParam indexParam = IndexParam.builder()
+                    .fieldName(fieldName)
+                    .indexType(indexType)
+                    .indexName(fieldName)
+                    .build();
+            indexParamList.add(indexParam);
+        }
+
+        BaseTest.milvusClientV2.createIndex(CreateIndexReq.builder()
+                .collectionName(collection)
+                .indexParams(indexParamList)
+                .build());
+    }
+
+    /**
+     * Drop Scalar Indexes
+     *
+     * @param collection     collection name
+     * @param fieldParamList scalar fields
+     */
+    public static void dropScalarCommonIndex(String collection, List<FieldParam> FieldParamList) {
+        List<String> fieldNames = FieldParamList.stream().map(FieldParam::getFieldName).collect(Collectors.toList());
+        fieldNames.forEach(x -> BaseTest.milvusClientV2.dropIndex(DropIndexReq.builder()
+                .collectionName(collection)
+                .fieldName(x)
+                .indexName(x)
+                .build()));
+    }
+
 
     /**
      * 为多向量查询提供AnnSearch
+     *
      * @param fieldParam 字段参数
-     * @param nq 传入的向量数
-     * @param topK 查询数量
-     * @param expr 表达式
+     * @param nq         传入的向量数
+     * @param topK       查询数量
+     * @param expr       表达式
      * @return AnnSearchReq
      */
     public static AnnSearchReq provideAnnSearch(FieldParam fieldParam, int nq, int topK, String expr) {
@@ -885,6 +932,7 @@ public class CommonFunction {
 
     /**
      * 根据索引类型提供查询参数
+     *
      * @param indexType index type
      * @return String 查询参数
      */
@@ -910,17 +958,17 @@ public class CommonFunction {
                 extraParam = "{\"nlist\": 128}";
                 break;
             case SCANN:
-                extraParam="{\"nlist\":1024,\"with_raw_data\":"+true+"}";
+                extraParam = "{\"nlist\":1024,\"with_raw_data\":" + true + "}";
                 break;
             case GPU_IVF_FLAT:
-                extraParam="{\"nlist\": 64}";
+                extraParam = "{\"nlist\": 64}";
                 break;
             case GPU_IVF_PQ:
-                extraParam="{\"nlist\": 64, \"m\": 16, \"nbits\": 8}";
+                extraParam = "{\"nlist\": 64, \"m\": 16, \"nbits\": 8}";
                 break;
             case SPARSE_INVERTED_INDEX:
             case SPARSE_WAND:
-                extraParam="{\"drop_ratio_search\":0.2}";
+                extraParam = "{\"drop_ratio_search\":0.2}";
                 break;
             default:
                 extraParam = "{\"nlist\":128}";

+ 2 - 0
tests/milvustestv2/src/main/java/com/zilliz/milvustestv2/params/FieldParam.java

@@ -1,6 +1,7 @@
 package com.zilliz.milvustestv2.params;
 
 import io.milvus.v2.common.DataType;
+import io.milvus.v2.common.IndexParam;
 import lombok.Builder;
 import lombok.Data;
 
@@ -14,4 +15,5 @@ public class FieldParam {
     int maxCapacity;
     DataType elementType;
 
+    IndexParam.IndexType indextype;
 }

+ 67 - 11
tests/milvustestv2/src/test/java/com/zilliz/milvustestv2/index/CreateIndexTest.java

@@ -5,6 +5,7 @@ import com.google.common.collect.Lists;
 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 io.milvus.v2.common.DataType;
 import io.milvus.v2.common.IndexParam;
 import io.milvus.v2.service.collection.request.DropCollectionReq;
@@ -15,7 +16,10 @@ import io.milvus.v2.service.vector.request.InsertReq;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
+import lombok.extern.slf4j.Slf4j;
+import org.testng.annotations.*;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 
@@ -23,26 +27,63 @@ import java.util.List;
  * @Author yongpeng.li
  * @Date 2024/2/20 16:34
  */
+@Slf4j
 public class CreateIndexTest extends BaseTest {
     String newCollectionName;
+
+    @DataProvider(name = "multiScalar")
+    public Object[][] providerScalarIndex() {
+        return new Object[][]{
+                {new ArrayList<FieldParam>() {{
+                    add(FieldParam.builder().fieldName(CommonData.fieldInt8).indextype(IndexParam.IndexType.BITMAP).build());
+                    add(FieldParam.builder().fieldName(CommonData.fieldInt16).indextype(IndexParam.IndexType.STL_SORT).build());
+                    add(FieldParam.builder().fieldName(CommonData.fieldInt64).indextype(IndexParam.IndexType.INVERTED).build());
+                    add(FieldParam.builder().fieldName(CommonData.fieldVarchar).indextype(IndexParam.IndexType.TRIE).build());
+                }}
+                },
+                {new ArrayList<FieldParam>() {{
+                    add(FieldParam.builder().fieldName(CommonData.fieldInt16).indextype(IndexParam.IndexType.BITMAP).build());
+                }}
+                },
+                {new ArrayList<FieldParam>() {{
+                    add(FieldParam.builder().fieldName(CommonData.fieldInt32).indextype(IndexParam.IndexType.BITMAP).build());
+                }}
+                },
+                {new ArrayList<FieldParam>() {{
+                    add(FieldParam.builder().fieldName(CommonData.fieldInt64).indextype(IndexParam.IndexType.BITMAP).build());
+                }}
+                },
+                {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.BITMAP).build());
+//                    add(FieldParam.builder().fieldName(CommonData.fieldBool).indextype(IndexParam.IndexType.BITMAP).build());
+//                    add(FieldParam.builder().fieldName(CommonData.fieldArray).indextype(IndexParam.IndexType.BITMAP).build());
+                }}
+                },
+        };
+    }
+
     @BeforeClass(alwaysRun = true)
-    public void providerCollection(){
+    public void providerCollection() {
         newCollectionName = CommonFunction.createNewCollection(CommonData.dim, null, DataType.FloatVector);
-        List<JsonObject> jsonObjects = CommonFunction.generateDefaultData(CommonData.numberEntities, CommonData.dim,DataType.FloatVector);
+        List<JsonObject> jsonObjects = CommonFunction.generateDefaultData(CommonData.numberEntities * 10, CommonData.dim, DataType.FloatVector);
         milvusClientV2.insert(InsertReq.builder().collectionName(newCollectionName).data(jsonObjects).build());
     }
 
     @AfterClass(alwaysRun = true)
-    public void cleanTestData(){
+    public void cleanTestData() {
         milvusClientV2.dropCollection(DropCollectionReq.builder().collectionName(newCollectionName).build());
     }
 
-    @Test(description = "Create vector index",groups = {"Smoke"})
-    public void createVectorIndex(){
+    @Test(description = "Create vector index", groups = {"Smoke"})
+    public void createVectorIndex() {
         IndexParam indexParam = IndexParam.builder()
                 .fieldName(CommonData.fieldFloatVector)
                 .indexType(IndexParam.IndexType.AUTOINDEX)
-               .extraParams(CommonFunction.provideExtraParam(IndexParam.IndexType.AUTOINDEX))
+                .extraParams(CommonFunction.provideExtraParam(IndexParam.IndexType.AUTOINDEX))
                 .metricType(IndexParam.MetricType.L2)
                 .build();
         milvusClientV2.createIndex(CreateIndexReq.builder()
@@ -53,8 +94,8 @@ public class CreateIndexTest extends BaseTest {
 
     }
 
-    @Test(description = "Create scalar index",groups = {"Smoke"},dependsOnMethods = {"createVectorIndex"})
-    public void createScalarIndex(){
+    @Test(description = "Create scalar index", groups = {"Smoke"}, dependsOnMethods = {"createVectorIndex"})
+    public void createScalarIndex() {
         milvusClientV2.releaseCollection(ReleaseCollectionReq.builder().collectionName(newCollectionName).build());
         IndexParam indexParam = IndexParam.builder()
                 .fieldName(CommonData.fieldVarchar)
@@ -69,8 +110,8 @@ public class CreateIndexTest extends BaseTest {
                 .build());
     }
 
-    @Test(description = "Create scalar index",groups = {"Smoke"},dependsOnMethods = {"createVectorIndex"})
-    public void createMultiScalarIndex(){
+    @Test(description = "Create scalar index", groups = {"Smoke"}, dependsOnMethods = {"createVectorIndex"})
+    public void createMultiScalarIndex() {
         milvusClientV2.releaseCollection(ReleaseCollectionReq.builder().collectionName(newCollectionName).build());
         IndexParam indexParam1 = IndexParam.builder()
                 .fieldName(CommonData.fieldVarchar)
@@ -84,12 +125,27 @@ public class CreateIndexTest extends BaseTest {
                 .fieldName(CommonData.fieldInt8)
                 .indexType(IndexParam.IndexType.STL_SORT)
                 .build();
+        IndexParam indexParam4 = IndexParam.builder()
+                .fieldName(CommonData.fieldInt16)
+                .indexType(IndexParam.IndexType.BITMAP)
+                .build();
         milvusClientV2.createIndex(CreateIndexReq.builder()
                 .collectionName(newCollectionName)
-                .indexParams(Lists.newArrayList(indexParam1,indexParam2,indexParam3))
+                .indexParams(Lists.newArrayList(indexParam1, indexParam2, indexParam3, indexParam4))
                 .build());
         milvusClientV2.loadCollection(LoadCollectionReq.builder()
                 .collectionName(newCollectionName)
                 .build());
     }
+
+    //    wait for bug fixed: https://github.com/milvus-io/milvus/issues/34314
+    @Test(description = "Create scalar index", groups = {"Smoke"}, dependsOnMethods = {"createVectorIndex"}, dataProvider = "multiScalar")
+    public void createAllBitmapIndex(List<FieldParam> FieldParamList) {
+        milvusClientV2.releaseCollection(ReleaseCollectionReq.builder().collectionName(newCollectionName).build());
+        CommonFunction.dropScalarCommonIndex(newCollectionName, FieldParamList);
+        CommonFunction.createScalarCommonIndex(newCollectionName, FieldParamList);
+        milvusClientV2.loadCollection(LoadCollectionReq.builder().collectionName(newCollectionName).build());
+        milvusClientV2.releaseCollection(ReleaseCollectionReq.builder().collectionName(newCollectionName).build());
+        CommonFunction.dropScalarCommonIndex(newCollectionName, FieldParamList);
+    }
 }

+ 77 - 10
tests/milvustestv2/src/test/java/com/zilliz/milvustestv2/vectorOperation/SearchTest.java

@@ -5,6 +5,7 @@ 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 com.zilliz.milvustestv2.utils.GenerateUtil;
 import com.zilliz.milvustestv2.utils.MathUtil;
@@ -20,8 +21,11 @@ import io.milvus.v2.common.DataType;
 import io.milvus.v2.service.collection.request.DescribeCollectionReq;
 import io.milvus.v2.service.collection.request.DropCollectionReq;
 import io.milvus.v2.service.collection.request.GetCollectionStatsReq;
+import io.milvus.v2.common.IndexParam;
+import io.milvus.v2.service.collection.request.*;
 import io.milvus.v2.service.collection.response.DescribeCollectionResp;
 import io.milvus.v2.service.collection.response.GetCollectionStatsResp;
+import io.milvus.v2.service.index.request.CreateIndexReq;
 import io.milvus.v2.service.vector.request.InsertReq;
 import io.milvus.v2.service.vector.request.QueryReq;
 import io.milvus.v2.service.vector.request.SearchReq;
@@ -30,12 +34,15 @@ import io.milvus.v2.service.vector.request.data.FloatVec;
 import io.milvus.v2.service.vector.response.InsertResp;
 import io.milvus.v2.service.vector.response.QueryResp;
 import io.milvus.v2.service.vector.response.SearchResp;
+import lombok.extern.slf4j.Slf4j;
 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.Collections;
 import java.util.List;
 import java.util.Map;
 
@@ -43,9 +50,11 @@ import java.util.Map;
  * @Author yongpeng.li
  * @Date 2024/2/19 17:03
  */
+@Slf4j
 public class SearchTest extends BaseTest {
     int topK = 10;
     private MilvusServiceClient milvusServiceClient;
+    String newCollectionName;
 
     @DataProvider(name = "filterAndExcept")
     public Object[][] providerData() {
@@ -75,22 +84,22 @@ public class SearchTest extends BaseTest {
         return new Object[][]{
                 {CommonData.defaultFloatVectorCollection, DataType.FloatVector},
 //                {CommonData.defaultBinaryVectorCollection,DataType.BinaryVector},
-                {CommonData.defaultFloat16VectorCollection,DataType.Float16Vector},
-                {CommonData.defaultBFloat16VectorCollection,DataType.BFloat16Vector},
-                {CommonData.defaultSparseFloatVectorCollection,DataType.SparseFloatVector},
+                {CommonData.defaultFloat16VectorCollection, DataType.Float16Vector},
+                {CommonData.defaultBFloat16VectorCollection, DataType.BFloat16Vector},
+                {CommonData.defaultSparseFloatVectorCollection, DataType.SparseFloatVector},
         };
     }
 
-    @DataProvider(name="VectorTypeWithFilter")
-    public Object[][] providerVectorTypeWithFilter(){
-        Object[][] vectorType=new Object[][]{
+    @DataProvider(name = "VectorTypeWithFilter")
+    public Object[][] providerVectorTypeWithFilter() {
+        Object[][] vectorType = new Object[][]{
                 {DataType.FloatVector},
                 {DataType.BinaryVector},
                 {DataType.Float16Vector},
                 {DataType.BFloat16Vector},
                 {DataType.SparseFloatVector}
         };
-        Object[][] filter=new Object[][]{
+        Object[][] filter = new Object[][]{
                 {CommonData.fieldVarchar + " like \"%0\" ", topK},
                 {CommonData.fieldInt64 + " < 10 ", topK},
                 {CommonData.fieldInt64 + " != 10 ", topK},
@@ -125,6 +134,47 @@ public class SearchTest extends BaseTest {
         };
     }
 
+    @BeforeClass(alwaysRun = true)
+    public void providerCollection() {
+        newCollectionName = CommonFunction.createNewCollection(CommonData.dim, null, DataType.FloatVector);
+        List<JsonObject> jsonObjects = CommonFunction.generateDefaultData(CommonData.numberEntities * 10, CommonData.dim, DataType.FloatVector);
+        milvusClientV2.insert(InsertReq.builder().collectionName(newCollectionName).data(jsonObjects).build());
+    }
+
+    @AfterClass(alwaysRun = true)
+    public void cleanTestData() {
+        milvusClientV2.dropCollection(DropCollectionReq.builder().collectionName(newCollectionName).build());
+    }
+
+    @Test(description = "Create vector and scalar index", groups = {"Smoke"})
+    public void createVectorAndScalarIndex() {
+        // Build Vector index
+        IndexParam indexParam = IndexParam.builder()
+                .fieldName(CommonData.fieldFloatVector)
+                .indexType(IndexParam.IndexType.AUTOINDEX)
+                .extraParams(CommonFunction.provideExtraParam(IndexParam.IndexType.AUTOINDEX))
+                .metricType(IndexParam.MetricType.L2)
+                .build();
+        milvusClientV2.createIndex(CreateIndexReq.builder()
+                .collectionName(newCollectionName)
+                .indexParams(Collections.singletonList(indexParam))
+                .build());
+
+        // 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.BITMAP).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(newCollectionName, FieldParamList);
+        log.info("Create Scalar index done{}, scalar index:{}", newCollectionName, FieldParamList);
+        milvusClientV2.loadCollection(LoadCollectionReq.builder().collectionName(newCollectionName).build());
+    }
+
     @Test(description = "search float vector collection", groups = {"Smoke"}, dataProvider = "filterAndExcept")
     public void searchFloatVectorCollection(String filter, int expect) {
         List<BaseVector> data = CommonFunction.providerBaseVector(CommonData.nq, CommonData.dim, DataType.FloatVector);
@@ -259,19 +309,36 @@ public class SearchTest extends BaseTest {
     }
 
     @Test(description = "search group by field name", groups = {"Smoke"}, dataProvider = "VectorTypeList")
-    public void searchByGroupByField(String collectionName,DataType vectorType){
+    public void searchByGroupByField(String collectionName, DataType vectorType) {
         List<BaseVector> data = CommonFunction.providerBaseVector(CommonData.nq, CommonData.dim, vectorType);
         SearchResp search = milvusClientV2.search(SearchReq.builder()
                 .collectionName(collectionName)
                 .outputFields(Lists.newArrayList("*"))
                 .consistencyLevel(ConsistencyLevel.STRONG)
-                        .groupByFieldName(CommonData.fieldInt8)
+                .groupByFieldName(CommonData.fieldInt8)
                 .data(data)
                 .topK(1000)
                 .build());
         Assert.assertEquals(search.getSearchResults().size(), CommonData.nq);
-        if(vectorType!=DataType.SparseFloatVector) {
+        if (vectorType != DataType.SparseFloatVector) {
             Assert.assertEquals(search.getSearchResults().get(0).size(), 127);
         }
     }
+
+    @Test(description = "search scalar index collection", groups = {"Smoke"}, dependsOnMethods = {"createVectorAndScalarIndex"}, dataProvider = "filterAndExcept")
+    public void searchScalarIndexCollection(String filter, int expect) {
+        List<BaseVector> data = CommonFunction.providerBaseVector(CommonData.nq, CommonData.dim, DataType.FloatVector);
+        SearchResp search = milvusClientV2.search(SearchReq.builder()
+                .collectionName(CommonData.defaultFloatVectorCollection)
+                .filter(filter)
+                .outputFields(Lists.newArrayList("*"))
+                .consistencyLevel(ConsistencyLevel.STRONG)
+                .annsField(CommonData.fieldFloatVector)
+                .data(data)
+                .topK(topK)
+                .build());
+        System.out.println(search);
+        Assert.assertEquals(search.getSearchResults().size(), CommonData.nq);
+        Assert.assertEquals(search.getSearchResults().get(0).size(), expect);
+    }
 }