test: add query expr test cases (#36073)

1. query with expr under different scalar index types
2. test framework supports preparing one piece of data and multiple
parameter queries

Signed-off-by: wangting0128 <[email protected]>

tests/python_client/base/client_base.py

@@ -18,7 +18,7 @@
 from common import common_type as ct
 from common.common_params import IndexPrams
 
-from pymilvus import ResourceGroupInfo
+from pymilvus import ResourceGroupInfo, DataType
 
 
 class Base:
@@ -35,6 +35,7 @@ class Base:
     resource_group_list = []
     high_level_api_wrap = None
     skip_connection = False
+
     def setup_class(self):
         log.info("[setup_class] Start setup class...")
 
@@ -44,6 +45,9 @@ def teardown_class(self):
     def setup_method(self, method):
         log.info(("*" * 35) + " setup " + ("*" * 35))
         log.info("[setup_method] Start setup test case %s." % method.__name__)
+        self._setup_objects()
+
+    def _setup_objects(self):
         self.connection_wrap = ApiConnectionsWrapper()
         self.utility_wrap = ApiUtilityWrapper()
         self.collection_wrap = ApiCollectionWrapper()
@@ -57,7 +61,9 @@ def setup_method(self, method):
     def teardown_method(self, method):
         log.info(("*" * 35) + " teardown " + ("*" * 35))
         log.info("[teardown_method] Start teardown test case %s..." % method.__name__)
+        self._teardown_objects()
 
+    def _teardown_objects(self):
         try:
             """ Drop collection before disconnect """
             if not self.connection_wrap.has_connection(alias=DefaultConfig.DEFAULT_USING)[0]:
@@ -80,7 +86,8 @@ def teardown_method(self, method):
             rgs_list = self.utility_wrap.list_resource_groups()[0]
             for rg_name in self.resource_group_list:
                 if rg_name is not None and rg_name in rgs_list:
-                    rg = self.utility_wrap.describe_resource_group(name=rg_name, check_task=ct.CheckTasks.check_nothing)[0]
+                    rg = \
+                    self.utility_wrap.describe_resource_group(name=rg_name, check_task=ct.CheckTasks.check_nothing)[0]
                     if isinstance(rg, ResourceGroupInfo):
                         if rg.num_available_node > 0:
                             self.utility_wrap.transfer_node(source=rg_name,
@@ -266,9 +273,9 @@ def init_collection_general(self, prefix="test", insert_data=False, nb=ct.defaul
                                                                      primary_field=primary_field)
         if vector_data_type == ct.sparse_vector:
             default_schema = cf.gen_default_sparse_schema(auto_id=auto_id, primary_field=primary_field,
-                                                                     enable_dynamic_field=enable_dynamic_field,
-                                                                     with_json=with_json,
-                                                                     multiple_dim_array=multiple_dim_array)
+                                                          enable_dynamic_field=enable_dynamic_field,
+                                                          with_json=with_json,
+                                                          multiple_dim_array=multiple_dim_array)
         if is_all_data_type:
             default_schema = cf.gen_collection_schema_all_datatype(auto_id=auto_id, dim=dim,
                                                                    primary_field=primary_field,
@@ -390,7 +397,8 @@ def init_user_with_privilege(self, privilege_object, object_name, privilege, db_
         self.utility_wrap.create_role()
 
         # grant privilege to the role
-        self.utility_wrap.role_grant(object=privilege_object, object_name=object_name, privilege=privilege, db_name=db_name)
+        self.utility_wrap.role_grant(object=privilege_object, object_name=object_name, privilege=privilege,
+                                     db_name=db_name)
 
         # bind the role to the user
         self.utility_wrap.role_add_user(tmp_user)
@@ -417,3 +425,54 @@ def show_indexes(self, collection_obj: ApiCollectionWrapper = None):
         indexes = {n.field_name: n.params for n in self.collection_wrap.indexes}
         log.info("[TestcaseBase] Collection: `{0}` index: {1}".format(collection_obj.name, indexes))
         return indexes
+
+
+class TestCaseClassBase(TestcaseBase):
+    """
+    Setup objects on class
+    """
+
+    def setup_class(self):
+        log.info("[setup_class] " + " Start setup class ".center(100, "~"))
+        self._setup_objects(self)
+
+    def teardown_class(self):
+        log.info("[teardown_class]" + " Start teardown class ".center(100, "~"))
+        self._teardown_objects(self)
+
+    def setup_method(self, method):
+        log.info(" setup ".center(80, "*"))
+        log.info("[setup_method] Start setup test case %s." % method.__name__)
+
+    def teardown_method(self, method):
+        log.info(" teardown ".center(80, "*"))
+        log.info("[teardown_method] Start teardown test case %s..." % method.__name__)
+
+    @property
+    def all_scalar_fields(self):
+        dtypes = [DataType.INT8, DataType.INT16, DataType.INT32, DataType.INT64, DataType.VARCHAR, DataType.BOOL,
+                  DataType.FLOAT, DataType.DOUBLE]
+        dtype_names = [f"{n.name}" for n in dtypes] + [f"ARRAY_{n.name}" for n in dtypes] + [DataType.JSON.name]
+        return dtype_names
+
+    @property
+    def all_index_scalar_fields(self):
+        return list(set(self.all_scalar_fields) - {DataType.JSON.name})
+
+    @property
+    def inverted_support_dtype_names(self):
+        return self.all_index_scalar_fields
+
+    @property
+    def inverted_not_support_dtype_names(self):
+        return [DataType.JSON.name]
+
+    @property
+    def bitmap_support_dtype_names(self):
+        dtypes = [DataType.INT8, DataType.INT16, DataType.INT32, DataType.INT64, DataType.BOOL, DataType.VARCHAR]
+        dtype_names = [f"{n.name}" for n in dtypes] + [f"ARRAY_{n.name}" for n in dtypes]
+        return dtype_names
+
+    @property
+    def bitmap_not_support_dtype_names(self):
+        return list(set(self.all_scalar_fields) - set(self.bitmap_support_dtype_names))

tests/python_client/common/code_mapping.py

@@ -38,3 +38,7 @@ class IndexErrorMessage(ExceptionsMessage):
     VectorMetricTypeExist = "metric type not set for vector index"
     CheckBitmapIndex = "bitmap index are only supported on bool, int, string and array field"
     CheckBitmapOnPK = "create bitmap index on primary key not supported"
+
+
+class QueryErrorMessage(ExceptionsMessage):
+    ParseExpressionFailed = "failed to create query plan: cannot parse expression: "

tests/python_client/common/common_func.py

@@ -21,6 +21,8 @@
 from customize.milvus_operator import MilvusOperator
 import pickle
 fake = Faker()
+
+from common.common_params import Expr
 """" Methods of processing data """
 
 
@@ -1293,6 +1295,11 @@ def gen_data_by_collection_field(field, nb=None, start=None):
             if nb is None:
                 return [np.float32(random.random()) for _ in range(max_capacity)]
             return [[np.float32(random.random()) for _ in range(max_capacity)] for _ in range(nb)]
+        if element_type == DataType.DOUBLE:
+            if nb is None:
+                return [np.float64(random.random()) for _ in range(max_capacity)]
+            return [[np.float64(random.random()) for _ in range(max_capacity)] for _ in range(nb)]
+
         if element_type == DataType.VARCHAR:
             max_length = field.params['max_length']
             max_length = min(20, max_length - 1)
@@ -1335,6 +1342,24 @@ def gen_values(schema: CollectionSchema, nb, start_id=0, default_values: dict =
     return data
 
 
+def gen_field_values(schema: CollectionSchema, nb, start_id=0, default_values: dict = {}) -> dict:
+    """
+    generate default value according to the collection fields,
+    which can replace the value of the specified field
+
+    return: <dict>
+        <field name>: <value list>
+    """
+    data = {}
+    for field in schema.fields:
+        default_value = default_values.get(field.name, None)
+        if default_value is not None:
+            data[field.name] = default_value
+        elif field.auto_id is False:
+            data[field.name] = gen_data_by_collection_field(field, nb, start_id * nb)
+    return data
+
+
 def gen_json_files_for_bulk_insert(data, schema, data_dir):
     for d in data:
         if len(d) > 0:
@@ -1746,6 +1771,48 @@ def gen_integer_overflow_expressions():
     return expressions
 
 
+def gen_modulo_expression(expr_fields):
+    exprs = []
+    for field in expr_fields:
+        exprs.extend([
+            (Expr.EQ(Expr.MOD(field, 10).subset, 1).value, field),
+            (Expr.LT(Expr.MOD(field, 17).subset, 9).value, field),
+            (Expr.LE(Expr.MOD(field, 100).subset, 50).value, field),
+            (Expr.GT(Expr.MOD(field, 50).subset, 40).value, field),
+            (Expr.GE(Expr.MOD(field, 29).subset, 15).value, field),
+            (Expr.NE(Expr.MOD(field, 29).subset, 10).value, field),
+        ])
+    return exprs
+
+
+def gen_varchar_expression(expr_fields):
+    exprs = []
+    for field in expr_fields:
+        exprs.extend([
+            (Expr.like(field, "a%").value, field, r'^a.*'),
+            (Expr.LIKE(field, "%b").value, field, r'.*b$'),
+            (Expr.AND(Expr.like(field, "%b").subset, Expr.LIKE(field, "z%").subset).value, field, r'^z.*b$'),
+            (Expr.And(Expr.like(field, "i%").subset, Expr.LIKE(field, "%j").subset).value, field, r'^i.*j$'),
+            (Expr.OR(Expr.like(field, "%h%").subset, Expr.LIKE(field, "%jo").subset).value, field, fr'(?:h.*|.*jo$)'),
+            (Expr.Or(Expr.like(field, "ip%").subset, Expr.LIKE(field, "%yu%").subset).value, field, fr'(?:^ip.*|.*yu)'),
+        ])
+    return exprs
+
+
+def gen_number_operation(expr_fields):
+    exprs = []
+    for field in expr_fields:
+        exprs.extend([
+            (Expr.LT(Expr.ADD(field, 23), 100).value, field),
+            (Expr.LT(Expr.ADD(-23, field), 121).value, field),
+            (Expr.LE(Expr.SUB(field, 123), 99).value, field),
+            (Expr.GT(Expr.MUL(field, 2), 88).value, field),
+            (Expr.GT(Expr.MUL(3, field), 137).value, field),
+            (Expr.GE(Expr.DIV(field, 30), 20).value, field),
+        ])
+    return exprs
+
+
 def l2(x, y):
     return np.linalg.norm(np.array(x) - np.array(y))
 

tests/python_client/common/common_params.py

@@ -1,6 +1,8 @@
 from dataclasses import dataclass
 from typing import List, Dict
 
+from pymilvus import DataType
+
 """ Define param names"""
 
 
@@ -55,6 +57,10 @@ def subset(self):
     def __repr__(self):
         return self.expr
 
+    @property
+    def value(self):
+        return self.expr
+
 
 class Expr:
     # BooleanConstant: 'true' | 'True' | 'TRUE' | 'false' | 'False' | 'FALSE'
@@ -344,6 +350,10 @@ class DefaultScalarIndexParams:
     def Default(field: str):
         return {field: IndexPrams()}
 
+    @staticmethod
+    def list_default(fields: List[str]) -> Dict[str, IndexPrams]:
+        return {n: IndexPrams() for n in fields}
+
     @staticmethod
     def Trie(field: str):
         return {field: IndexPrams(index_type=IndexName.Trie)}
@@ -356,6 +366,10 @@ def STL_SORT(field: str):
     def INVERTED(field: str):
         return {field: IndexPrams(index_type=IndexName.INVERTED)}
 
+    @staticmethod
+    def list_inverted(fields: List[str]) -> Dict[str, IndexPrams]:
+        return {n: IndexPrams(index_type=IndexName.INVERTED) for n in fields}
+
     @staticmethod
     def BITMAP(field: str):
         return {field: IndexPrams(index_type=IndexName.BITMAP)}