wehub-resource-sync
1112 lines
36 KiB
Python
1112 lines
36 KiB
Python
import logging
|
|
import pytest
|
|
|
|
|
|
from zvec import (
|
|
CollectionOption,
|
|
InvertIndexParam,
|
|
HnswIndexParam,
|
|
FieldSchema,
|
|
VectorSchema,
|
|
CollectionSchema,
|
|
Collection,
|
|
Doc,
|
|
Query,
|
|
StatusCode,
|
|
)
|
|
from distance_helper import *
|
|
from fixture_helper import *
|
|
from doc_helper import *
|
|
|
|
Maximum = 1024
|
|
|
|
DOCID_VALID_LIST = [
|
|
"1valid_Id",
|
|
"123.45",
|
|
"123abc",
|
|
"-!@#$%+=.123abc_+",
|
|
"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ123456789012",
|
|
]
|
|
DOCID_INVALID_LIST = [
|
|
None,
|
|
"",
|
|
"()qsd123",
|
|
" ",
|
|
"/&AS12",
|
|
"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890121",
|
|
]
|
|
|
|
FIELD_VALUE_VALID_LIST = [
|
|
(
|
|
"bool_field",
|
|
[
|
|
None,
|
|
True,
|
|
False,
|
|
],
|
|
),
|
|
(
|
|
"float_field",
|
|
[
|
|
None,
|
|
0.0,
|
|
-1.0,
|
|
1.0,
|
|
3.4028235e38,
|
|
-3.4028235e38,
|
|
1.17549435e-38,
|
|
-1.17549435e-38,
|
|
float("inf"),
|
|
float("-inf"),
|
|
],
|
|
),
|
|
(
|
|
"double_field",
|
|
[
|
|
None,
|
|
0.0,
|
|
-1.0,
|
|
1.0,
|
|
1.7976931348623157e308,
|
|
-1.7976931348623157e308,
|
|
2.2250738585072014e-308,
|
|
-2.2250738585072014e-308,
|
|
float("inf"),
|
|
float("-inf"),
|
|
],
|
|
),
|
|
(
|
|
"int32_field",
|
|
[
|
|
None,
|
|
0,
|
|
1,
|
|
-1,
|
|
2147483647,
|
|
-2147483648,
|
|
],
|
|
),
|
|
(
|
|
"int64_field",
|
|
[
|
|
None,
|
|
0,
|
|
1,
|
|
-1,
|
|
9223372036854775807,
|
|
-9223372036854775808,
|
|
],
|
|
),
|
|
(
|
|
"uint32_field",
|
|
[
|
|
None,
|
|
0,
|
|
1,
|
|
4294967295,
|
|
],
|
|
),
|
|
(
|
|
"uint64_field",
|
|
[
|
|
None,
|
|
0,
|
|
1,
|
|
18446744073709551615,
|
|
],
|
|
),
|
|
(
|
|
"string_field",
|
|
[
|
|
None,
|
|
"",
|
|
"a",
|
|
"test_name",
|
|
"这是一个中文名称测试",
|
|
"a" * 1000,
|
|
],
|
|
),
|
|
(
|
|
"array_bool_field",
|
|
[
|
|
None,
|
|
[],
|
|
[True],
|
|
[False, True],
|
|
[True, False, True, False] * 10,
|
|
],
|
|
),
|
|
(
|
|
"array_float_field",
|
|
[
|
|
None,
|
|
[],
|
|
[0.0],
|
|
[1.0, 2.0, 3.0],
|
|
[3.4028235e38, -3.4028235e38],
|
|
],
|
|
),
|
|
(
|
|
"array_double_field",
|
|
[
|
|
None,
|
|
[],
|
|
[0.0],
|
|
[1.0, 2.0, 3.0],
|
|
[1.7976931348623157e308, -1.7976931348623157e308],
|
|
],
|
|
),
|
|
(
|
|
"array_int32_field",
|
|
[
|
|
None,
|
|
[],
|
|
[0],
|
|
[1, 2, 3],
|
|
[2147483647, -2147483648],
|
|
],
|
|
),
|
|
(
|
|
"array_int64_field",
|
|
[
|
|
None,
|
|
[],
|
|
[0],
|
|
[1, 2, 3],
|
|
[9223372036854775807, -9223372036854775808],
|
|
],
|
|
),
|
|
(
|
|
"array_uint32_field",
|
|
[
|
|
None,
|
|
[],
|
|
[0],
|
|
[1, 2, 3],
|
|
[4294967295],
|
|
],
|
|
),
|
|
(
|
|
"array_uint64_field",
|
|
[
|
|
None,
|
|
[],
|
|
[0],
|
|
[1, 2, 3],
|
|
[18446744073709551615],
|
|
],
|
|
),
|
|
(
|
|
"array_string_field",
|
|
[
|
|
None,
|
|
[],
|
|
[""],
|
|
["a", "b", "c"],
|
|
["test_string", "测试字符串"],
|
|
["a" * 100] * 5,
|
|
],
|
|
),
|
|
]
|
|
FIELD_VALUE_INVALID_LIST = [
|
|
(
|
|
"bool_field",
|
|
["True", "False", "", "测试"],
|
|
),
|
|
("float_field", ["invalid", [1.0], {"value": 1.0}, "测试"]),
|
|
("double_field", ["invalid", [1.0], {"value": 1.0}, "测试"]),
|
|
(
|
|
"int32_field",
|
|
["invalid", [1], {"value": 1}, 2147483648, -2147483649, "测试"],
|
|
),
|
|
(
|
|
"int64_field",
|
|
[
|
|
"invalid",
|
|
[1],
|
|
{"value": 1},
|
|
9223372036854775808,
|
|
-9223372036854775809,
|
|
"测试",
|
|
],
|
|
),
|
|
(
|
|
"uint32_field",
|
|
["invalid", [1], {"value": 1}, 4294967296, -1, "测试"],
|
|
),
|
|
(
|
|
"uint64_field",
|
|
["invalid", [1], {"value": 1}, 18446744073709551616, -1, "测试"],
|
|
),
|
|
(
|
|
"string_field",
|
|
[
|
|
123,
|
|
12.34,
|
|
True,
|
|
["array"],
|
|
{"key": "value"},
|
|
],
|
|
),
|
|
(
|
|
"array_bool_field",
|
|
[True, False, [True, "invalid"], {"key": True}, "测试"],
|
|
),
|
|
(
|
|
"array_float_field",
|
|
[[1.0, "invalid"], [1.0, None], "invalid", [1.0, [2.0]], 1.0, "测试"],
|
|
),
|
|
(
|
|
"array_double_field",
|
|
[[1.0, "invalid"], [1.0, None], "invalid", [1.0, [2.0]], 1.0, "测试"],
|
|
),
|
|
(
|
|
"array_int32_field",
|
|
[[1, "invalid"], [1, None], "invalid", [1, [2]], 1, "测试"],
|
|
),
|
|
(
|
|
"array_int64_field",
|
|
[[1, "invalid"], [1, None], "invalid", [1, [2]], 1, "测试"],
|
|
),
|
|
(
|
|
"array_uint32_field",
|
|
[[1, "invalid"], [1, None], [1, -1], "invalid", [1, [2]], 1, "测试"],
|
|
),
|
|
(
|
|
"array_uint64_field",
|
|
[[1, "invalid"], [1, None], [1, -1], "invalid", [1, [2]], 1, "测试"],
|
|
),
|
|
(
|
|
"array_string_field",
|
|
[["valid", 123], ["valid", None], "invalid", [["nested"]], 123, "测试"],
|
|
),
|
|
]
|
|
|
|
VECTOR_VALUE_VALID_LIST = [
|
|
(
|
|
"vector_fp32_field",
|
|
[
|
|
[0.0] * 128,
|
|
[1.0] * 128,
|
|
[-1.0] * 128,
|
|
[float("inf")] * 128,
|
|
[float("-inf")] * 128,
|
|
[i / 128.0 for i in range(128)],
|
|
[-i / 128.0 for i in range(128)],
|
|
],
|
|
),
|
|
(
|
|
"vector_fp16_field",
|
|
[
|
|
[0.0] * 128,
|
|
[1.0] * 128,
|
|
[-1.0] * 128,
|
|
[float("inf")] * 128,
|
|
[float("-inf")] * 128,
|
|
[i / 128.0 for i in range(128)],
|
|
[-i / 128.0 for i in range(128)],
|
|
],
|
|
),
|
|
("vector_int8_field", [[100] * 128, [0] * 128, [-100] * 128]),
|
|
(
|
|
"sparse_vector_fp32_field",
|
|
[
|
|
{0: 1.0},
|
|
{0: 0.0, 1: 1.0, 2: -1.0},
|
|
{0: float("inf"), 1: float("-inf")},
|
|
{i: float(i) for i in range(10)},
|
|
{128: 1.0, 256: -1.0, 512: 0.5},
|
|
],
|
|
),
|
|
(
|
|
"sparse_vector_fp16_field",
|
|
[
|
|
{0: 1.0},
|
|
{0: 0.0, 1: 1.0, 2: -1.0},
|
|
{0: float("inf"), 1: float("-inf")},
|
|
{i: float(i) for i in range(10)},
|
|
{128: 1.0, 256: -1.0, 512: 0.5},
|
|
],
|
|
),
|
|
]
|
|
VECTOR_VALUE_INVALID_LIST = [
|
|
(
|
|
"vector_fp32_field",
|
|
[
|
|
None,
|
|
[],
|
|
[0.0] * 127,
|
|
[0.0] * 129,
|
|
[0.0] * 1000,
|
|
["invalid"],
|
|
[0, 1, 2],
|
|
[None] * 128,
|
|
],
|
|
),
|
|
(
|
|
"vector_fp16_field",
|
|
[
|
|
None,
|
|
[],
|
|
[0.0] * 127,
|
|
[0.0] * 129,
|
|
[0.0] * 1000,
|
|
["invalid"],
|
|
[0, 1, 2],
|
|
[None] * 128,
|
|
],
|
|
),
|
|
(
|
|
"vector_int8_field",
|
|
[
|
|
None,
|
|
[],
|
|
[1] * 127,
|
|
[10] * 129,
|
|
[0] * 1000,
|
|
["invalid"],
|
|
[0, 1, 2],
|
|
[None] * 128,
|
|
],
|
|
),
|
|
(
|
|
"sparse_vector_fp32_field",
|
|
[
|
|
None,
|
|
"invalid",
|
|
{None: 1.0},
|
|
{"0": 1.0},
|
|
{0: "invalid"},
|
|
{0: None},
|
|
{-1: 1.0},
|
|
],
|
|
),
|
|
(
|
|
"sparse_vector_fp16_field",
|
|
[
|
|
None,
|
|
"invalid",
|
|
{None: 1.0},
|
|
{"0": 1.0},
|
|
{0: "invalid"},
|
|
{0: None},
|
|
{-1: 1.0},
|
|
],
|
|
),
|
|
]
|
|
|
|
UPDATE_PARTIAL_VALUE = [
|
|
(
|
|
"partial_fields",
|
|
{"string_field": "partially_updated_test", "float_field": 95.5},
|
|
{},
|
|
),
|
|
("dense_vector_only", {}, {"vector_fp32_field": [0.3] * 128}),
|
|
("dense_vector_only", {}, {"vector_fp16_field": [0.6] * 128}),
|
|
("dense_vector_only", {}, {"vector_int8_field": [3] * 128}),
|
|
("sparse_vector_only", {}, {"sparse_vector_fp32_field": {1: 2.0, 2: 3.0, 4: 4.0}}),
|
|
(
|
|
"sparse_vector_only",
|
|
{},
|
|
{"sparse_vector_fp16_field": {10: 2.1, 20: 3.1, 40: 4.1}},
|
|
),
|
|
(
|
|
"fields_and_vectors",
|
|
{"string_field": "fully_updated_test", "bool_field": False},
|
|
{
|
|
"vector_fp32_field": [0.4] * 128,
|
|
"sparse_vector_fp32_field": {1: 3.0, 3: 5.0},
|
|
},
|
|
),
|
|
]
|
|
|
|
|
|
# ==================== helper ====================
|
|
def singledoc_and_check(
|
|
collection: Collection, insert_doc, operator="insert", is_delete=1
|
|
):
|
|
if operator == "insert":
|
|
result = collection.insert(insert_doc)
|
|
elif operator == "upsert":
|
|
result = collection.upsert(insert_doc)
|
|
elif operator == "update":
|
|
result = collection.update(insert_doc)
|
|
else:
|
|
logging.error("operator value is error!")
|
|
|
|
assert bool(result)
|
|
assert result.ok()
|
|
|
|
stats = collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 1
|
|
|
|
fetched_docs = collection.fetch([insert_doc.id])
|
|
assert len(fetched_docs) == 1
|
|
assert insert_doc.id in fetched_docs
|
|
|
|
fetched_doc = fetched_docs[insert_doc.id]
|
|
|
|
assert is_doc_equal(fetched_doc, insert_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"
|
|
)
|
|
|
|
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
|
|
if v != {}:
|
|
query_result = collection.query(
|
|
Query(field_name=v, vector=insert_doc.vectors[v]),
|
|
topk=10,
|
|
)
|
|
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 == insert_doc.id:
|
|
found_doc = doc
|
|
break
|
|
assert found_doc is not None, (
|
|
f"Inserted document {insert_doc.id} not found in query results"
|
|
)
|
|
assert is_doc_equal(found_doc, insert_doc, collection.schema, True, False)
|
|
if is_delete == 1:
|
|
collection.delete(insert_doc.id)
|
|
assert collection.stats.doc_count == 0, "Document should be deleted"
|
|
|
|
|
|
def updatedoc_partial_check(
|
|
collection, update_doc_partial, update_doc_full, operator="update", is_delete=1
|
|
):
|
|
if operator == "upsert":
|
|
result = collection.upsert(update_doc_partial)
|
|
elif operator == "update":
|
|
result = collection.update(update_doc_partial)
|
|
else:
|
|
logging.error("operator value is error!")
|
|
|
|
assert bool(result)
|
|
assert result.ok()
|
|
|
|
stats = collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 1
|
|
|
|
fetched_docs = collection.fetch([update_doc_partial.id])
|
|
assert len(fetched_docs) == 1, (
|
|
f"fetched_docs={fetched_docs},Expected 1 fetched document, but got {len(fetched_docs)}"
|
|
)
|
|
assert update_doc_partial.id in fetched_docs, (
|
|
f"Expected document ID {update_doc_partial.id} in fetched documents"
|
|
)
|
|
|
|
fetched_doc = fetched_docs[update_doc_partial.id]
|
|
assert is_doc_equal(fetched_doc, update_doc_full, 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"
|
|
)
|
|
|
|
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
|
|
if v != {}:
|
|
query_result = collection.query(
|
|
Query(field_name=v, vector=update_doc_full.vectors[v]),
|
|
topk=10,
|
|
)
|
|
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 == update_doc_partial.id:
|
|
found_doc = doc
|
|
break
|
|
assert found_doc is not None, (
|
|
f"Inserted document {update_doc_partial.id} not found in query results"
|
|
)
|
|
assert is_doc_equal(
|
|
found_doc, update_doc_full, collection.schema, True, False
|
|
)
|
|
if is_delete == 1:
|
|
collection.delete(update_doc_partial.id)
|
|
assert collection.stats.doc_count == 0, "Document should be deleted"
|
|
|
|
|
|
def batchdoc_and_check(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,
|
|
)
|
|
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, True, False)
|
|
|
|
|
|
# ==================== Tests ====================
|
|
# ----------------------------
|
|
# Collection Insert Test Case
|
|
# ----------------------------
|
|
|
|
|
|
class TestCollectionInsert:
|
|
def test_insert(self, full_collection: Collection):
|
|
single_doc = generate_doc(1, full_collection.schema)
|
|
singledoc_and_check(full_collection, single_doc)
|
|
|
|
@pytest.mark.parametrize("doc_num", [1, 5, Maximum])
|
|
def test_insert_batch(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)
|
|
|
|
def test_insert_duplicate(self, full_collection: Collection):
|
|
insert_doc = generate_doc(1, full_collection.schema)
|
|
|
|
result = full_collection.insert(insert_doc)
|
|
assert result.code().value == 0
|
|
assert result.ok()
|
|
|
|
# Verify documents were inserted
|
|
stats = full_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 1
|
|
|
|
insert_doc_duplicate = full_collection.insert(insert_doc)
|
|
assert bool(insert_doc_duplicate)
|
|
assert insert_doc_duplicate.code() == StatusCode.ALREADY_EXISTS, (
|
|
f"Second insert operation should fail with ALREADY_EXISTS, but got code {insert_doc_duplicate.code()}"
|
|
)
|
|
|
|
stats = full_collection.stats
|
|
assert stats is not None, "Collection stats should not be None"
|
|
assert stats.doc_count == 1, (
|
|
f"Document count should still be 1 after failed insert, but got {stats.doc_count}"
|
|
)
|
|
|
|
@pytest.mark.parametrize("doc_id", DOCID_VALID_LIST)
|
|
def test_insert_docid_valid(self, full_collection: Collection, doc_id):
|
|
insert_doc = generate_doc_random(doc_id, full_collection.schema)
|
|
singledoc_and_check(full_collection, insert_doc)
|
|
|
|
@pytest.mark.parametrize("doc_id", DOCID_INVALID_LIST)
|
|
def test_insert_docid_invalid(self, full_collection: Collection, doc_id):
|
|
insert_doc = generate_doc_random(doc_id, full_collection.schema)
|
|
|
|
with pytest.raises(Exception) as exc_info:
|
|
full_collection.insert(insert_doc)
|
|
|
|
assert exc_info.value is not None
|
|
stats = full_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 0
|
|
|
|
@pytest.mark.parametrize("field_name, field_values", FIELD_VALUE_VALID_LIST)
|
|
@pytest.mark.parametrize(
|
|
"full_schema_new",
|
|
[(True, True, HnswIndexParam()), (False, True, HnswIndexParam())],
|
|
indirect=True,
|
|
)
|
|
def test_insert_fields_valid(
|
|
self, full_collection_new: Collection, field_name: str, field_values, request
|
|
):
|
|
for i, field_value in enumerate(field_values):
|
|
doc_id = str(field_value) if field_name == "id" else str(i)
|
|
doc_fields, doc_vectors = generate_vectordict_random(
|
|
full_collection_new.schema
|
|
)
|
|
full_schema_params = request.getfixturevalue("full_schema_new")
|
|
target_field = None
|
|
for field in full_schema_params.fields:
|
|
if field.name == field_name:
|
|
target_field = field
|
|
break
|
|
doc_fields[field_name] = field_value
|
|
insert_doc = Doc(id=doc_id, fields=doc_fields, vectors=doc_vectors)
|
|
if target_field and not target_field.nullable and field_value is None:
|
|
with pytest.raises(Exception) as exc_info:
|
|
full_collection_new.insert(insert_doc)
|
|
assert exc_info.value is not None
|
|
else:
|
|
singledoc_and_check(full_collection_new, insert_doc)
|
|
|
|
@pytest.mark.parametrize("field_name, field_values", FIELD_VALUE_INVALID_LIST)
|
|
def test_insert_fields_invalid(
|
|
self, full_collection: Collection, field_name: str, field_values
|
|
):
|
|
for i, field_value in enumerate(field_values):
|
|
doc_id = str(field_value) if field_name == "id" else str(i)
|
|
doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema)
|
|
doc_fields[field_name] = field_value
|
|
insert_doc = Doc(id=doc_id, fields=doc_fields, vectors=doc_vectors)
|
|
|
|
with pytest.raises(Exception) as exc_info:
|
|
full_collection.insert(insert_doc)
|
|
assert exc_info.value is not None
|
|
stats = full_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 0
|
|
|
|
@pytest.mark.parametrize("vector_field, vector_values", VECTOR_VALUE_VALID_LIST)
|
|
def test_insert_vector_valid(
|
|
self, full_collection: Collection, vector_field: str, vector_values
|
|
):
|
|
for i, vector_value in enumerate(vector_values):
|
|
doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema)
|
|
|
|
doc_vectors[vector_field] = vector_value
|
|
|
|
insert_doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors)
|
|
|
|
singledoc_and_check(full_collection, insert_doc)
|
|
|
|
@pytest.mark.parametrize("vector_field, vector_values", VECTOR_VALUE_INVALID_LIST)
|
|
def test_insert_vector_invalid(
|
|
self, full_collection: Collection, vector_field: str, vector_values
|
|
):
|
|
for i, vector_value in enumerate(vector_values):
|
|
doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema)
|
|
doc_vectors[vector_field] = vector_value
|
|
insert_doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors)
|
|
with pytest.raises(Exception) as exc_info:
|
|
full_collection.insert(insert_doc)
|
|
|
|
assert exc_info.value is not None
|
|
stats = full_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 0
|
|
|
|
|
|
class TestCollectionUpdate:
|
|
def test_update(self, full_collection: Collection):
|
|
insert_doc = generate_doc(1, full_collection.schema)
|
|
singledoc_and_check(full_collection, insert_doc, is_delete=0)
|
|
updated_doc = generate_update_doc(1, full_collection.schema)
|
|
singledoc_and_check(full_collection, updated_doc, operator="update")
|
|
|
|
@pytest.mark.parametrize("doc_num", [1, 5, Maximum])
|
|
def test_update_batch(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)
|
|
multiple_update_docs = [
|
|
generate_update_doc(i, full_collection.schema) for i in range(doc_num)
|
|
]
|
|
batchdoc_and_check(
|
|
full_collection, multiple_update_docs, doc_num, operator="update"
|
|
)
|
|
|
|
def test_empty_collection_update(self, full_collection: Collection):
|
|
updated_doc = generate_update_doc(1, full_collection.schema)
|
|
result = full_collection.update(updated_doc)
|
|
assert bool(result), f"Expected 1 result, but got {len(result)}"
|
|
assert result.code() == StatusCode.NOT_FOUND, (
|
|
f"Update operation should fail with NOT_FOUND, but got code {result.code()}"
|
|
)
|
|
fetched_docs = full_collection.fetch([updated_doc.id])
|
|
assert len(fetched_docs) == 0
|
|
|
|
stats = full_collection.stats
|
|
assert stats is not None, "Collection stats should not be None"
|
|
assert stats.doc_count == 0, (
|
|
f"Document count should be 0, but got {stats.doc_count}"
|
|
)
|
|
|
|
@pytest.mark.parametrize("doc_num", [1, 5, Maximum])
|
|
def test_empty_collection_update_batch(self, full_collection: Collection, doc_num):
|
|
multiple_update_docs = [
|
|
generate_update_doc(i, full_collection.schema) for i in range(doc_num)
|
|
]
|
|
result = full_collection.update(multiple_update_docs)
|
|
assert len(result) == len(multiple_update_docs), (
|
|
f"Expected {len(multiple_update_docs)} results, but got {len(result)}"
|
|
)
|
|
for item in result:
|
|
assert item.code() == StatusCode.NOT_FOUND, (
|
|
f"Update operation should fail with NOT_FOUND, but got code {item.code()}"
|
|
)
|
|
|
|
stats = full_collection.stats
|
|
assert stats is not None, "Collection stats should not be None"
|
|
assert stats.doc_count == 0, (
|
|
f"Document count should be 0, but got {stats.doc_count}"
|
|
)
|
|
|
|
doc_ids = [doc.id for doc in multiple_update_docs]
|
|
fetched_docs = full_collection.fetch(doc_ids)
|
|
assert len(fetched_docs) == 0
|
|
|
|
@pytest.mark.parametrize("field_name, field_values", FIELD_VALUE_VALID_LIST)
|
|
@pytest.mark.parametrize(
|
|
"full_schema_new",
|
|
[(True, True, HnswIndexParam()), (False, True, HnswIndexParam())],
|
|
indirect=True,
|
|
)
|
|
def test_update_fields_valid(
|
|
self, full_collection_new: Collection, field_name: str, field_values, request
|
|
):
|
|
for i, field_value in enumerate(field_values):
|
|
insert_doc = generate_doc(i, full_collection_new.schema)
|
|
singledoc_and_check(full_collection_new, insert_doc, is_delete=0)
|
|
update_doc_fields, update_doc_vectors = generate_vectordict_random(
|
|
full_collection_new.schema
|
|
)
|
|
full_schema_params = request.getfixturevalue("full_schema_new")
|
|
target_field = None
|
|
for field in full_schema_params.fields:
|
|
if field.name == field_name:
|
|
target_field = field
|
|
break
|
|
update_doc_fields[field_name] = field_value
|
|
update_doc = Doc(
|
|
id=str(i), fields=update_doc_fields, vectors=update_doc_vectors
|
|
)
|
|
if target_field and not target_field.nullable and field_value is None:
|
|
with pytest.raises(Exception) as exc_info:
|
|
update_doc_fields[field_name] = field_value
|
|
full_collection_new.update(update_doc)
|
|
assert exc_info.value is not None
|
|
full_collection_new.delete(insert_doc.id)
|
|
else:
|
|
singledoc_and_check(
|
|
full_collection_new, update_doc, operator="update", is_delete=1
|
|
)
|
|
|
|
@pytest.mark.parametrize("field_name, field_values", FIELD_VALUE_INVALID_LIST)
|
|
def test_update_fields_invalid(
|
|
self, full_collection: Collection, field_name: str, field_values
|
|
):
|
|
for i, field_value in enumerate(field_values):
|
|
insert_doc = generate_doc(i, full_collection.schema)
|
|
singledoc_and_check(full_collection, insert_doc, is_delete=0)
|
|
update_doc_fields, update_doc_vectors = generate_vectordict_random(
|
|
full_collection.schema
|
|
)
|
|
update_doc_fields[field_name] = field_value
|
|
update_doc = Doc(
|
|
id=str(i), fields=update_doc_fields, vectors=update_doc_vectors
|
|
)
|
|
|
|
with pytest.raises(Exception) as exc_info:
|
|
full_collection.update(update_doc)
|
|
|
|
assert exc_info.value is not None
|
|
full_collection.delete(insert_doc.id)
|
|
stats = full_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 0
|
|
|
|
@pytest.mark.parametrize("vector_field, vector_values", VECTOR_VALUE_VALID_LIST)
|
|
def test_update_doc_vector_valid(
|
|
self,
|
|
full_collection: Collection,
|
|
collection_temp_dir,
|
|
collection_option,
|
|
vector_field: str,
|
|
vector_values,
|
|
):
|
|
for i, vector_value in enumerate(vector_values):
|
|
insert_doc = generate_doc(i, full_collection.schema)
|
|
singledoc_and_check(full_collection, insert_doc, is_delete=0)
|
|
update_doc_fields, update_doc_vectors = generate_vectordict_random(
|
|
full_collection.schema
|
|
)
|
|
update_doc_vectors[vector_field] = vector_value
|
|
update_doc = Doc(
|
|
id=str(i), fields=update_doc_fields, vectors=update_doc_vectors
|
|
)
|
|
singledoc_and_check(full_collection, update_doc, operator="update")
|
|
|
|
@pytest.mark.parametrize("vector_field, vector_values", VECTOR_VALUE_INVALID_LIST)
|
|
def test_update_doc_vector_invalid(
|
|
self,
|
|
full_collection: Collection,
|
|
collection_temp_dir,
|
|
collection_option,
|
|
vector_field: str,
|
|
vector_values,
|
|
):
|
|
for i, vector_value in enumerate(vector_values):
|
|
insert_doc = generate_doc(i, full_collection.schema)
|
|
singledoc_and_check(full_collection, insert_doc, is_delete=0)
|
|
update_doc_fields, update_doc_vectors = generate_vectordict_random(
|
|
full_collection.schema
|
|
)
|
|
update_doc_vectors[vector_field] = vector_value
|
|
update_doc = Doc(
|
|
id=str(i), fields=update_doc_fields, vectors=update_doc_vectors
|
|
)
|
|
with pytest.raises(Exception) as exc_info:
|
|
full_collection.update(update_doc)
|
|
assert exc_info.value is not None
|
|
full_collection.delete(insert_doc.id)
|
|
stats = full_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 0
|
|
|
|
@pytest.mark.parametrize(
|
|
"update_type, fields_to_update, vectors_to_update", UPDATE_PARTIAL_VALUE
|
|
)
|
|
def test_update_partial_fields(
|
|
self,
|
|
full_collection: Collection,
|
|
collection_temp_dir,
|
|
collection_option,
|
|
update_type: str,
|
|
fields_to_update: dict,
|
|
vectors_to_update: dict,
|
|
doc_id=1,
|
|
):
|
|
insert_doc = generate_doc(doc_id, full_collection.schema)
|
|
singledoc_and_check(full_collection, insert_doc, is_delete=0)
|
|
|
|
update_doc_fields, update_doc_vectors = insert_doc.fields, insert_doc.vectors
|
|
for k, v in fields_to_update.items():
|
|
update_doc_fields[k] = v
|
|
for k, v in vectors_to_update.items():
|
|
update_doc_vectors[k] = v
|
|
|
|
update_doc_full = Doc(
|
|
id=str(doc_id), fields=update_doc_fields, vectors=update_doc_vectors
|
|
)
|
|
|
|
update_doc_partial = Doc(
|
|
id=str(doc_id), fields=fields_to_update, vectors=vectors_to_update
|
|
)
|
|
|
|
updatedoc_partial_check(
|
|
full_collection,
|
|
update_doc_partial,
|
|
update_doc_full,
|
|
operator="update",
|
|
is_delete=1,
|
|
)
|
|
|
|
|
|
class TestCollectionUpsert:
|
|
def test_new_doc_upsert(self, full_collection: Collection):
|
|
single_doc = generate_doc(1, full_collection.schema)
|
|
singledoc_and_check(full_collection, single_doc, operator="upsert", is_delete=1)
|
|
|
|
@pytest.mark.parametrize("doc_num", [1, 5, Maximum])
|
|
def test_new_doc_upsert_batch(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="upsert")
|
|
|
|
def test_existing_doc_upsert(self, full_collection: Collection):
|
|
insert_doc = generate_doc(1, full_collection.schema)
|
|
singledoc_and_check(full_collection, insert_doc, is_delete=0)
|
|
updated_doc = generate_update_doc(1, full_collection.schema)
|
|
singledoc_and_check(full_collection, updated_doc, operator="upsert")
|
|
|
|
@pytest.mark.parametrize("doc_id", DOCID_VALID_LIST)
|
|
def test_upsert_docid_valid(self, full_collection: Collection, doc_id):
|
|
upsert_doc = generate_doc_random(doc_id, full_collection.schema)
|
|
singledoc_and_check(full_collection, upsert_doc, operator="upsert", is_delete=1)
|
|
|
|
@pytest.mark.parametrize("doc_id", DOCID_INVALID_LIST)
|
|
def test_upsert_docid_invalid(self, full_collection: Collection, doc_id):
|
|
upsert_doc = generate_doc_random(doc_id, full_collection.schema)
|
|
|
|
with pytest.raises(Exception) as exc_info:
|
|
full_collection.upsert(upsert_doc)
|
|
|
|
assert exc_info.value is not None
|
|
|
|
stats = full_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 0
|
|
|
|
@pytest.mark.parametrize("field_name, field_values", FIELD_VALUE_VALID_LIST)
|
|
@pytest.mark.parametrize(
|
|
"full_schema_new",
|
|
[(True, True, HnswIndexParam()), (False, True, HnswIndexParam())],
|
|
indirect=True,
|
|
)
|
|
def test_upsert_fields_valid(
|
|
self, full_collection_new: Collection, field_name: str, field_values, request
|
|
):
|
|
for i, field_value in enumerate(field_values):
|
|
doc_id = str(field_value) if field_name == "id" else str(i)
|
|
doc_fields, doc_vectors = generate_vectordict_random(
|
|
full_collection_new.schema
|
|
)
|
|
|
|
full_schema_params = request.getfixturevalue("full_schema_new")
|
|
target_field = None
|
|
for field in full_schema_params.fields:
|
|
if field.name == field_name:
|
|
target_field = field
|
|
break
|
|
doc_fields[field_name] = field_value
|
|
upsert_doc = Doc(id=doc_id, fields=doc_fields, vectors=doc_vectors)
|
|
if target_field and not target_field.nullable and field_value is None:
|
|
with pytest.raises(Exception) as exc_info:
|
|
full_collection_new.upsert(upsert_doc)
|
|
assert exc_info.value is not None
|
|
else:
|
|
singledoc_and_check(
|
|
full_collection_new, upsert_doc, operator="upsert", is_delete=1
|
|
)
|
|
|
|
@pytest.mark.parametrize("field_name, field_values", FIELD_VALUE_INVALID_LIST)
|
|
def test_upsert_fields_invalid(
|
|
self, full_collection: Collection, field_name: str, field_values
|
|
):
|
|
for i, field_value in enumerate(field_values):
|
|
doc_id = str(field_value) if field_name == "id" else str(i)
|
|
doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema)
|
|
doc_fields[field_name] = field_value
|
|
upsert_doc = Doc(id=doc_id, fields=doc_fields, vectors=doc_vectors)
|
|
|
|
with pytest.raises(Exception) as exc_info:
|
|
full_collection.upsert(upsert_doc)
|
|
assert exc_info.value is not None
|
|
stats = full_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 0
|
|
|
|
@pytest.mark.parametrize("vector_field, vector_values", VECTOR_VALUE_VALID_LIST)
|
|
def test_upsert_vector_valid(
|
|
self, full_collection: Collection, vector_field: str, vector_values
|
|
):
|
|
for i, vector_value in enumerate(vector_values):
|
|
doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema)
|
|
|
|
doc_vectors[vector_field] = vector_value
|
|
|
|
upsert_doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors)
|
|
|
|
singledoc_and_check(
|
|
full_collection, upsert_doc, operator="upsert", is_delete=1
|
|
)
|
|
|
|
@pytest.mark.parametrize("vector_field, vector_values", VECTOR_VALUE_INVALID_LIST)
|
|
def test_upsert_vector_invalid(
|
|
self, full_collection: Collection, vector_field: str, vector_values
|
|
):
|
|
for i, vector_value in enumerate(vector_values):
|
|
doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema)
|
|
doc_vectors[vector_field] = vector_value
|
|
upsert_doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors)
|
|
with pytest.raises(Exception) as exc_info:
|
|
full_collection.upsert(upsert_doc)
|
|
|
|
assert exc_info.value is not None
|
|
stats = full_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 0
|
|
|
|
|
|
class TestCollectionDelete:
|
|
@pytest.mark.parametrize("doc_num", [1, 5, Maximum])
|
|
def test_delete_batch(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_ids = [doc.id for doc in multiple_docs]
|
|
result = full_collection.delete(doc_ids)
|
|
assert len(result) == len(doc_ids)
|
|
for item in result:
|
|
assert item.ok()
|
|
|
|
def test_delete_non_exist(self, full_collection: Collection):
|
|
result = full_collection.delete("non_existing_id")
|
|
assert result.code().value == 1
|
|
assert result.code() == StatusCode.NOT_FOUND
|
|
|
|
@pytest.mark.parametrize("doc_num", [5])
|
|
def test_delete_batch_part_non_exist(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_ids = [doc.id for doc in multiple_docs]
|
|
doc_ids.extend([str(doc_num), str(doc_num + 1)])
|
|
result = full_collection.delete(doc_ids)
|
|
|
|
assert len(result) == len(doc_ids)
|
|
for i in range(len(result)):
|
|
if i < doc_num:
|
|
assert result[i].ok()
|
|
else:
|
|
assert result[i].code().value == 1
|
|
assert result[i].code() == StatusCode.NOT_FOUND
|
|
|
|
@pytest.mark.parametrize("doc_num", [5])
|
|
def test_delete_by_filter(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")
|
|
|
|
result = full_collection.delete_by_filter("int32_field > 0")
|
|
assert result is None
|
|
|
|
def test_delete_empty_ids(self, full_collection: Collection):
|
|
result = full_collection.delete([])
|
|
assert len(result) == 0
|