wehub-resource-sync
1435 lines
57 KiB
Python
1435 lines
57 KiB
Python
# Copyright 2025-present the zvec project
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
|
|
from distance_helper import *
|
|
from doc_helper import *
|
|
from fixture_helper import *
|
|
from params_helper import *
|
|
from zvec import StatusCode
|
|
from zvec.extension import QwenReRanker, RrfReRanker, WeightedReRanker
|
|
from zvec.model import Collection, Doc
|
|
from zvec.model.param import (
|
|
CollectionOption,
|
|
FlatIndexParam,
|
|
HnswIndexParam,
|
|
HnswQueryParam,
|
|
InvertIndexParam,
|
|
IVFIndexParam,
|
|
IVFQueryParam,
|
|
)
|
|
from zvec.model.schema import FieldSchema, VectorSchema
|
|
from zvec.typing import DataType, MetricType, QuantizeType, StatusCode
|
|
|
|
|
|
# ==================== helper ====================
|
|
def batchdoc_and_check(
|
|
collection: Collection, multiple_docs, doc_num, operator="insert"
|
|
):
|
|
if operator == "insert":
|
|
result = collection.insert(multiple_docs)
|
|
elif operator == "upsert":
|
|
result = collection.upsert(multiple_docs)
|
|
|
|
elif operator == "update":
|
|
result = collection.update(multiple_docs)
|
|
else:
|
|
logging.error("operator value is error!")
|
|
|
|
assert len(result) == len(multiple_docs)
|
|
for item in result:
|
|
assert item.ok(), (
|
|
f"result={result},Insert operation failed with code {item.code()}"
|
|
)
|
|
|
|
stats = collection.stats
|
|
assert stats is not None, "Collection stats should not be None"
|
|
assert stats.doc_count == len(multiple_docs), (
|
|
f"Document count should be {len(multiple_docs)} after insert, but got {stats.doc_count}"
|
|
)
|
|
|
|
doc_ids = [doc.id for doc in multiple_docs]
|
|
fetched_docs = collection.fetch(doc_ids)
|
|
assert len(fetched_docs) == len(multiple_docs), (
|
|
f"fetched_docs={fetched_docs},Expected {len(multiple_docs)} fetched documents, but got {len(fetched_docs)}"
|
|
)
|
|
|
|
for original_doc in multiple_docs:
|
|
assert original_doc.id in fetched_docs, (
|
|
f"Expected document ID {original_doc.id} in fetched documents"
|
|
)
|
|
fetched_doc = fetched_docs[original_doc.id]
|
|
|
|
assert is_doc_equal(fetched_doc, original_doc, collection.schema)
|
|
|
|
assert hasattr(fetched_doc, "score"), "Document should have a score attribute"
|
|
assert fetched_doc.score == 0.0, (
|
|
"Fetch operation should return default score of 0.0"
|
|
)
|
|
|
|
first_doc = multiple_docs[doc_num - 1]
|
|
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
|
|
query_result = collection.query(
|
|
Query(field_name=v, vector=first_doc.vectors[v]),
|
|
topk=1024,
|
|
include_vector=True,
|
|
)
|
|
assert len(query_result) > 0, (
|
|
f"Expected at least 1 query result, but got {len(query_result)}"
|
|
)
|
|
|
|
found_doc = None
|
|
|
|
for doc in query_result:
|
|
if doc.id == first_doc.id:
|
|
found_doc = doc
|
|
break
|
|
assert found_doc is not None, (
|
|
f"Inserted document {first_doc.id} not found in query results"
|
|
)
|
|
|
|
assert is_doc_equal(found_doc, first_doc, collection.schema)
|
|
assert hasattr(found_doc, "score")
|
|
assert isinstance(found_doc.score, (int, float))
|
|
|
|
|
|
def batchdoc_and_check_ivf(
|
|
collection: Collection, multiple_docs, doc_num, operator="insert"
|
|
):
|
|
if operator == "insert":
|
|
result = collection.insert(multiple_docs)
|
|
elif operator == "upsert":
|
|
result = collection.upsert(multiple_docs)
|
|
|
|
elif operator == "update":
|
|
result = collection.update(multiple_docs)
|
|
else:
|
|
logging.error("operator value is error!")
|
|
|
|
assert len(result) == len(multiple_docs)
|
|
for item in result:
|
|
assert item.ok(), (
|
|
f"result={result},Insert operation failed with code {item.code()}"
|
|
)
|
|
|
|
stats = collection.stats
|
|
assert stats is not None, "Collection stats should not be None"
|
|
assert stats.doc_count == len(multiple_docs), (
|
|
f"Document count should be {len(multiple_docs)} after insert, but got {stats.doc_count}"
|
|
)
|
|
|
|
doc_ids = [doc.id for doc in multiple_docs]
|
|
fetched_docs = collection.fetch(doc_ids)
|
|
assert len(fetched_docs) == len(multiple_docs), (
|
|
f"fetched_docs={fetched_docs},Expected {len(multiple_docs)} fetched documents, but got {len(fetched_docs)}"
|
|
)
|
|
|
|
for original_doc in multiple_docs:
|
|
assert original_doc.id in fetched_docs, (
|
|
f"Expected document ID {original_doc.id} in fetched documents"
|
|
)
|
|
fetched_doc = fetched_docs[original_doc.id]
|
|
|
|
assert is_doc_equal(fetched_doc, original_doc, collection.schema)
|
|
|
|
assert hasattr(fetched_doc, "score"), "Document should have a score attribute"
|
|
assert fetched_doc.score == 0.0, (
|
|
"Fetch operation should return default score of 0.0"
|
|
)
|
|
|
|
first_doc = multiple_docs[doc_num - 1]
|
|
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
|
|
if v in ["vector_fp16_field", "vector_fp32_field"]:
|
|
query_result = collection.query(
|
|
Query(field_name=v, vector=first_doc.vectors[v]),
|
|
topk=1024,
|
|
include_vector=True,
|
|
)
|
|
assert len(query_result) > 0, (
|
|
f"Expected at least 1 query result, but got {len(query_result)}"
|
|
)
|
|
|
|
found_doc = None
|
|
|
|
for doc in query_result:
|
|
if doc.id == first_doc.id:
|
|
found_doc = doc
|
|
break
|
|
assert found_doc is not None, (
|
|
f"Inserted document {first_doc.id} not found in query results"
|
|
)
|
|
|
|
assert is_doc_equal(found_doc, first_doc, collection.schema)
|
|
assert hasattr(found_doc, "score")
|
|
assert isinstance(found_doc.score, (int, float))
|
|
|
|
|
|
def single_querydoc_check(
|
|
multiple_docs,
|
|
query_result,
|
|
full_collection: Collection,
|
|
is_by_vector=0,
|
|
query_vector=None,
|
|
data_type=None,
|
|
vector_name=None,
|
|
metric_type=MetricType.IP,
|
|
id_include_vector: bool = False,
|
|
is_output_fields=0,
|
|
):
|
|
for original_doc in multiple_docs:
|
|
for doc in query_result:
|
|
if doc.id == original_doc.id:
|
|
found_doc = doc
|
|
if is_output_fields == 0:
|
|
assert is_doc_equal(
|
|
found_doc,
|
|
original_doc,
|
|
full_collection.schema,
|
|
True,
|
|
id_include_vector,
|
|
)
|
|
assert hasattr(found_doc, "score")
|
|
# assert found_doc.score >= 0.0
|
|
if not id_include_vector:
|
|
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
|
|
assert found_doc.vector(v) == {}
|
|
else:
|
|
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
|
|
assert found_doc.vector(v) != {}
|
|
if is_by_vector:
|
|
prev_score = float("inf")
|
|
for i, doc in enumerate(query_result):
|
|
doc_vector = full_collection.fetch(doc.id)[doc.id].vector(
|
|
vector_name
|
|
)
|
|
expected_score = distance(
|
|
query_vector, doc_vector, metric_type, data_type, k
|
|
)
|
|
if (
|
|
full_collection.schema.vector(vector_name).data_type
|
|
!= DataType.VECTOR_FP16
|
|
):
|
|
assert abs(doc.score - expected_score) < 0.001, (
|
|
f"{data_type} {vector_name} :Expected score {expected_score:.6f}, but got {doc.score:.6f} for document {doc.id}"
|
|
)
|
|
assert doc.score <= prev_score, (
|
|
f"{data_type} {vector_name} :Scores should be in descending order. Current: {doc.score}, Previous: {prev_score}"
|
|
)
|
|
prev_score = doc.score
|
|
|
|
|
|
def multi_querydoc_check(multiple_docs, query_result, full_collection):
|
|
for original_doc in multiple_docs:
|
|
for doc in query_result:
|
|
if doc.id == original_doc.id:
|
|
found_doc = doc
|
|
assert is_doc_equal(
|
|
found_doc, original_doc, full_collection.schema, False, False
|
|
)
|
|
assert hasattr(found_doc, "score"), (
|
|
"Document should have a score attribute"
|
|
)
|
|
assert found_doc.score >= 0.0, (
|
|
"Fetch operation should return default score of 0.0"
|
|
)
|
|
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
|
|
assert found_doc.vector(v) == {}
|
|
|
|
|
|
# ==================== Tests ====================
|
|
class TestCollectionFetch:
|
|
def test_fetch_non_existing(self, full_collection: Collection):
|
|
result = full_collection.fetch(ids=["non_existing_id1", "non_existing_id2"])
|
|
assert len(result) == 0
|
|
|
|
@pytest.mark.parametrize("doc_num", [3])
|
|
def test_fetch_partial_non_existing(self, full_collection: Collection, doc_num):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection.schema) for i in range(doc_num)
|
|
]
|
|
batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert")
|
|
|
|
fetch_id_list = [doc.id for doc in multiple_docs]
|
|
fetch_id_list.append("non_existing_id")
|
|
result = full_collection.fetch(ids=fetch_id_list)
|
|
|
|
assert len(result) == doc_num
|
|
assert "non_existing_id" not in result.keys()
|
|
|
|
def test_fetch_empty_ids(self, full_collection: Collection):
|
|
result = full_collection.fetch(ids=[])
|
|
assert len(result) == 0, (
|
|
f"Expected 0 results for empty ID list, but got {len(result)}"
|
|
)
|
|
|
|
@pytest.mark.parametrize("doc_num", [3])
|
|
def test_fetch_with_output_fields(self, full_collection: Collection, doc_num):
|
|
"""Test that fetch respects output_fields parameter."""
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection.schema) for i in range(doc_num)
|
|
]
|
|
result = full_collection.insert(multiple_docs)
|
|
for item in result:
|
|
assert item.ok(), f"Insert failed: {item.code()}"
|
|
|
|
doc_id = multiple_docs[0].id
|
|
|
|
# Case 1: output_fields=None -> all scalar fields returned
|
|
fetched_all = full_collection.fetch(ids=[doc_id], output_fields=None)
|
|
assert doc_id in fetched_all
|
|
doc_all = fetched_all[doc_id]
|
|
assert doc_all is not None
|
|
assert doc_all.has_field("int32_field"), (
|
|
"int32_field should be present when output_fields=None"
|
|
)
|
|
assert doc_all.has_field("string_field"), (
|
|
"string_field should be present when output_fields=None"
|
|
)
|
|
|
|
# Case 2: output_fields=["int32_field"] -> only int32_field returned
|
|
fetched_partial = full_collection.fetch(
|
|
ids=[doc_id], output_fields=["int32_field"]
|
|
)
|
|
assert doc_id in fetched_partial
|
|
doc_partial = fetched_partial[doc_id]
|
|
assert doc_partial is not None
|
|
assert doc_partial.has_field("int32_field"), "int32_field should be present"
|
|
assert not doc_partial.has_field("string_field"), (
|
|
'string_field should not be present when output_fields=["int32_field"]'
|
|
)
|
|
assert not doc_partial.has_field("float_field"), (
|
|
'float_field should not be present when output_fields=["int32_field"]'
|
|
)
|
|
|
|
# Case 3: output_fields=[] (empty) -> no scalar fields returned
|
|
fetched_empty = full_collection.fetch(ids=[doc_id], output_fields=[])
|
|
assert doc_id in fetched_empty
|
|
doc_empty = fetched_empty[doc_id]
|
|
assert doc_empty is not None
|
|
assert doc_empty.id == doc_id, "pk should still be set"
|
|
assert not doc_empty.has_field("int32_field"), (
|
|
"int32_field should not be present when output_fields=[]"
|
|
)
|
|
assert not doc_empty.has_field("string_field"), (
|
|
"string_field should not be present when output_fields=[]"
|
|
)
|
|
|
|
# Case 4: multiple output_fields
|
|
fetched_multi = full_collection.fetch(
|
|
ids=[doc_id], output_fields=["int32_field", "float_field"]
|
|
)
|
|
assert doc_id in fetched_multi
|
|
doc_multi = fetched_multi[doc_id]
|
|
assert doc_multi is not None
|
|
assert doc_multi.has_field("int32_field")
|
|
assert doc_multi.has_field("float_field")
|
|
assert not doc_multi.has_field("string_field")
|
|
|
|
@pytest.mark.parametrize("doc_num", [3])
|
|
def test_fetch_with_include_vector(self, full_collection: Collection, doc_num):
|
|
"""Test that fetch respects include_vector parameter."""
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection.schema) for i in range(doc_num)
|
|
]
|
|
result = full_collection.insert(multiple_docs)
|
|
for item in result:
|
|
assert item.ok(), f"Insert failed: {item.code()}"
|
|
|
|
doc_id = multiple_docs[0].id
|
|
|
|
# Case 1: include_vector=True (default) -> vector data returned
|
|
fetched_with_vec = full_collection.fetch(ids=[doc_id])
|
|
assert doc_id in fetched_with_vec
|
|
doc_with_vec = fetched_with_vec[doc_id]
|
|
assert doc_with_vec is not None
|
|
assert doc_with_vec.has_field("int32_field"), (
|
|
"scalar fields should still be present"
|
|
)
|
|
assert doc_with_vec.vector("vector_fp32_field"), (
|
|
"vector should be present when include_vector=True (default)"
|
|
)
|
|
|
|
# Case 2: include_vector=False -> no vector data returned
|
|
fetched_no_vec = full_collection.fetch(ids=[doc_id], include_vector=False)
|
|
assert doc_id in fetched_no_vec
|
|
doc_no_vec = fetched_no_vec[doc_id]
|
|
assert doc_no_vec is not None
|
|
assert doc_no_vec.has_field("int32_field"), (
|
|
"scalar fields should still be present"
|
|
)
|
|
assert not doc_no_vec.vector("vector_fp32_field"), (
|
|
"vector should not be present when include_vector=False"
|
|
)
|
|
|
|
# Case 3: include_vector=False with output_fields
|
|
fetched_combo = full_collection.fetch(
|
|
ids=[doc_id], output_fields=["int32_field"], include_vector=False
|
|
)
|
|
assert doc_id in fetched_combo
|
|
doc_combo = fetched_combo[doc_id]
|
|
assert doc_combo is not None
|
|
assert doc_combo.has_field("int32_field")
|
|
assert not doc_combo.has_field("string_field")
|
|
assert not doc_combo.vector("vector_fp32_field"), (
|
|
"vector should not be present when include_vector=False"
|
|
)
|
|
|
|
|
|
class TestCollectionQuery:
|
|
@pytest.mark.parametrize("doc_num", [5])
|
|
def test_query_with_no_condition(self, full_collection: Collection, doc_num):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection.schema) for i in range(doc_num)
|
|
]
|
|
batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert")
|
|
query_result = full_collection.query()
|
|
assert len(query_result) == doc_num
|
|
single_querydoc_check(multiple_docs, query_result, full_collection)
|
|
|
|
@pytest.mark.parametrize("doc_num", [10])
|
|
def test_query_with_filter_empty(self, full_collection: Collection, doc_num):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection.schema) for i in range(doc_num)
|
|
]
|
|
batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert")
|
|
result1 = full_collection.query(filter="")
|
|
assert len(result1) == doc_num
|
|
single_querydoc_check(multiple_docs, result1, full_collection)
|
|
result2 = full_collection.query(filter=None)
|
|
assert len(result2) == doc_num
|
|
single_querydoc_check(multiple_docs, result2, full_collection)
|
|
ids1 = set(doc.id for doc in result1)
|
|
ids2 = set(doc.id for doc in result2)
|
|
assert ids1 == ids2
|
|
|
|
@pytest.mark.parametrize("field_name", ["int32_field"])
|
|
@pytest.mark.parametrize("doc_num", [10])
|
|
def test_query_with_filter_single_condition(
|
|
self, full_collection: Collection, doc_num, field_name
|
|
):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection.schema) for i in range(doc_num)
|
|
]
|
|
batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert")
|
|
filter = field_name + " > 5"
|
|
query_result = full_collection.query(filter=filter)
|
|
assert len(query_result) == doc_num - 6
|
|
|
|
returned_doc_ids = set()
|
|
for doc in query_result:
|
|
returned_doc_ids.add(doc.id)
|
|
|
|
expected_doc_ids = set(str(i) for i in range(6, doc_num))
|
|
|
|
for doc in query_result:
|
|
assert doc.id in expected_doc_ids
|
|
assert int(doc.field(field_name)) > 5
|
|
|
|
single_querydoc_check(multiple_docs, query_result, full_collection)
|
|
|
|
@pytest.mark.parametrize("field_name", ["int32_field"])
|
|
@pytest.mark.parametrize(
|
|
"filter",
|
|
[
|
|
"int32_field > 3 and int32_field < 9",
|
|
"int32_field >= 5 and int32_field <= 7",
|
|
],
|
|
)
|
|
@pytest.mark.parametrize("doc_num", [10])
|
|
def test_query_with_filter_and(
|
|
self, full_collection: Collection, doc_num, field_name, filter
|
|
):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection.schema) for i in range(doc_num)
|
|
]
|
|
batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert")
|
|
filter = field_name + " > 3 and " + field_name + " < 9"
|
|
query_result = full_collection.query(filter=filter)
|
|
if filter == "int32_field > 3 and int32_field < 9":
|
|
assert len(query_result) == doc_num - 4 - 1
|
|
expected_doc_ids = set(str(i) for i in range(4, 9))
|
|
|
|
for doc in query_result:
|
|
assert doc.id in expected_doc_ids
|
|
field_value = int(doc.field(field_name))
|
|
assert field_value > 3 and field_value < 9
|
|
else:
|
|
assert len(query_result) == 3
|
|
expected_doc_ids = set(str(i) for i in range(5, 8))
|
|
|
|
for doc in query_result:
|
|
assert doc.id in expected_doc_ids
|
|
field_value = int(doc.field(field_name))
|
|
assert field_value >= 5 and field_value <= 7
|
|
|
|
single_querydoc_check(multiple_docs, query_result, full_collection)
|
|
|
|
@pytest.mark.parametrize("field_name", ["int32_field"])
|
|
@pytest.mark.parametrize(
|
|
"filter",
|
|
[
|
|
"int32_field < 3 or int32_field > 8",
|
|
"int32_field = 3 or int32_field = 7",
|
|
"int32_field <= 3 or int32_field >= 8",
|
|
],
|
|
)
|
|
@pytest.mark.parametrize("doc_num", [10])
|
|
def test_query_with_filter_or(
|
|
self, full_collection: Collection, doc_num, field_name, filter
|
|
):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection.schema) for i in range(doc_num)
|
|
]
|
|
batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert")
|
|
query_result = full_collection.query(filter=filter)
|
|
if filter == "int32_field < 3 or int32_field > 8":
|
|
assert len(query_result) == 4
|
|
expected_doc_ids = set([str(0), str(1), str(2), str(9)])
|
|
for doc in query_result:
|
|
assert doc.id in expected_doc_ids
|
|
field_value = int(doc.field(field_name))
|
|
assert field_value < 3 or field_value > 8
|
|
elif filter == "int32_field = 3 or int32_field = 7":
|
|
assert len(query_result) == 2
|
|
expected_doc_ids = set([str(3), str(7)])
|
|
for doc in query_result:
|
|
assert doc.id in expected_doc_ids
|
|
field_value = int(doc.field(field_name))
|
|
assert field_value == 3 or field_value == 7
|
|
else:
|
|
assert len(query_result) == 6
|
|
expected_doc_ids = set([str(0), str(1), str(2), str(3), str(8), str(9)])
|
|
for doc in query_result:
|
|
assert doc.id in expected_doc_ids
|
|
field_value = int(doc.field(field_name))
|
|
assert field_value <= 3 or field_value >= 8
|
|
|
|
single_querydoc_check(multiple_docs, query_result, full_collection)
|
|
|
|
@pytest.mark.parametrize("field_names", [("int32_field", "bool_field")])
|
|
@pytest.mark.parametrize(
|
|
"filter",
|
|
[
|
|
"(int32_field < 3 or int32_field > 8) and bool_field = false",
|
|
"(int32_field > 2 and int32_field < 5) or (int32_field > 7 and bool_field = true)",
|
|
],
|
|
)
|
|
@pytest.mark.parametrize("doc_num", [10])
|
|
def test_query_with_filter_parentheses(
|
|
self, full_collection: Collection, doc_num, field_names, filter
|
|
):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection.schema) for i in range(doc_num)
|
|
]
|
|
batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert")
|
|
query_result = full_collection.query(filter=filter)
|
|
if filter == "(int32_field < 3 or int32_field > 8) and bool_field = false":
|
|
assert len(query_result) == 2
|
|
expected_doc_ids = set([str(1), str(9)])
|
|
for doc in query_result:
|
|
assert doc.id in expected_doc_ids
|
|
assert (
|
|
int(doc.field(field_names[0])) < 3
|
|
or int(doc.field(field_names[0])) > 8
|
|
) and doc.field(field_names[1]) == False
|
|
else:
|
|
assert len(query_result) == 3
|
|
expected_doc_ids = set([str(3), str(4), str(8)])
|
|
for doc in query_result:
|
|
assert doc.id in expected_doc_ids
|
|
assert (
|
|
(
|
|
int(doc.field(field_names[0])) > 2
|
|
and int(doc.field(field_names[0])) < 5
|
|
)
|
|
or (doc.field(field_names[0])) > 7
|
|
and doc.field(field_names[1]) == True
|
|
)
|
|
single_querydoc_check(multiple_docs, query_result, full_collection)
|
|
|
|
@pytest.mark.parametrize(
|
|
"filter",
|
|
[
|
|
"int32_field >",
|
|
"int32_field = 'string'",
|
|
"nonexistent_field = 5",
|
|
"int32_field > 5 and",
|
|
"int32_field > > 5",
|
|
],
|
|
)
|
|
@pytest.mark.parametrize("doc_num", [10])
|
|
def test_query_filter_invalid(self, full_collection: Collection, doc_num, filter):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection.schema) for i in range(doc_num)
|
|
]
|
|
batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert")
|
|
with pytest.raises(Exception) as exc_info:
|
|
full_collection.query(filter=filter)
|
|
if filter in ["int32_field = 'string'", "nonexistent_field = 5"]:
|
|
assert "Analyze SQL info failed" in str(exc_info.value)
|
|
else:
|
|
assert "Invalid filter" in str(exc_info.value)
|
|
|
|
@pytest.mark.parametrize("field_name", ["int32_field"])
|
|
@pytest.mark.parametrize("topk_value", [1, 5, 10, 50, 100, 500, 1000, 1024])
|
|
def test_query_with_filter_topk_valid(
|
|
self, full_collection: Collection, topk_value: int, field_name
|
|
):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection.schema) for i in range(topk_value)
|
|
]
|
|
batchdoc_and_check(
|
|
full_collection, multiple_docs, topk_value, operator="insert"
|
|
)
|
|
filter = (
|
|
field_name + f" >={topk_value - 1} and " + field_name + f" <={topk_value}"
|
|
)
|
|
print("filter:\n")
|
|
print(filter)
|
|
query_result = full_collection.query(filter=filter, topk=topk_value)
|
|
assert len(query_result) == 1
|
|
expected_doc_ids = [str(topk_value - 1)]
|
|
|
|
for doc in query_result:
|
|
assert doc.id in expected_doc_ids
|
|
field_value = int(doc.field(field_name))
|
|
assert field_value >= topk_value - 1 and field_value <= topk_value
|
|
single_querydoc_check(multiple_docs, query_result, full_collection)
|
|
|
|
@pytest.mark.parametrize("field_name", ["int32_field"])
|
|
@pytest.mark.parametrize("topk_value", [1, 5, 10, 50, 100, 500, 1000, 1024])
|
|
def test_query_without_filter_topk_valid(
|
|
self, full_collection: Collection, topk_value: int, field_name
|
|
):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection.schema) for i in range(topk_value)
|
|
]
|
|
batchdoc_and_check(
|
|
full_collection, multiple_docs, topk_value, operator="insert"
|
|
)
|
|
|
|
query_result = full_collection.query(topk=topk_value)
|
|
assert len(query_result) == topk_value
|
|
single_querydoc_check(multiple_docs, query_result, full_collection)
|
|
|
|
@pytest.mark.parametrize("doc_num", [10])
|
|
def test_query_with_include_vector(self, full_collection: Collection, doc_num):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection.schema) for i in range(doc_num)
|
|
]
|
|
batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert")
|
|
|
|
query_result = full_collection.query(include_vector=True)
|
|
assert len(query_result) > 0
|
|
single_querydoc_check(
|
|
multiple_docs, query_result, full_collection, id_include_vector=1
|
|
)
|
|
|
|
@pytest.mark.parametrize("output_fields", [["int32_field", "int64_field"]])
|
|
@pytest.mark.parametrize("doc_num", [10])
|
|
def test_query_with_output_fields(
|
|
self, full_collection: Collection, doc_num, output_fields
|
|
):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection.schema) for i in range(doc_num)
|
|
]
|
|
batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert")
|
|
query_result = full_collection.query(output_fields=output_fields)
|
|
assert len(query_result) > 0
|
|
for doc in query_result:
|
|
field_names = doc.field_names()
|
|
assert field_names == output_fields
|
|
|
|
@pytest.mark.parametrize(
|
|
"filter",
|
|
[
|
|
"int32_field >= 10 and int32_field <= 20",
|
|
"int32_field = 3 and int32_field = 8",
|
|
],
|
|
)
|
|
@pytest.mark.parametrize("doc_num", [10])
|
|
def test_query_empty_result(self, full_collection: Collection, doc_num, filter):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection.schema) for i in range(doc_num)
|
|
]
|
|
batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert")
|
|
result = full_collection.query(filter=filter)
|
|
assert len(result) == 0
|
|
|
|
@pytest.mark.parametrize(
|
|
"full_schema_new",
|
|
[(True, True, HnswIndexParam()), (False, True, FlatIndexParam())],
|
|
indirect=True,
|
|
)
|
|
@pytest.mark.parametrize("doc_num", [10])
|
|
def test_query_by_id(
|
|
self, full_collection_new: Collection, doc_num, full_schema_new
|
|
):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection_new.schema) for i in range(doc_num)
|
|
]
|
|
batchdoc_and_check(
|
|
full_collection_new, multiple_docs, doc_num, operator="insert"
|
|
)
|
|
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
|
|
query_result = full_collection_new.query(Query(field_name=v, id="1"))
|
|
assert len(query_result) > 0
|
|
query_doc = full_collection_new.fetch(ids=["1"])
|
|
query_vector = query_doc["1"].vector(v)
|
|
single_querydoc_check(
|
|
multiple_docs,
|
|
query_result,
|
|
full_collection_new,
|
|
is_by_vector=1,
|
|
query_vector=query_vector,
|
|
data_type=k,
|
|
vector_name=v,
|
|
)
|
|
|
|
@pytest.mark.parametrize("doc_num", [10])
|
|
def test_query_by_id_ivf(self, full_collection_ivf: Collection, doc_num):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection_ivf.schema) for i in range(doc_num)
|
|
]
|
|
batchdoc_and_check_ivf(
|
|
full_collection_ivf, multiple_docs, doc_num, operator="insert"
|
|
)
|
|
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
|
|
if v in ["vector_fp16_field", "vector_fp32_field"]:
|
|
query_result = full_collection_ivf.query(Query(field_name=v, id="1"))
|
|
assert len(query_result) > 0
|
|
query_doc = full_collection_ivf.fetch(ids=["1"])
|
|
query_vector = query_doc["1"].vector(v)
|
|
single_querydoc_check(
|
|
multiple_docs,
|
|
query_result,
|
|
full_collection_ivf,
|
|
is_by_vector=1,
|
|
query_vector=query_vector,
|
|
data_type=k,
|
|
vector_name=v,
|
|
)
|
|
|
|
@pytest.mark.parametrize(
|
|
"full_schema_new",
|
|
[(True, True, HnswIndexParam()), (False, True, FlatIndexParam())],
|
|
indirect=True,
|
|
)
|
|
@pytest.mark.parametrize("doc_num", [10])
|
|
@pytest.mark.parametrize("topk", [None, 1024])
|
|
@pytest.mark.parametrize("filter", [None, "int32_field >= 3 and int32_field <= 7"])
|
|
def test_query_by_vector(
|
|
self, full_collection_new: Collection, doc_num, full_schema_new, topk, filter
|
|
):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection_new.schema) for i in range(doc_num)
|
|
]
|
|
batchdoc_and_check(
|
|
full_collection_new, multiple_docs, doc_num, operator="insert"
|
|
)
|
|
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
|
|
doc_fields, doc_vectors = generate_vectordict_random(
|
|
full_collection_new.schema
|
|
)
|
|
query_vector = doc_vectors[v]
|
|
if topk and filter:
|
|
query_result = full_collection_new.query(
|
|
Query(field_name=v, vector=query_vector),
|
|
filter=filter,
|
|
topk=topk,
|
|
)
|
|
elif topk and not filter:
|
|
query_result = full_collection_new.query(
|
|
Query(field_name=v, vector=query_vector), topk=topk
|
|
)
|
|
elif not topk and filter:
|
|
query_result = full_collection_new.query(
|
|
Query(field_name=v, vector=query_vector),
|
|
filter=filter,
|
|
)
|
|
else:
|
|
query_result = full_collection_new.query(
|
|
Query(field_name=v, vector=query_vector)
|
|
)
|
|
assert len(query_result) > 0, (
|
|
f"Expected at least 1 query result, but got {len(query_result)}"
|
|
)
|
|
single_querydoc_check(
|
|
multiple_docs,
|
|
query_result,
|
|
full_collection_new,
|
|
is_by_vector=1,
|
|
query_vector=query_vector,
|
|
data_type=k,
|
|
vector_name=v,
|
|
)
|
|
|
|
@pytest.mark.parametrize("doc_num", [10])
|
|
def test_query_by_vector_ivf(self, full_collection_ivf: Collection, doc_num):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection_ivf.schema) for i in range(doc_num)
|
|
]
|
|
batchdoc_and_check_ivf(
|
|
full_collection_ivf, multiple_docs, doc_num, operator="insert"
|
|
)
|
|
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
|
|
if v in ["vector_fp16_field", "vector_fp32_field"]:
|
|
doc_fields, doc_vectors = generate_vectordict_random(
|
|
full_collection_ivf.schema
|
|
)
|
|
query_vector = doc_vectors[v]
|
|
query_result = full_collection_ivf.query(
|
|
Query(field_name=v, vector=query_vector),
|
|
topk=1024,
|
|
)
|
|
assert len(query_result) > 0, (
|
|
f"Expected at least 1 query result, but got {len(query_result)}"
|
|
)
|
|
single_querydoc_check(
|
|
multiple_docs,
|
|
query_result,
|
|
full_collection_ivf,
|
|
is_by_vector=1,
|
|
query_vector=query_vector,
|
|
data_type=k,
|
|
vector_name=v,
|
|
)
|
|
|
|
@pytest.mark.parametrize("doc_num", [10])
|
|
def test_query_multivector_rrf(self, full_collection: Collection, doc_num):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection.schema) for i in range(doc_num)
|
|
]
|
|
batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert")
|
|
|
|
doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema)
|
|
single_query_results = {}
|
|
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
|
|
single_query_results[v] = full_collection.query(
|
|
Query(field_name=v, vector=doc_vectors[v])
|
|
)
|
|
expected_rrf_scores = calculate_multi_vector_rrf_scores(single_query_results)
|
|
multi_query_vectors = []
|
|
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
|
|
multi_query_vectors.append(Query(field_name=v, vector=doc_vectors[v]))
|
|
|
|
rrf_reranker = RrfReRanker()
|
|
multi_query_result = full_collection.query(
|
|
multi_query_vectors,
|
|
topk=3,
|
|
reranker=rrf_reranker,
|
|
)
|
|
assert len(multi_query_result) > 0, (
|
|
f"Expected at least 1 result, but got {len(multi_query_result)}"
|
|
)
|
|
|
|
multi_querydoc_check(multiple_docs, multi_query_result, full_collection)
|
|
|
|
prev_score = float("inf")
|
|
for i, doc in enumerate(multi_query_result):
|
|
doc_id = doc.id
|
|
assert doc_id in expected_rrf_scores, (
|
|
f"Document {doc_id} should be in expected RRF scores"
|
|
)
|
|
expected_score = expected_rrf_scores[doc_id]
|
|
actual_score = doc.score
|
|
assert abs(actual_score - expected_score) < 1e-6, (
|
|
f"RRF score mismatch for document {doc_id}: expected {expected_score}, got {actual_score}"
|
|
)
|
|
assert doc.score <= prev_score, (
|
|
f"Scores should be in descending order. Current: {doc.score}, Previous: {prev_score}"
|
|
)
|
|
prev_score = doc.score
|
|
|
|
@pytest.mark.parametrize(
|
|
"weights",
|
|
[
|
|
{
|
|
"vector_fp32_field": 0.3,
|
|
"vector_fp16_field": 0.2,
|
|
"vector_int8_field": 0.3,
|
|
"sparse_vector_fp32_field": 0.1,
|
|
"sparse_vector_fp16_field": 0.1,
|
|
}
|
|
],
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"metric_type", [MetricType.L2, MetricType.IP, MetricType.COSINE]
|
|
)
|
|
@pytest.mark.parametrize("doc_num", [10])
|
|
def test_query_multivector_weighted(
|
|
self, full_collection: Collection, doc_num, weights, metric_type
|
|
):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection.schema) for i in range(doc_num)
|
|
]
|
|
batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert")
|
|
doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema)
|
|
|
|
# Weights are positional, aligned with the multi_query_vectors order
|
|
# (DEFAULT_VECTOR_FIELD_NAME insertion order). Metric normalization is
|
|
# automatic from each field's schema.
|
|
weights_list = [weights[v] for v in DEFAULT_VECTOR_FIELD_NAME.values()]
|
|
weighted_reranker = WeightedReRanker(weights_list)
|
|
|
|
single_query_results = {}
|
|
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
|
|
single_query_results[v] = full_collection.query(
|
|
Query(field_name=v, vector=doc_vectors[v])
|
|
)
|
|
expected_weighted_scores = calculate_multi_vector_weighted_scores(
|
|
single_query_results, weights, MetricType.IP
|
|
)
|
|
|
|
multi_query_vectors = []
|
|
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
|
|
multi_query_vectors.append(Query(field_name=v, vector=doc_vectors[v]))
|
|
|
|
multi_query_result = full_collection.query(
|
|
multi_query_vectors,
|
|
topk=3,
|
|
reranker=weighted_reranker,
|
|
)
|
|
assert len(multi_query_result) > 0, (
|
|
f"Expected at least 1 result, but got {len(multi_query_result)}"
|
|
)
|
|
|
|
multi_querydoc_check(multiple_docs, multi_query_result, full_collection)
|
|
|
|
prev_score = float("inf")
|
|
for i, doc in enumerate(multi_query_result):
|
|
doc_id = doc.id
|
|
assert doc_id in expected_weighted_scores, (
|
|
f"Document {doc_id} should be in expected scores"
|
|
)
|
|
expected_score = expected_weighted_scores[doc_id]
|
|
actual_score = doc.score
|
|
assert abs(actual_score - expected_score) < 1e-6, (
|
|
f"score mismatch for document {doc_id}: expected {expected_score}, got {actual_score}"
|
|
)
|
|
assert doc.score <= prev_score, (
|
|
f"Scores should be in descending order. Current: {doc.score}, Previous: {prev_score}"
|
|
)
|
|
prev_score = doc.score
|
|
|
|
@pytest.mark.parametrize("topk", [5])
|
|
@pytest.mark.parametrize("doc_num", [10])
|
|
@pytest.mark.parametrize("filter", ["int32_field >= 3 and int32_field <= 7"])
|
|
def test_query_consistency(
|
|
self, full_collection: Collection, filter, doc_num, topk
|
|
):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection.schema) for i in range(doc_num)
|
|
]
|
|
batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert")
|
|
results = []
|
|
for i in range(5):
|
|
query_result = full_collection.query(filter=filter, topk=topk)
|
|
single_querydoc_check(multiple_docs, query_result, full_collection)
|
|
|
|
results.append(query_result)
|
|
assert len(results) == 5
|
|
expected_count = len(results[0])
|
|
for i, result in enumerate(results):
|
|
assert len(result) == expected_count
|
|
|
|
expected_ids = set(doc.id for doc in results[0])
|
|
for i, result in enumerate(results):
|
|
result_ids = set(doc.id for doc in result)
|
|
assert result_ids == expected_ids
|
|
|
|
for i, result in enumerate(results):
|
|
result_ids = [doc.id for doc in result]
|
|
expected_sorted_ids = sorted(result_ids, key=lambda x: int(x))
|
|
assert result_ids == expected_sorted_ids
|
|
|
|
@pytest.mark.parametrize("ef", [0, 100, 1024, 2048])
|
|
@pytest.mark.parametrize("doc_num", [10])
|
|
@pytest.mark.parametrize("topk", [1024])
|
|
@pytest.mark.parametrize("filter", ["int32_field >= 3 and int32_field <= 7"])
|
|
@pytest.mark.parametrize(
|
|
"full_schema_new", [(True, True, HnswIndexParam())], indirect=True
|
|
)
|
|
def test_query_vector_with_HnswQueryParam_valid(
|
|
self,
|
|
full_collection_new: Collection,
|
|
doc_num,
|
|
full_schema_new,
|
|
topk,
|
|
filter,
|
|
ef,
|
|
):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection_new.schema) for i in range(doc_num)
|
|
]
|
|
batchdoc_and_check(
|
|
full_collection_new, multiple_docs, doc_num, operator="insert"
|
|
)
|
|
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
|
|
doc_fields, doc_vectors = generate_vectordict_random(
|
|
full_collection_new.schema
|
|
)
|
|
query_vector = doc_vectors[v]
|
|
query_result = full_collection_new.query(
|
|
Query(field_name=v, vector=query_vector, param=HnswQueryParam(ef=ef)),
|
|
filter=filter,
|
|
topk=topk,
|
|
)
|
|
assert len(query_result) > 0, (
|
|
f"Expected at least 1 query result, but got {len(query_result)}"
|
|
)
|
|
single_querydoc_check(
|
|
multiple_docs,
|
|
query_result,
|
|
full_collection_new,
|
|
is_by_vector=1,
|
|
query_vector=query_vector,
|
|
data_type=k,
|
|
vector_name=v,
|
|
)
|
|
|
|
@pytest.mark.parametrize("ef", [None, "invalid", 10.5])
|
|
@pytest.mark.parametrize("doc_num", [10])
|
|
@pytest.mark.parametrize("topk", [10])
|
|
@pytest.mark.parametrize("filter", ["int32_field >= 3 and int32_field <= 7"])
|
|
def test_query_vector_with_HnswQueryParam_invalid(
|
|
self, full_collection: Collection, doc_num, topk, ef, filter
|
|
):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection.schema) for i in range(doc_num)
|
|
]
|
|
batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert")
|
|
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
|
|
doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema)
|
|
query_vector = doc_vectors[v]
|
|
with pytest.raises(Exception) as exc_info:
|
|
full_collection.query(
|
|
Query(
|
|
field_name=v, vector=query_vector, param=HnswQueryParam(ef=ef)
|
|
),
|
|
filter=filter,
|
|
topk=topk,
|
|
)
|
|
assert INCOMPATIBLE_CONSTRUCTOR_ERROR_MSG in str(exc_info.value)
|
|
|
|
@pytest.mark.parametrize("nprobe", [1, 10, 100, 2048])
|
|
@pytest.mark.parametrize("doc_num", [10])
|
|
@pytest.mark.parametrize("topk", [10])
|
|
@pytest.mark.parametrize("filter", ["int32_field >= 3 and int32_field <= 7"])
|
|
@pytest.mark.parametrize(
|
|
"full_schema_ivf", [(True, True, IVFIndexParam())], indirect=True
|
|
)
|
|
def test_query_vector_with_IVFQueryParam_valid(
|
|
self, full_collection_ivf: Collection, nprobe, doc_num, topk, filter
|
|
):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection_ivf.schema) for i in range(doc_num)
|
|
]
|
|
batchdoc_and_check_ivf(
|
|
full_collection_ivf, multiple_docs, doc_num, operator="insert"
|
|
)
|
|
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
|
|
doc_fields, doc_vectors = generate_vectordict_random(
|
|
full_collection_ivf.schema
|
|
)
|
|
if v in ["vector_fp32_field"]:
|
|
query_vector = doc_vectors[v]
|
|
|
|
query_result = full_collection_ivf.query(
|
|
Query(
|
|
field_name=v,
|
|
vector=query_vector,
|
|
param=IVFQueryParam(nprobe=nprobe),
|
|
),
|
|
filter=filter,
|
|
topk=topk,
|
|
)
|
|
assert len(query_result) > 0
|
|
single_querydoc_check(
|
|
multiple_docs,
|
|
query_result,
|
|
full_collection_ivf,
|
|
is_by_vector=1,
|
|
query_vector=query_vector,
|
|
data_type=k,
|
|
vector_name=v,
|
|
)
|
|
|
|
@pytest.mark.parametrize("nprobe", [None, 10.5])
|
|
@pytest.mark.parametrize("doc_num", [10])
|
|
@pytest.mark.parametrize("topk", [10])
|
|
@pytest.mark.parametrize("filter", ["int32_field >= 3 and int32_field <= 7"])
|
|
def test_query_vector_with_IVFQueryParam_invalid(
|
|
self, full_collection_ivf: Collection, nprobe, doc_num, topk, filter
|
|
):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection_ivf.schema) for i in range(doc_num)
|
|
]
|
|
batchdoc_and_check_ivf(
|
|
full_collection_ivf, multiple_docs, doc_num, operator="insert"
|
|
)
|
|
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
|
|
doc_fields, doc_vectors = generate_vectordict_random(
|
|
full_collection_ivf.schema
|
|
)
|
|
if v in ["vector_fp32_field"]:
|
|
print("v:\n")
|
|
print(v)
|
|
query_vector = doc_vectors[v]
|
|
with pytest.raises(Exception) as exc_info:
|
|
full_collection_ivf.query(
|
|
Query(
|
|
field_name=v,
|
|
vector=query_vector,
|
|
param=IVFQueryParam(nprobe=nprobe),
|
|
),
|
|
# filter=filter,
|
|
topk=topk,
|
|
)
|
|
assert INCOMPATIBLE_CONSTRUCTOR_ERROR_MSG in str(exc_info.value)
|
|
|
|
@pytest.mark.parametrize("filter", ["int32_field >= 3 and int32_field <= 7"])
|
|
@pytest.mark.parametrize("doc_num", [10])
|
|
def test_query_vector_with_param_invalid(
|
|
self, full_collection: Collection, doc_num, filter
|
|
):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection.schema) for i in range(doc_num)
|
|
]
|
|
batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert")
|
|
|
|
with pytest.raises(Exception) as exc_info:
|
|
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
|
|
doc_fields, doc_vectors = generate_vectordict_random(
|
|
full_collection.schema
|
|
)
|
|
query_vector = doc_vectors[v]
|
|
if v in ["vector_fp16_field", "vector_fp32_field"]:
|
|
full_collection.query(
|
|
Query(
|
|
field_name=v, vector=query_vector, param=HnswIndexParam()
|
|
),
|
|
filter=filter,
|
|
)
|
|
assert INCOMPATIBLE_FUNCTION_ERROR_MSG in str(exc_info.value)
|
|
|
|
@pytest.mark.parametrize("doc_num", [10])
|
|
@pytest.mark.parametrize(
|
|
"test_case_name,vector_query,expected_error_msg",
|
|
[
|
|
(
|
|
"Non-existent vector field name",
|
|
lambda ref_dense_vector: Query(
|
|
field_name="nonexistent_vector", vector=ref_dense_vector
|
|
),
|
|
"Expected exception for non-existent vector field name",
|
|
),
|
|
(
|
|
"Invalid vector data type for dense vector (string instead of list)",
|
|
lambda ref_dense_vector: Query(
|
|
field_name="vector_fp32_field", vector="invalid_vector_data"
|
|
),
|
|
"Expected exception for invalid dense vector data type",
|
|
),
|
|
(
|
|
"Invalid vector data type for sparse vector (list instead of dict)",
|
|
lambda ref_dense_vector: Query(
|
|
field_name="sparse_fp32", vector=[1.0, 2.0, 3.0]
|
|
),
|
|
"Expected exception for invalid sparse vector data type",
|
|
),
|
|
(
|
|
"Empty vector data for dense vector",
|
|
lambda ref_dense_vector: Query(
|
|
field_name="vector_fp32_field", vector=[]
|
|
),
|
|
"Expected exception for empty dense vector data",
|
|
),
|
|
(
|
|
"Invalid dimension for dense vector",
|
|
lambda ref_dense_vector: Query(
|
|
field_name="vector_fp32_field", vector=[1.0, 2.0]
|
|
), # Only 2 dimensions instead of 128
|
|
"Expected exception for invalid dense vector dimension",
|
|
),
|
|
(
|
|
"Non-existent document ID for by_id query",
|
|
lambda ref_dense_vector: Query(
|
|
field_name="vector_fp32_field", id="999"
|
|
), # Non-existent ID
|
|
"Expected exception for non-existent document ID",
|
|
),
|
|
(
|
|
"Neither vector nor id specified",
|
|
lambda ref_dense_vector: Query(
|
|
field_name="vector_fp32_field"
|
|
), # Neither vector nor id
|
|
"Expected exception for specifying neither vector nor id",
|
|
),
|
|
],
|
|
)
|
|
def test_query_vector_with_vectors_invalid(
|
|
self,
|
|
full_collection: Collection,
|
|
doc_num,
|
|
test_case_name,
|
|
vector_query,
|
|
expected_error_msg,
|
|
):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection.schema) for i in range(doc_num)
|
|
]
|
|
batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert")
|
|
ref_doc_result = full_collection.fetch(ids=["5"])
|
|
assert "5" in ref_doc_result
|
|
ref_doc = ref_doc_result["5"]
|
|
ref_dense_vector = ref_doc.vector("vector_fp32_field")
|
|
|
|
with pytest.raises(Exception) as exc_info:
|
|
full_collection.query([vector_query(ref_dense_vector)])
|
|
assert exc_info.value is not None, expected_error_msg
|
|
|
|
@pytest.mark.parametrize("filter", ["int32_field >= 3 and int32_field <= 7"])
|
|
@pytest.mark.parametrize("doc_num", [10])
|
|
def test_query_invalid_param_incompatible_type(
|
|
self, full_collection: Collection, doc_num, filter
|
|
):
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection.schema) for i in range(doc_num)
|
|
]
|
|
batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert")
|
|
|
|
with pytest.raises(Exception) as exc_info:
|
|
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
|
|
doc_fields, doc_vectors = generate_vectordict_random(
|
|
full_collection.schema
|
|
)
|
|
query_vector = doc_vectors[v]
|
|
full_collection.query(
|
|
Query(field_name=v, vector=query_vector),
|
|
filter=filter,
|
|
param=HnswIndexParam(),
|
|
topk=3,
|
|
)
|
|
|
|
assert "query() got an unexpected keyword argument 'param'" in str(
|
|
exc_info.value
|
|
)
|
|
|
|
|
|
class TestRRFScoreCalculation:
|
|
class MockDoc:
|
|
def __init__(self, id, score=0.0):
|
|
self._id = id
|
|
self._score = score
|
|
|
|
@property
|
|
def id(self):
|
|
return self._id
|
|
|
|
@property
|
|
def score(self):
|
|
return self._score
|
|
|
|
@score.setter
|
|
def score(self, score):
|
|
self._score = score
|
|
|
|
def test_rrf_score_calculation_formula(self):
|
|
k = 60
|
|
|
|
assert abs(calculate_rrf_score(0, k) - 1.0 / 61) < 1e-10, (
|
|
"RRF score for rank 0 should be 1/61"
|
|
)
|
|
assert abs(calculate_rrf_score(1, k) - 1.0 / 62) < 1e-10, (
|
|
"RRF score for rank 1 should be 1/62"
|
|
)
|
|
assert abs(calculate_rrf_score(2, k) - 1.0 / 63) < 1e-10, (
|
|
"RRF score for rank 2 should be 1/63"
|
|
)
|
|
assert abs(calculate_rrf_score(10, k) - 1.0 / 71) < 1e-10, (
|
|
"RRF score for rank 10 should be 1/71"
|
|
)
|
|
|
|
k = 10
|
|
assert abs(calculate_rrf_score(0, k) - 1.0 / 11) < 1e-10, (
|
|
"RRF score for rank 0 with k=10 should be 1/11"
|
|
)
|
|
assert abs(calculate_rrf_score(1, k) - 1.0 / 12) < 1e-10, (
|
|
"RRF score for rank 1 with k=10 should be 1/12"
|
|
)
|
|
|
|
def test_multi_vector_rrf_scores(self):
|
|
query1_results = [self.MockDoc("1"), self.MockDoc("2"), self.MockDoc("3")]
|
|
query2_results = [self.MockDoc("3"), self.MockDoc("1"), self.MockDoc("4")]
|
|
query3_results = [self.MockDoc("2"), self.MockDoc("4"), self.MockDoc("5")]
|
|
query_results = {
|
|
"vector1": query1_results,
|
|
"vector2": query2_results,
|
|
"vector3": query3_results,
|
|
}
|
|
rrf_scores = calculate_multi_vector_rrf_scores(query_results, k=60)
|
|
|
|
expected_doc1_score = 1.0 / 61 + 1.0 / 62
|
|
assert abs(rrf_scores["1"] - expected_doc1_score) < 1e-10, (
|
|
f"RRF score for doc1 mismatch: expected {expected_doc1_score}, got {rrf_scores['1']}"
|
|
)
|
|
expected_doc2_score = 1.0 / 62 + 1.0 / 61
|
|
assert abs(rrf_scores["2"] - expected_doc2_score) < 1e-10, (
|
|
f"RRF score for doc2 mismatch: expected {expected_doc2_score}, got {rrf_scores['2']}"
|
|
)
|
|
expected_doc3_score = 1.0 / 63 + 1.0 / 61
|
|
assert abs(rrf_scores["3"] - expected_doc3_score) < 1e-10, (
|
|
f"RRF score for doc3 mismatch: expected {expected_doc3_score}, got {rrf_scores['3']}"
|
|
)
|
|
expected_doc4_score = 1.0 / 63 + 1.0 / 62
|
|
assert abs(rrf_scores["4"] - expected_doc4_score) < 1e-10, (
|
|
f"RRF score for doc4 mismatch: expected {expected_doc4_score}, got {rrf_scores['4']}"
|
|
)
|
|
|
|
expected_doc5_score = 1.0 / 63
|
|
assert abs(rrf_scores["5"] - expected_doc5_score) < 1e-10, (
|
|
f"RRF score for doc5 mismatch: expected {expected_doc5_score}, got {rrf_scores['5']}"
|
|
)
|
|
sorted_scores = sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)
|
|
expected_order = ["1", "2", "3", "4", "5"]
|
|
actual_order = [item[0] for item in sorted_scores]
|
|
assert actual_order == expected_order, (
|
|
f"RRF score ranking mismatch: expected {expected_order}, got {actual_order}"
|
|
)
|
|
|
|
|
|
class TestCollectionConcurrencyOperations:
|
|
@pytest.mark.parametrize("doc_num", [10])
|
|
def test_concurrent_insert_update_upsert_query(
|
|
self, full_collection: Collection, doc_num
|
|
):
|
|
import threading
|
|
|
|
results = []
|
|
errors = []
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection.schema) for i in range(1000, 1010)
|
|
]
|
|
|
|
batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert")
|
|
|
|
def insert_operation(thread_id):
|
|
try:
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection.schema)
|
|
for i in range(thread_id, thread_id + 5)
|
|
]
|
|
result = full_collection.insert(multiple_docs)
|
|
results.append(("insert", thread_id, len(result)))
|
|
except Exception as e:
|
|
errors.append(("insert", thread_id, str(e)))
|
|
|
|
def update_operation(thread_id):
|
|
try:
|
|
multiple_docs = [
|
|
generate_doc_random(i, full_collection.schema)
|
|
for i in range(1000, 1001)
|
|
]
|
|
result = full_collection.update(multiple_docs)
|
|
results.append(("update", thread_id, len(result)))
|
|
except Exception as e:
|
|
errors.append(("update", thread_id, str(e)))
|
|
|
|
def upsert_operation(thread_id):
|
|
try:
|
|
multiple_docs = [
|
|
generate_doc(i, full_collection.schema)
|
|
for i in range(thread_id, thread_id + 5)
|
|
]
|
|
result = full_collection.upsert(multiple_docs)
|
|
results.append(("upsert", thread_id, len(result)))
|
|
except Exception as e:
|
|
errors.append(("upsert", thread_id, str(e)))
|
|
|
|
def query_operation(thread_id):
|
|
try:
|
|
if thread_id % 3 == 0:
|
|
result = full_collection.query(filter="int32_field > 1", topk=5)
|
|
elif thread_id % 3 == 1:
|
|
result = full_collection.query(filter="bool_field = true", topk=3)
|
|
else:
|
|
query_vector = [0.1] * 128
|
|
result = full_collection.query(
|
|
Query(field_name="vector_fp32_field", vector=query_vector),
|
|
topk=3,
|
|
)
|
|
|
|
results.append(("query", thread_id, len(result)))
|
|
except Exception as e:
|
|
errors.append(("query", thread_id, str(e)))
|
|
|
|
def delete_operation(thread_id):
|
|
try:
|
|
# Delete some existing documents
|
|
delete_ids = (
|
|
[f"{thread_id + 1}", f"{thread_id + 2}"]
|
|
if thread_id < 5
|
|
else [f"{thread_id % 5 + 1}"]
|
|
)
|
|
result = full_collection.delete(delete_ids)
|
|
results.append(("delete", thread_id, len(result)))
|
|
except Exception as e:
|
|
errors.append(("delete", thread_id, str(e)))
|
|
|
|
threads = []
|
|
for i in range(1):
|
|
thread = threading.Thread(target=insert_operation, args=(i,))
|
|
threads.append(thread)
|
|
thread.start()
|
|
for i in range(1):
|
|
thread = threading.Thread(target=update_operation, args=(i,))
|
|
threads.append(thread)
|
|
thread.start()
|
|
for i in range(1):
|
|
thread = threading.Thread(target=upsert_operation, args=(i,))
|
|
threads.append(thread)
|
|
thread.start()
|
|
for i in range(1):
|
|
thread = threading.Thread(target=query_operation, args=(i,))
|
|
threads.append(thread)
|
|
thread.start()
|
|
for i in range(1):
|
|
thread = threading.Thread(target=delete_operation, args=(i,))
|
|
threads.append(thread)
|
|
thread.start()
|
|
|
|
for thread in threads:
|
|
thread.join()
|
|
|
|
insert_results = [r for r in results if r[0] == "insert"]
|
|
update_results = [r for r in results if r[0] == "update"]
|
|
upsert_results = [r for r in results if r[0] == "upsert"]
|
|
query_results = [r for r in results if r[0] == "query"]
|
|
delete_results = [r for r in results if r[0] == "delete"]
|
|
|
|
assert (
|
|
len(insert_results)
|
|
+ len(update_results)
|
|
+ len(upsert_results)
|
|
+ len(query_results)
|
|
+ len(delete_results)
|
|
> 0
|
|
), f"No operations succeeded. Errors: {errors}"
|
|
|
|
critical_errors = [
|
|
e for e in errors if "critical" in e[2].lower() or "fatal" in e[2].lower()
|
|
]
|
|
assert len(critical_errors) == 0, f"Critical errors occurred: {critical_errors}"
|