Skip to content

Commit

Permalink
test: add search group by test case for bitmap
Browse files Browse the repository at this point in the history
Signed-off-by: wangting0128 <[email protected]>
  • Loading branch information
wangting0128 committed Sep 23, 2024
1 parent 1fb8b46 commit fb025ed
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 6 deletions.
2 changes: 1 addition & 1 deletion tests/python_client/testcases/test_index.py
Original file line number Diff line number Diff line change
Expand Up @@ -2338,7 +2338,7 @@ def test_bitmap_on_primary_key_field(self, request, primary_field, auto_id):
1. create an empty collection
2. build `BITMAP` index on primary key field
expected:
1. Primary key filed does not support building bitmap index
1. Primary key field does not support building bitmap index
"""
# init params
collection_name = f"{request.function.__name__}_{primary_field}_{auto_id}"
Expand Down
40 changes: 35 additions & 5 deletions tests/python_client/testcases/test_mix_scenes.py
Original file line number Diff line number Diff line change
Expand Up @@ -586,6 +586,36 @@ def test_bitmap_index_query_count(self):
self.collection_wrap.query(expr='', output_fields=['count(*)'], check_task=CheckTasks.check_query_results,
check_items={"exp_res": [{"count(*)": self.nb}]})

@pytest.mark.tags(CaseLabel.L2)
@pytest.mark.parametrize("limit", [10, 1000])
@pytest.mark.parametrize("group_by_field", ['INT8', 'INT16', 'INT32', 'INT64', 'BOOL', 'VARCHAR'])
@pytest.mark.parametrize(
"dim, search_params, vector_field",
[(3, {"metric_type": MetricType.L2, "ef": 32}, DataType.FLOAT16_VECTOR.name),
(1000, {"metric_type": MetricType.IP, "drop_ratio_search": 0.2}, DataType.SPARSE_FLOAT_VECTOR.name)])
def test_bitmap_index_search_group_by(self, limit, group_by_field, dim, search_params, vector_field):
"""
target:
1. check search iterator with BITMAP index built on scalar fields
method:
1. prepare some data and build `BITMAP index` on scalar fields
2. search group by scalar fields and check result
expected:
1. search group by with BITMAP index
"""
res, _ = self.collection_wrap.search(cf.gen_vectors(nb=1, dim=dim, vector_data_type=vector_field), vector_field,
search_params, limit, group_by_field=group_by_field,
output_fields=[group_by_field])
output_values = [i.fields for r in res for i in r]

# check output field
assert len([True for i in output_values if set(i.keys()) != {group_by_field}]) == 0, f"res: {output_values}"

# check `group_by_field` field values are unique
values = [v for i in output_values for k, v in i.items()]

assert len(values) == len(set(values)), f"values: {values}, output_values:{output_values}"

@pytest.mark.tags(CaseLabel.L2)
@pytest.mark.parametrize("batch_size", [10, 1000])
def test_bitmap_index_search_iterator(self, batch_size):
Expand All @@ -601,7 +631,7 @@ def test_bitmap_index_search_iterator(self, batch_size):
search_params, vector_field = {"metric_type": "L2", "ef": 32}, DataType.FLOAT16_VECTOR.name
self.collection_wrap.search_iterator(
cf.gen_vectors(nb=1, dim=3, vector_data_type=vector_field), vector_field, search_params, batch_size,
expr='int64_pk > 15', check_task=CheckTasks.check_search_iterator, check_items={"batch_size": batch_size})
expr='INT16 > 15', check_task=CheckTasks.check_search_iterator, check_items={"batch_size": batch_size})

@pytest.mark.tags(CaseLabel.L2)
def test_bitmap_index_hybrid_search(self):
Expand Down Expand Up @@ -659,7 +689,7 @@ def setup_class(self):

# create a collection with fields
self.collection_wrap.init_collection(
name=cf.gen_unique_str("test_bitmap_index_dql_expr"),
name=cf.gen_unique_str("test_bitmap_index_offset_cache"),
schema=cf.set_collection_schema(
fields=[self.primary_field, DataType.FLOAT_VECTOR.name, *self().all_scalar_fields],
field_params={
Expand Down Expand Up @@ -825,7 +855,7 @@ def setup_class(self):

# create a collection with fields
self.collection_wrap.init_collection(
name=cf.gen_unique_str("test_bitmap_index_dql_expr"),
name=cf.gen_unique_str("test_bitmap_index_bitmap"),
schema=cf.set_collection_schema(
fields=[self.primary_field, DataType.FLOAT_VECTOR.name, *self().all_scalar_fields],
field_params={
Expand Down Expand Up @@ -991,7 +1021,7 @@ def setup_class(self):

# create a collection with fields
self.collection_wrap.init_collection(
name=cf.gen_unique_str("test_bitmap_index_unicode"),
name=cf.gen_unique_str("test_index_unicode_string"),
schema=cf.set_collection_schema(
fields=[self.primary_field, DataType.FLOAT_VECTOR.name,
f"{DataType.VARCHAR.name}_BITMAP", f"{DataType.ARRAY.name}_{DataType.VARCHAR.name}_BITMAP",
Expand Down Expand Up @@ -1065,7 +1095,7 @@ def test_index_unicode_string_query(self, expr, expr_field, limit, rex):
@pytest.mark.tags(CaseLabel.L2)
@pytest.mark.parametrize("obj", cf.gen_varchar_unicode_expression_array(
['ARRAY_VARCHAR_BITMAP', 'ARRAY_VARCHAR_INVERTED', 'ARRAY_VARCHAR_NoIndex']))
@pytest.mark.parametrize("limit", [1])
@pytest.mark.parametrize("limit", [1, 10, 3000])
def test_index_unicode_string_array_query(self, limit, obj):
"""
target:
Expand Down

0 comments on commit fb025ed

Please sign in to comment.