Add test cases of AutoIndex master (#24563)

Signed-off-by: nico <cheng.yuan@zilliz.com>

tests/python_client/base/collection_wrapper.py

@@ -241,7 +241,7 @@ class ApiCollectionWrapper:
         return res, check_result
 
     @trace()
-    def create_index(self, field_name, index_params, index_name=None, check_task=None, check_items=None, **kwargs):
+    def create_index(self, field_name, index_params=None, index_name=None, check_task=None, check_items=None, **kwargs):
         disktimeout = 600
         timeout = kwargs.get("timeout", disktimeout * 2)
         index_name = INDEX_NAME if index_name is None else index_name

tests/python_client/common/common_func.py

@@ -550,6 +550,23 @@ def gen_simple_index():
     return index_params
 
 
+def gen_autoindex_params():
+    index_params = [
+        {},
+        {"metric_type": "IP"},
+        {"metric_type": "L2"},
+        {"metric_type": "COSINE"},
+        {"index_type": "AUTOINDEX"},
+        {"index_type": "AUTOINDEX", "metric_type": "L2"},
+        {"index_type": "AUTOINDEX", "metric_type": "COSINE"},
+        {"index_type": "IVF_FLAT", "metric_type": "L2", "nlist": "1024", "m": "100"},
+        {"index_type": "DISKANN", "metric_type": "L2"},
+        {"index_type": "IVF_PQ", "nlist": "128", "m": "16", "nbits": "8", "metric_type": "IP"},
+        {"index_type": "IVF_SQ8", "nlist": "128", "metric_type": "COSINE"}
+    ]
+    return index_params
+
+
 def gen_invalid_field_types():
     field_types = [
         6,
@@ -627,6 +644,17 @@ def gen_search_param(index_type, metric_type="L2"):
     return search_params
 
 
+def gen_autoindex_search_params():
+    search_params = [
+        {},
+        {"metric_type": "IP"},
+        {"nlist": "1024"},
+        {"efSearch": "100"},
+        {"search_k": "1000"}
+    ]
+    return search_params
+
+
 def gen_invalid_search_param(index_type, metric_type="L2"):
     search_params = []
     if index_type in ["FLAT", "IVF_FLAT", "IVF_SQ8", "IVF_PQ"] \

tests/python_client/testcases/test_index.py

@@ -1,5 +1,6 @@
 from time import sleep
 import pytest
+import copy
 
 from base.client_base import TestcaseBase
 from base.index_wrapper import ApiIndexWrapper
@@ -18,6 +19,7 @@ prefix = "index"
 default_schema = cf.gen_default_collection_schema()
 default_field_name = ct.default_float_vec_field_name
 default_index_params = {"index_type": "IVF_SQ8", "metric_type": "L2", "params": {"nlist": 64}}
+default_autoindex_params = {"index_type": "AUTOINDEX", "metric_type": "IP"}
 
 # copied from pymilvus
 uid = "test_index"
@@ -1287,6 +1289,20 @@ class TestIndexInvalid(TestcaseBase):
                                   check_items={"err_code": 1,
                                                "err_msg": "invalid index params"})
 
+    @pytest.mark.tags(CaseLabel.L1)
+    def test_create_index_json(self):
+        """
+        target: test create index on json fields
+        method: 1.create collection, and create index
+        expected: create index raise an error
+        """
+        collection_w, _, _, insert_ids = self.init_collection_general(prefix, True,
+                                                                      dim=ct.default_dim, is_index=False)[0:4]
+        collection_w.create_index(ct.default_json_field_name, index_params=ct.default_flat_index,
+                                  check_task=CheckTasks.err_res,
+                                  check_items={ct.err_code: 1,
+                                               ct.err_msg: "create index on json field is not supported"})
+
 
 @pytest.mark.tags(CaseLabel.GPU)
 class TestNewIndexAsync(TestcaseBase):
@@ -1862,16 +1878,67 @@ class TestIndexDiskann(TestcaseBase):
                                   check_items={ct.err_code: 1,
                                                ct.err_msg: "invalid index params"})
 
+
+@pytest.mark.tags(CaseLabel.GPU)
+class TestAutoIndex(TestcaseBase):
+    """ Test case of Auto index """
+
+    @pytest.mark.tags(CaseLabel.L0)
+    def test_create_autoindex_with_no_params(self):
+        """
+        target: test create auto index with no params
+        method: create index with only one field name
+        expected: create successfully
+        """
+        collection_w = self.init_collection_general(prefix, is_index=False)[0]
+        collection_w.create_index(field_name)
+        actual_index_params = collection_w.index()[0].params
+        assert default_autoindex_params == actual_index_params
+
     @pytest.mark.tags(CaseLabel.L1)
-    def test_create_index_json(self):
+    @pytest.mark.parametrize("index_params", cf.gen_autoindex_params())
+    def test_create_autoindex_with_params(self, index_params):
         """
-        target: test create index on json fields
+        target: test create auto index with params
-        method: 1.create collection, and create index
+        method: create index with different params
-        expected: create index raise an error
+        expected: create successfully
         """
-        collection_w, _, _, insert_ids = self.init_collection_general(prefix, True,
+        collection_w = self.init_collection_general(prefix, is_index=False)[0]
-                                                                      dim=ct.default_dim, is_index=False)[0:4]
+        collection_w.create_index(field_name, index_params)
-        collection_w.create_index(ct.default_json_field_name, index_params=ct.default_flat_index,
+        actual_index_params = collection_w.index()[0].params
+        log.info(collection_w.index()[0].params)
+        expect_autoindex_params = copy.copy(default_autoindex_params)
+        if index_params.get("index_type"):
+            if index_params["index_type"] != 'AUTOINDEX':
+                expect_autoindex_params = index_params
+        if index_params.get("metric_type"):
+            expect_autoindex_params["metric_type"] = index_params["metric_type"]
+        assert actual_index_params == expect_autoindex_params
+
+    @pytest.mark.tags(CaseLabel.L2)
+    def test_create_autoindex_with_invalid_params(self):
+        """
+        target: test create auto index with invalid params
+        method: create index with invalid params
+        expected: raise exception
+        """
+        collection_w = self.init_collection_general(prefix, is_index=False)[0]
+        index_params = {"metric_type": "L2", "nlist": "1024", "M": "100"}
+        collection_w.create_index(field_name, index_params,
                                   check_task=CheckTasks.err_res,
-                                  check_items={ct.err_code: 1,
+                                  check_items={"err_code": 1,
-                                               ct.err_msg: "create index on json field is not supported"})
+                                               "err_msg": "only metric type can be "
+                                                          "passed when use AutoIndex"})
+
+    @pytest.mark.tags(CaseLabel.L2)
+    def test_create_autoindex_on_binary_vectors(self):
+        """
+        target: test create auto index on binary vectors
+        method: create index on binary vectors
+        expected: raise exception
+        """
+        collection_w = self.init_collection_general(prefix, is_binary=True, is_index=False)[0]
+        collection_w.create_index(binary_field_name, {},
+                                  check_task=CheckTasks.err_res,
+                                  check_items={"err_code": 1,
+                                               "err_msg": "float vector is only supported"})

tests/python_client/testcases/test_search.py

@@ -2030,6 +2030,30 @@ class TestCollectionSearch(TestcaseBase):
                                              "ids": insert_ids,
                                              "limit": default_limit,
                                              "_async": _async})
+
+    @pytest.mark.tags(CaseLabel.L1)
+    @pytest.mark.parametrize("search_params", cf.gen_autoindex_search_params())
+    @pytest.mark.skip("issue #24533 #24555")
+    def test_search_default_search_params_fit_for_autoindex(self, search_params, auto_id, _async):
+        """
+        target: test search using autoindex
+        method: test search using autoindex and its corresponding search params
+        expected: search successfully
+        """
+        # 1. initialize with data
+        collection_w = self.init_collection_general(prefix, True, auto_id=auto_id, is_index=False)[0]
+        # 2. create index and load
+        collection_w.create_index("float_vector", {})
+        collection_w.load()
+        # 3. search
+        log.info("Searching with search params: {}".format(search_params))
+        collection_w.search(vectors[:default_nq], default_search_field,
+                            search_params, default_limit,
+                            default_search_exp, _async=_async,
+                            check_task=CheckTasks.check_search_results,
+                            check_items={"nq": default_nq,
+                                         "limit": default_limit,
+                                         "_async": _async})
             
     @pytest.mark.tags(CaseLabel.L2)
     @pytest.mark.tags(CaseLabel.GPU)