Skip to content

Commit

Permalink
test: Add query iterator test and seperate a new test file (#37085)
Browse files Browse the repository at this point in the history 
related issue: #37084

Signed-off-by: yanliang567 <[email protected]>
  • Loading branch information
@yanliang567
yanliang567 authored Oct 24, 2024
1 parent d8db3e8 commit 0d376f1
Show file tree
Hide file tree
Showing 4 changed files with 306 additions and 186 deletions.
4 changes: 3 additions & 1 deletion tests/python_client/check/func_check.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -465,8 +465,10 @@ def check_query_iterator(query_res, func_name, check_items):
log.info("search iteration finished, close")
query_iterator.close()
break
pk_name = ct.default_int64_field_name if res[0].get(ct.default_int64_field_name, None) is not None \
else ct.default_string_field_name
for i in range(len(res)):
pk_list.append(res[i][ct.default_int64_field_name])
pk_list.append(res[i][pk_name])
if check_items.get("limit", None):
assert len(res) <= check_items["limit"]
assert len(pk_list) == len(set(pk_list))
Expand Down
4 changes: 2 additions & 2 deletions tests/python_client/requirements.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,8 +27,8 @@ pytest-parallel
pytest-random-order

# pymilvus
pymilvus==2.5.0rc101
pymilvus[bulk_writer]==2.5.0rc101
pymilvus==2.5.0rc103
pymilvus[bulk_writer]==2.5.0rc103

# for customize config test
python-benedict==0.24.3
Expand Down
183 changes: 0 additions & 183 deletions tests/python_client/testcases/test_query.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4027,189 +4027,6 @@ def test_count_query_search_after_release_partition_load(self):
check_items={ct.err_code: 1, ct.err_msg: 'not loaded'})


class TestQueryIterator(TestcaseBase):
"""
******************************************************************
The following cases are used to test query iterator
******************************************************************
"""

@pytest.mark.tags(CaseLabel.L0)
def test_query_iterator_normal(self):
"""
target: test query iterator normal
method: 1. query iterator
2. check the result, expect pk
expected: query successfully
"""
# 1. initialize with data
batch_size = 100
collection_w = self.init_collection_general(prefix, True, is_index=False)[0]
collection_w.create_index(ct.default_float_vec_field_name, {"metric_type": "L2"})
collection_w.load()
# 2. search iterator
expr = "int64 >= 0"
collection_w.query_iterator(batch_size, expr=expr,
check_task=CheckTasks.check_query_iterator,
check_items={"count": ct.default_nb,
"batch_size": batch_size})

@pytest.mark.tags(CaseLabel.L1)
def test_query_iterator_using_default_batch_size(self):
"""
target: test query iterator normal
method: 1. query iterator
2. check the result, expect pk
expected: query successfully
"""
# 1. initialize with data
collection_w = self.init_collection_general(prefix, True)[0]
# 2. query iterator
collection_w.query_iterator(check_task=CheckTasks.check_query_iterator,
check_items={"count": ct.default_nb,
"batch_size": ct.default_batch_size})

@pytest.mark.tags(CaseLabel.L2)
@pytest.mark.parametrize("offset", [500, 1000, 1777])
def test_query_iterator_with_offset(self, offset):
"""
target: test query iterator normal
method: 1. query iterator
2. check the result, expect pk
expected: query successfully
"""
# 1. initialize with data
batch_size = 300
collection_w = self.init_collection_general(prefix, True, is_index=False)[0]
collection_w.create_index(ct.default_float_vec_field_name, {"metric_type": "L2"})
collection_w.load()
# 2. search iterator
expr = "int64 >= 0"
collection_w.query_iterator(batch_size, expr=expr, offset=offset,
check_task=CheckTasks.check_query_iterator,
check_items={"count": ct.default_nb - offset,
"batch_size": batch_size})

@pytest.mark.tags(CaseLabel.L2)
@pytest.mark.parametrize("vector_data_type", ["FLOAT_VECTOR", "FLOAT16_VECTOR", "BFLOAT16_VECTOR"])
def test_query_iterator_output_different_vector_type(self, vector_data_type):
"""
target: test query iterator with output fields
method: 1. query iterator output different vector type
2. check the result, expect pk
expected: query successfully
"""
# 1. initialize with data
batch_size = 400
collection_w = self.init_collection_general(prefix, True,
vector_data_type=vector_data_type)[0]
# 2. query iterator
expr = "int64 >= 0"
collection_w.query_iterator(batch_size, expr=expr,
output_fields=[ct.default_float_vec_field_name],
check_task=CheckTasks.check_query_iterator,
check_items={"count": ct.default_nb,
"batch_size": batch_size})

@pytest.mark.tags(CaseLabel.L1)
@pytest.mark.parametrize("batch_size", [10, 100, 777, 2000])
def test_query_iterator_with_different_batch_size(self, batch_size):
"""
target: test query iterator normal
method: 1. query iterator
2. check the result, expect pk
expected: query successfully
"""
# 1. initialize with data
offset = 500
collection_w = self.init_collection_general(prefix, True, is_index=False)[0]
collection_w.create_index(ct.default_float_vec_field_name, {"metric_type": "L2"})
collection_w.load()
# 2. search iterator
expr = "int64 >= 0"
collection_w.query_iterator(batch_size=batch_size, expr=expr, offset=offset,
check_task=CheckTasks.check_query_iterator,
check_items={"count": ct.default_nb - offset,
"batch_size": batch_size})

@pytest.mark.tags(CaseLabel.L2)
@pytest.mark.parametrize("offset", [0, 10, 100, 1000])
@pytest.mark.parametrize("limit", [0, 100, 1500, 2000, 10000])
def test_query_iterator_with_different_limit(self, limit, offset):
"""
target: test query iterator normal
method: 1. query iterator
2. check the result, expect pk
expected: query successfully
"""
# 1. initialize with data
collection_w = self.init_collection_general(prefix, True)[0]
# 2. query iterator
Count = limit if limit + offset <= ct.default_nb else ct.default_nb - offset
collection_w.query_iterator(limit=limit, expr="", offset=offset,
check_task=CheckTasks.check_query_iterator,
check_items={"count": max(Count, 0),
"batch_size": ct.default_batch_size})

@pytest.mark.tags(CaseLabel.L2)
def test_query_iterator_invalid_batch_size(self):
"""
target: test query iterator invalid limit and offset
method: query iterator using invalid limit and offset
expected: raise exception
"""
# 1. initialize with data
nb = 17000 # set nb > 16384
collection_w = self.init_collection_general(prefix, True, nb=nb)[0]
# 2. search iterator
expr = "int64 >= 0"
error = {"err_code": 1, "err_msg": "batch size cannot be less than zero"}
collection_w.query_iterator(batch_size=-1, expr=expr, check_task=CheckTasks.err_res, check_items=error)

@pytest.mark.tags(CaseLabel.L0)
@pytest.mark.parametrize("batch_size", [100, 500])
@pytest.mark.parametrize("auto_id", [True, False])
def test_query_iterator_empty_expr(self, auto_id, batch_size):
"""
target: test query iterator with empty expression
method: query iterator empty expression with a limit
expected: return topK results by order
"""
# 1. initialize with data
collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, auto_id=auto_id)[0:4]

# 2. query with limit
collection_w.query_iterator(batch_size=batch_size,
check_task=CheckTasks.check_query_iterator,
check_items={"batch_size": batch_size,
"count": ct.default_nb,
"exp_ids": insert_ids})

@pytest.mark.tags(CaseLabel.L2)
@pytest.mark.parametrize("offset", [100, 1000])
@pytest.mark.parametrize("batch_size", [500, 1000])
def test_query_iterator_expr_empty_with_random_pk_pagination(self, batch_size, offset):
"""
target: test query iterator with empty expression
method: create a collection using random pk, query empty expression with a limit
expected: return topK results by order
"""
# 1. initialize with data
collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, random_primary_key=True)[0:4]

# 2. query with empty expr and check the result
exp_ids = sorted(insert_ids)
collection_w.query_iterator(batch_size, output_fields=[ct.default_string_field_name],
check_task=CheckTasks.check_query_iterator,
check_items={"batch_size": batch_size, "count": ct.default_nb, "exp_ids": exp_ids})

# 3. query with pagination
exp_ids = sorted(insert_ids)[offset:]
collection_w.query_iterator(batch_size, offset=offset, output_fields=[ct.default_string_field_name],
check_task=CheckTasks.check_query_iterator,
check_items={"batch_size": batch_size, "count": ct.default_nb - offset, "exp_ids": exp_ids})


class TestQueryNoneAndDefaultData(TestcaseBase):
"""
test Query interface with none and default data
Expand Down
Loading

0 comments on commit 0d376f1

Please sign in to comment.