From a2961b264aa3d5e95fcb97f59051282dac194032 Mon Sep 17 00:00:00 2001
From: wangting0128 <ting.wang@zilliz.com>
Date: Fri, 27 Sep 2024 11:48:31 +0800
Subject: [PATCH] test: add parition key & bitmap test case

Signed-off-by: wangting0128 <ting.wang@zilliz.com>
---
 tests/python_client/common/common_func.py     |  4 ++
 .../testcases/test_mix_scenes.py              | 52 +++++++++++++++++++
 2 files changed, 56 insertions(+)

diff --git a/tests/python_client/common/common_func.py b/tests/python_client/common/common_func.py
index e41a47baa007d..9e72939b807de 100644
--- a/tests/python_client/common/common_func.py
+++ b/tests/python_client/common/common_func.py
@@ -3004,6 +3004,9 @@ def set_collection_schema(fields: list, field_params: dict = {}, **kwargs):
                 is_primary: bool
                 description: str
                 max_length: int = 65535
+            varchar_2:
+                max_length: int = 100
+                is_partition_key: bool
             array_int8_1:
                 max_capacity: int = 100
             array_varchar_1:
@@ -3016,6 +3019,7 @@ def set_collection_schema(fields: list, field_params: dict = {}, **kwargs):
             primary_field: str
             auto_id: bool
             enable_dynamic_field: bool
+            num_partitions: int
     """
     field_schemas = [set_field_schema(field=field, params=field_params.get(field, {})) for field in fields]
     return ApiCollectionSchemaWrapper().init_collection_schema(fields=field_schemas, **kwargs)[0]
diff --git a/tests/python_client/testcases/test_mix_scenes.py b/tests/python_client/testcases/test_mix_scenes.py
index ac96bc5c1b02d..79fb9928025c4 100644
--- a/tests/python_client/testcases/test_mix_scenes.py
+++ b/tests/python_client/testcases/test_mix_scenes.py
@@ -1884,6 +1884,58 @@ def test_bitmap_offset_cache_and_mmap(self, request):
             check_items={"nq": nq, "ids": insert_data.get(DataType.INT64.name), "limit": limit,
                          "output_fields": scalar_fields})
 
+    @pytest.mark.tags(CaseLabel.L2)
+    @pytest.mark.parametrize("scalar_field, data_type, expr_data", [('INT64', 'int', 3), ('VARCHAR', 'str', '3')])
+    def test_bitmap_partition_keys(self, request, scalar_field, data_type, expr_data):
+        """
+        target:
+            1. build BITMAP index on partition key field
+        method:
+            1. create a collection with scalar field that enable partition key
+            2. insert some data and build BITMAP index
+            4. load collection
+            5. query via partition key field expr
+        expected:
+            1. build index and query are successful
+        """
+        # init params
+        collection_name, primary_field, nb = f"{request.function.__name__}_{scalar_field}", "int64_pk", 10000
+
+        # connect to server before testing
+        self._connect()
+
+        # create a collection with fields that can build `BITMAP` index
+        self.collection_wrap.init_collection(
+            name=collection_name,
+            schema=cf.set_collection_schema(
+                fields=[primary_field, DataType.FLOAT_VECTOR.name, scalar_field],
+                field_params={primary_field: FieldParams(is_primary=True).to_dict,
+                              scalar_field: FieldParams(is_partition_key=True).to_dict},
+            )
+        )
+
+        # prepare data (> 1024 triggering index building)
+        self.collection_wrap.insert(data=cf.gen_values(self.collection_wrap.schema, nb=nb, default_values={
+            scalar_field: [eval(f"{data_type}({random.randint(1, 4)})") for _ in range(nb)]
+        }), check_task=CheckTasks.check_insert_result)
+
+        # flush collection, segment sealed
+        self.collection_wrap.flush()
+
+        # build `BITMAP` index
+        self.build_multi_index(index_params={
+            **DefaultVectorIndexParams.HNSW(DataType.FLOAT_VECTOR.name),
+            **DefaultScalarIndexParams.BITMAP(scalar_field)
+        })
+
+        # load collection
+        self.collection_wrap.load()
+
+        # query before upsert
+        expr = f'{scalar_field} == {expr_data}' if scalar_field == 'INT64' else f'{scalar_field} == "{expr_data}"'
+        res, _ = self.collection_wrap.query(expr=expr, output_fields=[scalar_field], limit=100)
+        assert set([r.get(scalar_field) for r in res]) == {expr_data}
+
 
 @pytest.mark.xdist_group("TestGroupSearch")
 class TestGroupSearch(TestCaseClassBase):