wehub-resource-sync
1721 lines
62 KiB
Python
1721 lines
62 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 may obtain a copy of the 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 fixture_helper import *
|
|
from doc_helper import *
|
|
from params_helper import *
|
|
|
|
|
|
class TestDDL:
|
|
def test_collection_stats(self, basic_collection: Collection):
|
|
assert basic_collection.stats is not None
|
|
stats = basic_collection.stats
|
|
assert stats.doc_count == 0
|
|
assert len(stats.index_completeness) == 2
|
|
assert stats.index_completeness["dense"] == 1
|
|
assert stats.index_completeness["sparse"] == 1
|
|
|
|
def test_collection_destroy(
|
|
self, basic_collection: Collection, collection_temp_dir, collection_option
|
|
):
|
|
doc = generate_doc(1, basic_collection.schema)
|
|
|
|
result = basic_collection.insert(doc)
|
|
assert bool(result)
|
|
assert result.ok()
|
|
|
|
stats = basic_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 1
|
|
|
|
basic_collection.destroy()
|
|
|
|
with pytest.raises(Exception) as exc_info:
|
|
stats = basic_collection.stats
|
|
assert ACCESS_DESTROYED_COLLECTION_ERROR_MSG in str(exc_info.value)
|
|
|
|
with pytest.raises(Exception) as exc_info:
|
|
zvec.open(path=collection_temp_dir, option=collection_option)
|
|
assert COLLECTION_PATH_NOT_EXIST_ERROR_MSG in str(exc_info.value)
|
|
|
|
def test_collection_flush(self, basic_collection: Collection):
|
|
doc = generate_doc(1, basic_collection.schema)
|
|
|
|
result = basic_collection.insert(doc)
|
|
assert bool(result)
|
|
assert result.ok()
|
|
|
|
basic_collection.flush()
|
|
|
|
fetched_docs = basic_collection.fetch(["1"])
|
|
assert "1" in fetched_docs
|
|
assert fetched_docs["1"].id == "1"
|
|
|
|
|
|
class TestIndexDDL:
|
|
@pytest.mark.parametrize("field_name", DEFAULT_SCALAR_FIELD_NAME.values())
|
|
@pytest.mark.parametrize("index_type", SUPPORT_SCALAR_INDEX_TYPES)
|
|
def test_scalar_index_operation(
|
|
self,
|
|
full_collection: Collection,
|
|
field_name: str,
|
|
index_type: IndexType,
|
|
):
|
|
# INSERT 0~5 Doc
|
|
docs = [generate_doc(i, full_collection.schema) for i in range(5)]
|
|
|
|
result = full_collection.insert(docs)
|
|
assert len(result) == 5
|
|
for item in result:
|
|
assert item.ok()
|
|
|
|
stats = full_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 5
|
|
|
|
if field_name in ["bool_field"]:
|
|
query_filter = f"{field_name} = true"
|
|
elif field_name in ["double_field", "float_field"]:
|
|
query_filter = f"{field_name} >= 3.0"
|
|
elif field_name in [
|
|
"int32_field",
|
|
"int64_field",
|
|
"uint32_field",
|
|
"uint64_field",
|
|
]:
|
|
query_filter = f"{field_name} >= 30"
|
|
elif field_name in ["string_field"]:
|
|
query_filter = f"{field_name} >= 'test_3'"
|
|
elif field_name in ["array_bool_field"]:
|
|
query_filter = f"{field_name} contain_any (false)"
|
|
elif field_name in ["array_double_field", "array_float_field"]:
|
|
query_filter = f"{field_name} contain_any (3.0, 4.0)"
|
|
elif field_name in [
|
|
"array_int64_field",
|
|
"array_int32_field",
|
|
"array_uint64_field",
|
|
"array_uint32_field",
|
|
]:
|
|
query_filter = f"{field_name} contain_any (3, 4)"
|
|
elif field_name == "array_string_field":
|
|
query_filter = f"{field_name} contain_any ('test_3', 'test_4')"
|
|
else:
|
|
assert False, f"Unsupported field type for index creation: {field_name}"
|
|
|
|
query_result_before = full_collection.query(filter=query_filter, topk=10)
|
|
|
|
if index_type not in DEFAULT_INDEX_PARAMS:
|
|
pytest.fail(f"Unsupported index type for index creation: {index_type}")
|
|
index_param = DEFAULT_INDEX_PARAMS[index_type]
|
|
|
|
full_collection.create_index(
|
|
field_name=field_name, index_param=index_param, option=IndexOption()
|
|
)
|
|
stats_after_create = full_collection.stats
|
|
assert stats_after_create is not None
|
|
assert stats_after_create.doc_count == 5
|
|
|
|
query_result_after = full_collection.query(filter=query_filter, topk=10)
|
|
|
|
assert len(query_result_before) == len(query_result_after), (
|
|
f"Query result count mismatch for {field_name} with index type {index_type}: before={len(query_result_before)}, after={len(query_result_after)}"
|
|
)
|
|
|
|
before_ids = set(doc.id for doc in query_result_before)
|
|
after_ids = set(doc.id for doc in query_result_after)
|
|
assert before_ids == after_ids, (
|
|
f"Query result IDs mismatch for {field_name} with index type {index_type}: before={before_ids}, after={after_ids}"
|
|
)
|
|
|
|
# INSERT 5~8 Doc
|
|
new_docs = [generate_doc(i, full_collection.schema) for i in range(5, 8)]
|
|
|
|
result = full_collection.insert(new_docs)
|
|
assert len(result) == 3
|
|
for item in result:
|
|
assert item.ok()
|
|
|
|
stats_after_insert1 = full_collection.stats
|
|
assert stats_after_insert1 is not None
|
|
assert stats_after_insert1.doc_count == 8
|
|
|
|
fetched_docs = full_collection.fetch([f"{i}" for i in range(5, 8)])
|
|
assert len(fetched_docs) == 3
|
|
|
|
for i in range(5, 8):
|
|
doc_id = f"{i}"
|
|
assert doc_id in fetched_docs
|
|
|
|
query_result = full_collection.query(filter=query_filter, topk=20)
|
|
assert len(query_result) >= len(query_result_before)
|
|
|
|
full_collection.drop_index(field_name=field_name)
|
|
|
|
# Insert 8~10 Doc
|
|
more_docs = [generate_doc(i, full_collection.schema) for i in range(8, 10)]
|
|
|
|
result = full_collection.insert(more_docs)
|
|
assert len(result) == 2
|
|
for item in result:
|
|
assert item.ok()
|
|
|
|
stats_after_insert2 = full_collection.stats
|
|
assert stats_after_insert2 is not None
|
|
assert stats_after_insert2.doc_count == 10
|
|
|
|
fetched_docs = full_collection.fetch([f"{i}" for i in range(8, 10)])
|
|
assert len(fetched_docs) == 2
|
|
|
|
for i in range(8, 10):
|
|
doc_id = f"{i}"
|
|
assert doc_id in fetched_docs
|
|
|
|
query_result = full_collection.query(filter=query_filter, topk=20)
|
|
assert len(query_result) >= len(query_result_before)
|
|
|
|
final_stats = full_collection.stats
|
|
assert final_stats is not None
|
|
assert final_stats.doc_count == 10
|
|
full_collection.destroy()
|
|
|
|
@pytest.mark.parametrize("field_name", DEFAULT_SCALAR_FIELD_NAME.values())
|
|
@pytest.mark.parametrize("index_type", SUPPORT_SCALAR_INDEX_TYPES)
|
|
def test_duplicate_create_index(
|
|
self, full_collection: Collection, field_name: str, index_type: IndexType
|
|
):
|
|
docs = [generate_doc(i, full_collection.schema) for i in range(10)]
|
|
|
|
result = full_collection.insert(docs)
|
|
assert bool(result)
|
|
for item in result:
|
|
assert item.ok()
|
|
|
|
stats = full_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 10
|
|
|
|
if field_name in ["bool_field"]:
|
|
query_filter = f"{field_name} = true"
|
|
elif field_name in ["double_field", "float_field"]:
|
|
query_filter = f"{field_name} >= 3.0"
|
|
elif field_name in [
|
|
"int32_field",
|
|
"int64_field",
|
|
"uint32_field",
|
|
"uint64_field",
|
|
]:
|
|
query_filter = f"{field_name} >= 30"
|
|
elif field_name in ["string_field"]:
|
|
query_filter = f"{field_name} >= 'test_3'"
|
|
elif field_name in ["array_bool_field"]:
|
|
query_filter = f"{field_name} contain_any (false)"
|
|
elif field_name in ["array_double_field", "array_float_field"]:
|
|
query_filter = f"{field_name} contain_any (3.0, 4.0)"
|
|
elif field_name in [
|
|
"array_int64_field",
|
|
"array_int32_field",
|
|
"array_uint64_field",
|
|
"array_uint32_field",
|
|
]:
|
|
query_filter = f"{field_name} contain_any (3, 4)"
|
|
elif field_name == "array_string_field":
|
|
query_filter = f"{field_name} contain_any ('test_3', 'test_4')"
|
|
else:
|
|
assert False, f"Unsupported field type for index creation: {field_name}"
|
|
|
|
query_result_before = full_collection.query(filter=query_filter, topk=5)
|
|
|
|
if index_type not in DEFAULT_INDEX_PARAMS:
|
|
pytest.fail(f"Unsupported index type for index creation: {index_type}")
|
|
index_param = DEFAULT_INDEX_PARAMS[index_type]
|
|
|
|
full_collection.create_index(
|
|
field_name=field_name, index_param=index_param, option=IndexOption()
|
|
)
|
|
|
|
query_result_after = full_collection.query(filter=query_filter, topk=5)
|
|
|
|
assert len(query_result_before) == len(query_result_after), (
|
|
f"Query result count mismatch: before={len(query_result_before)}, after={len(query_result_after)}"
|
|
)
|
|
|
|
before_ids = set(doc.id for doc in query_result_before)
|
|
after_ids = set(doc.id for doc in query_result_after)
|
|
assert before_ids == after_ids, (
|
|
f"Query result IDs mismatch: before={before_ids}, after={after_ids}"
|
|
)
|
|
|
|
full_collection.create_index(
|
|
field_name=field_name, index_param=index_param, option=IndexOption()
|
|
)
|
|
|
|
def test_optimize(self, full_collection: Collection):
|
|
docs = [generate_doc(i, full_collection.schema) for i in range(10)]
|
|
|
|
result = full_collection.insert(docs)
|
|
assert bool(result)
|
|
for item in result:
|
|
assert item.ok()
|
|
|
|
stats = full_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 10
|
|
|
|
full_collection.optimize(option=OptimizeOption())
|
|
|
|
fetched_docs = full_collection.fetch(["1"])
|
|
assert "1" in fetched_docs
|
|
assert fetched_docs["1"].id == "1"
|
|
|
|
@pytest.mark.parametrize(
|
|
"vector_type, index_type", SUPPORT_VECTOR_DATA_TYPE_INDEX_MAP_PARAMS
|
|
)
|
|
def test_vector_index_operation(
|
|
self,
|
|
full_collection: Collection,
|
|
vector_type: DataType,
|
|
index_type: IndexType,
|
|
):
|
|
vector_field_name = DEFAULT_VECTOR_FIELD_NAME[vector_type]
|
|
|
|
docs = [generate_doc(i, full_collection.schema) for i in range(5)]
|
|
|
|
result = full_collection.insert(docs)
|
|
assert len(result) == 5, (
|
|
f"Expected 5 insertion results, got {len(result)} for vector type {vector_type} and index type {index_type}"
|
|
)
|
|
for i, item in enumerate(result):
|
|
assert item.ok(), (
|
|
f"Before create_index,result={result},Insertion result {i} is not OK for vector type {vector_type} and index type {index_type} and result={result}"
|
|
)
|
|
|
|
stats = full_collection.stats
|
|
assert stats is not None, (
|
|
f"stats is None for vector type {vector_type} and index type {index_type}"
|
|
)
|
|
assert stats.doc_count == 5, (
|
|
f"doc_count!=5 for vector type {vector_type} and index type {index_type}"
|
|
)
|
|
|
|
if index_type not in DEFAULT_INDEX_PARAMS:
|
|
pytest.fail(
|
|
f"Unsupported index type {index_type} for vector type {vector_type} in test_vector_all_data_types_index_create_drop_validation"
|
|
)
|
|
index_param = DEFAULT_INDEX_PARAMS[index_type]
|
|
|
|
full_collection.create_index(
|
|
field_name=vector_field_name,
|
|
index_param=index_param,
|
|
option=IndexOption(),
|
|
)
|
|
|
|
stats_after_create = full_collection.stats
|
|
assert stats_after_create is not None, (
|
|
f"stats_after_create_index is None for vector type {vector_type} and index type {index_type}"
|
|
)
|
|
|
|
new_docs = [generate_doc(i, full_collection.schema) for i in range(5, 8)]
|
|
|
|
result = full_collection.insert(new_docs)
|
|
assert len(result) == 3, (
|
|
f"Expected 3 insertion results, got {len(result)} for vector type {vector_type} and index type {index_type}"
|
|
)
|
|
for i, item in enumerate(result):
|
|
assert item.ok(), (
|
|
f"Before drop_index,result={result},BInsertion result {i} is not OK for vector type {vector_type} and index type {index_type} and "
|
|
)
|
|
|
|
stats_after_insert1 = full_collection.stats
|
|
assert stats_after_insert1 is not None, (
|
|
f"stats_after_insert1 is None for vector type {vector_type} and index type {index_type}"
|
|
)
|
|
assert stats_after_insert1.doc_count == 8, (
|
|
f"Expected 8 documents, got {stats_after_insert1.doc_count} for vector type {vector_type} and index type {index_type}"
|
|
)
|
|
|
|
fetched_docs = full_collection.fetch([f"{i}" for i in range(5, 8)])
|
|
assert len(fetched_docs) == 3, (
|
|
f"Expected 3 fetched documents, got {len(fetched_docs)} for vector type {vector_type} and index type {index_type}"
|
|
)
|
|
|
|
for i in range(5, 8):
|
|
doc_id = f"{i}"
|
|
assert doc_id in fetched_docs, (
|
|
f"Document ID {doc_id} not found in fetched results for vector type {vector_type} and index type {index_type}"
|
|
)
|
|
assert fetched_docs[doc_id].id == doc_id, (
|
|
f"Document {doc_id} has incorrect ID field value for vector type {vector_type} and index type {index_type}"
|
|
)
|
|
|
|
full_collection.drop_index(field_name=vector_field_name)
|
|
|
|
more_docs = [generate_doc(i, full_collection.schema) for i in range(8, 10)]
|
|
result = full_collection.insert(more_docs)
|
|
assert len(result) == 2, (
|
|
f"Expected 2 insertion results, got {len(result)} for vector type {vector_type} and index type {index_type}"
|
|
)
|
|
for i, item in enumerate(result):
|
|
assert item.ok(), (
|
|
f"After drop_index,Insertion result {i} is not OK for vector type {vector_type} and index type {index_type} and result={result}"
|
|
)
|
|
|
|
# Verify document count after second insertion
|
|
stats_after_insert2 = full_collection.stats
|
|
assert stats_after_insert2 is not None, (
|
|
f"stats_after_insert2 is None for vector type {vector_type} and index type {index_type}"
|
|
)
|
|
assert stats_after_insert2.doc_count == 10, (
|
|
f"Expected 10 documents, got {stats_after_insert2.doc_count} for vector type {vector_type} and index type {index_type}"
|
|
)
|
|
|
|
# Fetch data
|
|
fetched_docs = full_collection.fetch([f"{i}" for i in range(8, 10)])
|
|
assert len(fetched_docs) == 2, (
|
|
f"Expected 2 fetched documents, got {len(fetched_docs)} for vector type {vector_type} and index type {index_type}"
|
|
)
|
|
|
|
# Verify fetched documents have correct data
|
|
for i in range(8, 10):
|
|
doc_id = f"{i}"
|
|
assert doc_id in fetched_docs, (
|
|
f"Document ID {doc_id} not found in fetched results for vector type {vector_type} and index type {index_type}"
|
|
)
|
|
assert fetched_docs[doc_id].id == doc_id, (
|
|
f"Document {doc_id} has incorrect ID field value for vector type {vector_type} and index type {index_type}"
|
|
)
|
|
|
|
# Final verification
|
|
final_stats = full_collection.stats
|
|
assert final_stats is not None, (
|
|
f"final_stats is None for vector type {vector_type} and index type {index_type}"
|
|
)
|
|
assert final_stats.doc_count == 10, (
|
|
f"Expected 10 documents, got {final_stats.doc_count} for vector type {vector_type} and index type {index_type}"
|
|
)
|
|
full_collection.destroy()
|
|
|
|
@staticmethod
|
|
def create_collection(
|
|
collection_path, collection_option: CollectionOption
|
|
) -> Collection:
|
|
schema = CollectionSchema(
|
|
name="test_collection_invalid_vector_index",
|
|
fields=[
|
|
FieldSchema(
|
|
"id",
|
|
DataType.INT64,
|
|
nullable=False,
|
|
index_param=InvertIndexParam(enable_range_optimization=True),
|
|
),
|
|
FieldSchema(
|
|
"name",
|
|
DataType.STRING,
|
|
nullable=True,
|
|
index_param=InvertIndexParam(),
|
|
),
|
|
],
|
|
vectors=[
|
|
VectorSchema(
|
|
"dense",
|
|
DataType.VECTOR_FP32,
|
|
dimension=128,
|
|
index_param=HnswIndexParam(),
|
|
),
|
|
],
|
|
)
|
|
coll = zvec.create_and_open(
|
|
path=collection_path, schema=schema, option=collection_option
|
|
)
|
|
assert coll is not None, "Failed to create and open collection"
|
|
return coll
|
|
|
|
@staticmethod
|
|
def check_error_message(exc_info, invalid_name):
|
|
if type(invalid_name) is str:
|
|
assert INDEX_NON_EXISTENT_COLUMN_ERROR_MSG in str(exc_info.value), (
|
|
"Error message is unreasonable: e=" + str(exc_info.value)
|
|
)
|
|
else:
|
|
assert INCOMPATIBLE_FUNCTION_ERROR_MSG in str(exc_info.value), (
|
|
"Error message is unreasonable: e=" + str(exc_info.value)
|
|
)
|
|
|
|
@pytest.mark.parametrize(
|
|
"invalid_field_name,invalid_vector_name",
|
|
[
|
|
("", ""), # Empty string
|
|
(" ", " "), # Space only
|
|
("v" * 33, "v" * 33), # Too long (33 characters, exceeds 32)
|
|
("vector name", "vector_name"), # Contains space
|
|
("vector@name", "vector@name"), # Contains special character
|
|
("vector/name", "vector/name"), # Contains slash
|
|
("vector\\name", "vector\\name"), # Contains backslash
|
|
("vector.name", "vector.name"), # Contains dot
|
|
("vector$data", "vector$data"), # Contains dollar sign
|
|
("vector+name", "vector+name"), # Contains plus sign
|
|
("vector=name", "vector=name"), # Contains equals sign
|
|
(None, None), # None value,
|
|
(1, 1),
|
|
(1.1, 1.1),
|
|
],
|
|
)
|
|
def test_invalid_field_and_vector_name(
|
|
self,
|
|
collection_temp_dir,
|
|
collection_option: CollectionOption,
|
|
invalid_field_name: Any,
|
|
invalid_vector_name: Any,
|
|
):
|
|
coll = self.create_collection(collection_temp_dir, collection_option)
|
|
with pytest.raises(Exception) as exc_info:
|
|
coll.create_index(
|
|
field_name=invalid_vector_name,
|
|
index_param=HnswIndexParam(),
|
|
option=IndexOption(),
|
|
)
|
|
self.check_error_message(exc_info, invalid_vector_name)
|
|
with pytest.raises(Exception) as exc_info:
|
|
coll.create_index(
|
|
field_name=invalid_field_name,
|
|
index_param=InvertIndexParam(),
|
|
option=IndexOption(),
|
|
)
|
|
self.check_error_message(exc_info, invalid_field_name)
|
|
coll.destroy()
|
|
coll = self.create_collection(collection_temp_dir, collection_option)
|
|
with pytest.raises(Exception) as exc_info:
|
|
coll.drop_index(field_name=invalid_vector_name)
|
|
self.check_error_message(exc_info, invalid_vector_name)
|
|
with pytest.raises(Exception) as exc_info:
|
|
coll.drop_index(field_name=invalid_field_name)
|
|
self.check_error_message(exc_info, invalid_field_name)
|
|
coll.destroy()
|
|
|
|
@pytest.mark.parametrize(
|
|
"field_name,vector_name",
|
|
[
|
|
("2", "3"),
|
|
("col", "co1"),
|
|
("ID", "IM"),
|
|
("name-1", "name2"),
|
|
("Weigt_12", "Weigt_13"),
|
|
("123age", "123agl"),
|
|
],
|
|
)
|
|
def test_valid_field_and_vector_name(
|
|
self,
|
|
collection_temp_dir,
|
|
collection_option: CollectionOption,
|
|
field_name: str,
|
|
vector_name: str,
|
|
):
|
|
schema = zvec.CollectionSchema(
|
|
name="test_index_names",
|
|
fields=[
|
|
FieldSchema(
|
|
"id",
|
|
DataType.INT64,
|
|
nullable=False,
|
|
index_param=InvertIndexParam(enable_range_optimization=True),
|
|
),
|
|
FieldSchema(field_name, DataType.STRING, nullable=True),
|
|
],
|
|
vectors=[
|
|
VectorSchema(
|
|
vector_name,
|
|
DataType.VECTOR_FP32,
|
|
dimension=128,
|
|
index_param=HnswIndexParam(),
|
|
)
|
|
],
|
|
)
|
|
|
|
coll = zvec.create_and_open(
|
|
path=collection_temp_dir, schema=schema, option=collection_option
|
|
)
|
|
|
|
assert coll is not None, (
|
|
f"Failed to create and open collection with field_name={field_name}, vector_name={vector_name}"
|
|
)
|
|
|
|
# Insert some data
|
|
docs = [
|
|
Doc(
|
|
id=f"{i}",
|
|
fields={"id": i, field_name: f"value_{i}"},
|
|
vectors={vector_name: [float(j % 10) for j in range(128)]},
|
|
)
|
|
for i in range(5)
|
|
]
|
|
|
|
result = coll.insert(docs)
|
|
assert len(result) == 5, (
|
|
f"Expected 5 insertion results, got {len(result)} for field_name={field_name}, vector_name={vector_name}"
|
|
)
|
|
for item in result:
|
|
assert item.ok(), (
|
|
f"Insertion failed for field_name={field_name}, vector_name={vector_name}: {item}"
|
|
)
|
|
|
|
# Create index on field
|
|
coll.create_index(
|
|
field_name=field_name,
|
|
index_param=InvertIndexParam(),
|
|
option=IndexOption(),
|
|
)
|
|
|
|
# Create index on vector
|
|
coll.create_index(
|
|
field_name=vector_name,
|
|
index_param=HnswIndexParam(),
|
|
option=IndexOption(),
|
|
)
|
|
|
|
# Verify indexes were created successfully
|
|
stats = coll.stats
|
|
assert stats is not None, (
|
|
f"Stats is None for field_name={field_name}, vector_name={vector_name}"
|
|
)
|
|
|
|
coll.destroy()
|
|
|
|
def test_compicated_workflow(
|
|
self,
|
|
collection_temp_dir,
|
|
basic_schema: CollectionSchema,
|
|
collection_option: CollectionOption,
|
|
):
|
|
"""
|
|
Test the complete workflow:
|
|
1. Create collection
|
|
2. Create index
|
|
3. Insert doc
|
|
4. Upsert
|
|
5. Update doc
|
|
6. Fetch doc
|
|
7. Query doc
|
|
8. Drop index
|
|
9. Insert doc
|
|
10. Update doc
|
|
11. Upsert doc
|
|
12. Fetch doc
|
|
13. Query doc
|
|
14. Flush
|
|
15. Destroy
|
|
"""
|
|
# Step 1: Create collection
|
|
coll = zvec.create_and_open(
|
|
path=collection_temp_dir,
|
|
schema=basic_schema,
|
|
option=collection_option,
|
|
)
|
|
|
|
assert coll is not None, "Failed to create and open collection"
|
|
assert coll.path == collection_temp_dir
|
|
assert coll.schema.name == basic_schema.name
|
|
assert coll.stats.doc_count == 0
|
|
|
|
# Step 2: Create index
|
|
coll.create_index(
|
|
field_name="name", index_param=InvertIndexParam(), option=IndexOption()
|
|
)
|
|
# Verify index was created
|
|
stats = coll.stats
|
|
assert stats is not None, "coll.stats is None!"
|
|
|
|
# Step 3: Insert doc
|
|
doc1 = Doc(
|
|
id="1",
|
|
fields={"id": 1, "name": "test1", "weight": 80.5},
|
|
vectors={
|
|
"dense": np.random.random(128).tolist(),
|
|
"sparse": {1: 1.0, 2: 2.0},
|
|
},
|
|
)
|
|
|
|
result = coll.insert(doc1)
|
|
assert bool(result)
|
|
assert result.ok()
|
|
assert coll.stats.doc_count == 1
|
|
|
|
# Step 4: Upsert (existing doc)
|
|
doc1_updated = Doc(
|
|
id="1",
|
|
fields={"id": 1, "name": "test1_updated", "weight": 85.0},
|
|
vectors={
|
|
"dense": np.random.random(128).tolist(),
|
|
"sparse": {1: 1.5, 2: 2.5},
|
|
},
|
|
)
|
|
|
|
result = coll.upsert(doc1_updated)
|
|
assert bool(result)
|
|
assert result.ok()
|
|
assert coll.stats.doc_count == 1
|
|
|
|
# Step 5: Update doc
|
|
doc2 = Doc(
|
|
id="2",
|
|
fields={"id": 2, "name": "test2", "weight": 90.0},
|
|
vectors={
|
|
"dense": np.random.random(128).tolist(),
|
|
"sparse": {1: 3.0, 2: 4.0},
|
|
},
|
|
)
|
|
|
|
# First insert doc2
|
|
result = coll.insert(doc2)
|
|
assert bool(result)
|
|
assert result.ok()
|
|
assert coll.stats.doc_count == 2
|
|
|
|
# Then update it
|
|
doc2_updated = Doc(
|
|
id="2",
|
|
fields={"id": 2, "name": "test2_updated", "weight": 95.0},
|
|
vectors={
|
|
"dense": np.random.random(128).tolist(),
|
|
"sparse": {1: 3.5, 2: 4.5},
|
|
},
|
|
)
|
|
|
|
result = coll.update(doc2_updated)
|
|
assert bool(result)
|
|
assert result.ok()
|
|
assert coll.stats.doc_count == 2
|
|
|
|
# Step 6: Fetch doc
|
|
fetched_docs = coll.fetch(["1", "2"])
|
|
assert len(fetched_docs) == 2
|
|
assert "1" in fetched_docs
|
|
assert "2" in fetched_docs
|
|
assert fetched_docs["1"].field("name") == "test1_updated"
|
|
assert fetched_docs["2"].field("name") == "test2_updated"
|
|
|
|
# Step 7: Query doc
|
|
query_result = coll.query(filter="id >= 1", topk=10)
|
|
assert len(query_result) == 2
|
|
|
|
# Step 8: Drop index
|
|
coll.drop_index(field_name="name")
|
|
|
|
# Step 9: Insert doc
|
|
doc3 = Doc(
|
|
id="3",
|
|
fields={"id": 3, "name": "test3", "weight": 100.0},
|
|
vectors={
|
|
"dense": np.random.random(128).tolist(),
|
|
"sparse": {1: 5.0, 2: 6.0},
|
|
},
|
|
)
|
|
|
|
result = coll.insert(doc3)
|
|
assert bool(result)
|
|
assert result.ok()
|
|
assert coll.stats.doc_count == 3
|
|
|
|
# Step 10: Update doc
|
|
doc3_updated = Doc(
|
|
id="3",
|
|
fields={"id": 3, "name": "test3_updated", "weight": 105.0},
|
|
vectors={
|
|
"dense": np.random.random(128).tolist(),
|
|
"sparse": {1: 5.5, 2: 6.5},
|
|
},
|
|
)
|
|
|
|
result = coll.update(doc3_updated)
|
|
assert bool(result)
|
|
assert result.ok()
|
|
assert coll.stats.doc_count == 3
|
|
|
|
# Step 11: Upsert doc
|
|
doc4 = Doc(
|
|
id="4",
|
|
fields={"id": 4, "name": "test4", "weight": 110.0},
|
|
vectors={
|
|
"dense": np.random.random(128).tolist(),
|
|
"sparse": {1: 7.0, 2: 8.0},
|
|
},
|
|
)
|
|
|
|
result = coll.upsert(doc4)
|
|
assert bool(result)
|
|
assert result.ok()
|
|
assert coll.stats.doc_count == 4
|
|
|
|
# Step 12: Fetch doc
|
|
fetched_docs = coll.fetch(["3", "4"])
|
|
assert len(fetched_docs) == 2
|
|
assert "3" in fetched_docs
|
|
assert "4" in fetched_docs
|
|
assert fetched_docs["3"].field("name") == "test3_updated"
|
|
assert fetched_docs["4"].field("name") == "test4"
|
|
|
|
# Step 13: Query doc
|
|
query_result = coll.query(filter="id >= 3", topk=10)
|
|
assert len(query_result) == 2
|
|
|
|
# Step 14: Flush
|
|
coll.flush()
|
|
|
|
# Verify data is still accessible after flush
|
|
fetched_docs = coll.fetch(["1", "2", "3", "4"])
|
|
assert len(fetched_docs) == 4
|
|
|
|
# Step 15: Destroy
|
|
coll.destroy()
|
|
|
|
@pytest.mark.parametrize(
|
|
"data_type, index_param", VALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP_PARAMS
|
|
)
|
|
def test_vector_index_params(
|
|
self,
|
|
collection_temp_dir,
|
|
collection_option: CollectionOption,
|
|
data_type: DataType,
|
|
index_param,
|
|
single_vector_schema,
|
|
):
|
|
vector_name = DEFAULT_VECTOR_FIELD_NAME[data_type]
|
|
dimension = DEFAULT_VECTOR_DIMENSION
|
|
|
|
coll = zvec.create_and_open(
|
|
path=collection_temp_dir,
|
|
schema=single_vector_schema,
|
|
option=collection_option,
|
|
)
|
|
|
|
assert coll is not None, (
|
|
f"Failed to create and open collection, {data_type}, {index_param}"
|
|
)
|
|
|
|
docs = {str(i): generate_doc(i, single_vector_schema) for i in range(5)}
|
|
result = coll.insert(docs.values())
|
|
assert len(result) == len(docs), (
|
|
f"Expected 5 results, got {len(result)}, {data_type}, {index_param}"
|
|
)
|
|
for item in result:
|
|
assert item.ok(), f"Insertion failed for, {data_type}, {index_param}"
|
|
|
|
def check_result(
|
|
label: str, metric_type: MetricType, quantize_type: QuantizeType
|
|
):
|
|
query_vector = [1] * dimension
|
|
if data_type in [DataType.SPARSE_VECTOR_FP16, DataType.SPARSE_VECTOR_FP32]:
|
|
query_vector = {1: 1}
|
|
|
|
fetch_result = coll.fetch([str(i) for i in range(len(docs))])
|
|
assert len(fetch_result) == len(docs), (
|
|
f"{label}, Expected 5 fetched docs, got {len(fetch_result)}, {data_type}, {index_param}"
|
|
)
|
|
for i in range(len(docs)):
|
|
doc_id = str(i)
|
|
assert doc_id in fetch_result, (
|
|
f"{label}, Document ID '{doc_id}' not found, {data_type}, {index_param}"
|
|
)
|
|
fetched_doc = fetch_result[doc_id]
|
|
# Verify doc equal
|
|
assert is_doc_equal(fetched_doc, docs[doc_id], single_vector_schema), (
|
|
f"{label}, doc not equal, insert: {docs[doc_id]}, fetched: {fetched_doc}, {data_type}, {index_param}"
|
|
)
|
|
|
|
query_result: list[Doc] = coll.query(
|
|
Query(field_name=vector_name, vector=query_vector),
|
|
include_vector=False,
|
|
topk=len(docs),
|
|
)
|
|
assert len(query_result) == len(docs), (
|
|
f"{label}, Expected {len(docs)} result, got {len(query_result)}, {data_type}, {index_param}"
|
|
)
|
|
inserted_ids = [str(i) for i in range(len(docs))]
|
|
queried_ids = [doc.id for doc in query_result]
|
|
assert set(inserted_ids) == set(queried_ids), (
|
|
f"{label}, inserted_ids != queried_ids, insert: {inserted_ids}, query: {queried_ids}, {data_type}, {index_param}"
|
|
)
|
|
|
|
last_score = None
|
|
for i, doc in enumerate(query_result):
|
|
# Get the document's vector for comparison
|
|
expect_doc = generate_doc(int(doc.id), single_vector_schema)
|
|
doc_vector = expect_doc.vector(vector_name)
|
|
expected_score = distance(
|
|
doc_vector,
|
|
query_vector,
|
|
metric_type,
|
|
data_type,
|
|
quantize_type,
|
|
)
|
|
print(f"query: {doc}, expect_core: {expected_score}")
|
|
if quantize_type is QuantizeType.UNDEFINED:
|
|
assert is_float_equal(doc.score, expected_score), (
|
|
f"{label} top{i} pk{doc.id} score {doc.score:6f} expected:{expected_score:6f}, {data_type}, {index_param}"
|
|
)
|
|
if last_score is not None:
|
|
if metric_type == MetricType.IP:
|
|
assert last_score >= doc.score, (
|
|
f"{label}, score not sorted, last_score: {last_score}, current_score: {doc.score}, {data_type}, {index_param}"
|
|
)
|
|
else:
|
|
assert last_score <= doc.score, (
|
|
f"{label}, score not sorted, last_score: {last_score}, current_score: {doc.score}, {data_type}, {index_param}"
|
|
)
|
|
last_score = doc.score
|
|
|
|
# default metric_type=IP, quantize_type=None
|
|
check_result("pre_create_index", MetricType.IP, QuantizeType.UNDEFINED)
|
|
|
|
# create index
|
|
coll.create_index(
|
|
field_name=vector_name,
|
|
index_param=index_param,
|
|
option=IndexOption(),
|
|
)
|
|
check_result(
|
|
"post_create_index", index_param.metric_type, index_param.quantize_type
|
|
)
|
|
|
|
coll.drop_index(field_name=vector_name)
|
|
check_result("post_drop_index", MetricType.IP, QuantizeType.UNDEFINED)
|
|
|
|
new_docs = {str(i): generate_doc(i, single_vector_schema) for i in range(5, 8)}
|
|
new_result = coll.insert(new_docs.values())
|
|
assert len(new_result) == len(new_docs), (
|
|
f"Expected {len(new_docs)} insertion results for new docs, got {len(new_result)} for vector {vector_name}"
|
|
)
|
|
for item in new_result:
|
|
assert item.ok(), (
|
|
f"New document insertion failed for vector {vector_name}: {item}"
|
|
)
|
|
docs |= new_docs
|
|
coll.create_index(
|
|
field_name=vector_name,
|
|
index_param=index_param,
|
|
option=IndexOption(),
|
|
)
|
|
|
|
check_result(
|
|
"post_create_index2", index_param.metric_type, index_param.quantize_type
|
|
)
|
|
coll.destroy()
|
|
|
|
|
|
class TestColumnDDL:
|
|
def test_add_column(self, basic_collection: Collection):
|
|
basic_collection.add_column(
|
|
field_schema=FieldSchema("income", DataType.INT32),
|
|
expression="'weight' * 2", # Simple expression
|
|
)
|
|
doc = Doc(
|
|
id="1",
|
|
fields={"id": 1, "name": "test", "weight": 80.5, "income": 1},
|
|
vectors={
|
|
"dense": np.random.random(128).tolist(),
|
|
"sparse": {1: 1.0, 2: 2.0},
|
|
},
|
|
)
|
|
|
|
result = basic_collection.insert(doc)
|
|
assert bool(result), f"Expected 1 result, but got {len(result)}"
|
|
assert result.ok(), (
|
|
f"result={result},Insert operation failed with code = {result.code()}"
|
|
)
|
|
|
|
stats = basic_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 1
|
|
|
|
def test_add_column_with_default_option(self, basic_collection: Collection):
|
|
# Add a new column with default option
|
|
basic_collection.add_column(
|
|
field_schema=FieldSchema("test_column_default", DataType.INT32),
|
|
expression="100",
|
|
option=AddColumnOption(), # Default option
|
|
)
|
|
# Verify column was added by inserting data
|
|
doc = Doc(
|
|
id="1",
|
|
fields={"id": 1, "name": "test", "weight": 80.5, "test_column_default": 1},
|
|
vectors={
|
|
"dense": np.random.random(128).tolist(),
|
|
"sparse": {1: 1.0, 2: 2.0},
|
|
},
|
|
)
|
|
|
|
result = basic_collection.insert(doc)
|
|
assert bool(result), f"Expected 1 result, but got {len(result)}"
|
|
assert result.ok(), (
|
|
f"result={result},Insert operation failed with code = {result.code()}"
|
|
)
|
|
# Verify document was inserted
|
|
stats = basic_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 1
|
|
|
|
@pytest.mark.parametrize("concurrency", [0, 1, 4, 8])
|
|
def test_add_column_with_various_concurrency_options(
|
|
self, basic_collection: Collection, concurrency
|
|
):
|
|
field_name = f"test_column_concurrent_{concurrency}"
|
|
basic_collection.add_column(
|
|
field_schema=FieldSchema(field_name, DataType.INT32),
|
|
expression="100",
|
|
option=AddColumnOption(concurrency=concurrency),
|
|
)
|
|
|
|
doc = Doc(
|
|
id="1",
|
|
fields={"id": 1, "name": "test", "weight": 80.5, field_name: 200},
|
|
vectors={
|
|
"dense": np.random.random(128).tolist(),
|
|
"sparse": {1: 1.0, 2: 2.0},
|
|
},
|
|
)
|
|
result = basic_collection.insert(doc)
|
|
assert bool(result), f"Expected 1 result, but got {len(result)}"
|
|
assert result.ok(), (
|
|
f"result={result},Insert operation failed with code = {result.code()}"
|
|
)
|
|
|
|
stats = basic_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 1
|
|
|
|
@pytest.mark.parametrize("data_type", SUPPORT_ADD_COLUMN_DATA_TYPE)
|
|
def test_add_column_valid_data_types(self, basic_collection: Collection, data_type):
|
|
field_name = f"test_field_{data_type.name.lower()}"
|
|
|
|
# Add a new column with specific data type
|
|
basic_collection.add_column(
|
|
field_schema=FieldSchema(field_name, data_type),
|
|
expression="1" if data_type != DataType.STRING else "'test'",
|
|
)
|
|
|
|
# Verify column was added by inserting data
|
|
if data_type == DataType.STRING:
|
|
field_value = "test_value"
|
|
elif data_type in [DataType.ARRAY_STRING]:
|
|
field_value = ["test_value"]
|
|
elif data_type in [DataType.ARRAY_INT32, DataType.ARRAY_INT64]:
|
|
field_value = [1, 2, 3]
|
|
elif data_type in [DataType.ARRAY_FLOAT, DataType.ARRAY_DOUBLE]:
|
|
field_value = [1.1, 2.2, 3.3]
|
|
elif data_type == DataType.ARRAY_BOOL:
|
|
field_value = [True, False]
|
|
elif data_type in [DataType.FLOAT, DataType.DOUBLE]:
|
|
field_value = 1.5
|
|
elif data_type in [DataType.INT32, DataType.INT64]:
|
|
field_value = 100
|
|
elif data_type == DataType.BOOL:
|
|
field_value = True
|
|
else:
|
|
field_value = 1
|
|
|
|
doc = Doc(
|
|
id="1",
|
|
fields={
|
|
"id": 1,
|
|
"name": "test",
|
|
"weight": 80.5,
|
|
field_name: field_value,
|
|
},
|
|
vectors={
|
|
"dense": np.random.random(128).tolist(),
|
|
"sparse": {1: 1.0, 2: 2.0},
|
|
},
|
|
)
|
|
result = basic_collection.insert(doc)
|
|
assert bool(result), f"Expected 1 result, but got {len(result)}"
|
|
assert result.ok(), (
|
|
f"result={result},Insert operation failed with code = {result.code()}"
|
|
)
|
|
|
|
# Verify document was inserted
|
|
stats = basic_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 1
|
|
|
|
@pytest.mark.parametrize("data_type", NOT_SUPPORT_ADD_COLUMN_DATA_TYPE)
|
|
def test_add_column_invalid_data_types(
|
|
self, basic_collection: Collection, data_type
|
|
):
|
|
with pytest.raises(Exception) as exc_info:
|
|
field_name = f"test_field_{data_type.name.lower()}"
|
|
|
|
# Add a new column with specific data type
|
|
basic_collection.add_column(
|
|
field_schema=FieldSchema(field_name, data_type),
|
|
expression="1" if data_type != DataType.STRING else "'test'",
|
|
)
|
|
|
|
assert NOT_SUPPORT_ADD_COLUMN_ERROR_MSG in str(exc_info.value)
|
|
|
|
@pytest.mark.parametrize("nullable", [True, False])
|
|
def test_add_column_with_nullable_options(
|
|
self, basic_collection: Collection, nullable
|
|
):
|
|
field_name = f"test_field_nullable_{str(nullable).lower()}"
|
|
|
|
# Add a new column with specific nullable option
|
|
basic_collection.add_column(
|
|
field_schema=FieldSchema(field_name, DataType.INT32, nullable=nullable),
|
|
expression="100",
|
|
)
|
|
|
|
# Verify column was added by inserting data
|
|
doc = Doc(
|
|
id="1",
|
|
fields={"id": 1, "name": "test", "weight": 80.5, field_name: 200},
|
|
vectors={
|
|
"dense": np.random.random(128).tolist(),
|
|
"sparse": {1: 1.0, 2: 2.0},
|
|
},
|
|
)
|
|
result = basic_collection.insert(doc)
|
|
assert bool(result), f"Expected 1 result, but got {len(result)}"
|
|
assert result.ok(), (
|
|
f"result={result},Insert operation failed with code = {result.code()}"
|
|
)
|
|
|
|
# Verify document was inserted
|
|
stats = basic_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 1
|
|
|
|
# Verify column was added by inserting data
|
|
doc = Doc(
|
|
id="2",
|
|
fields={"id": 2, "name": "test", "weight": 80.5, field_name: None},
|
|
vectors={
|
|
"dense": np.random.random(128).tolist(),
|
|
"sparse": {1: 1.0, 2: 2.0},
|
|
},
|
|
)
|
|
if nullable:
|
|
result = basic_collection.insert(doc)
|
|
assert bool(result), f"Expected 1 result, but got {len(result)}"
|
|
assert result.ok(), (
|
|
f"result={result},Insert operation failed with code = {result.code()}"
|
|
)
|
|
else:
|
|
with pytest.raises(ValueError) as e:
|
|
basic_collection.insert(doc)
|
|
assert (
|
|
"Field 'test_field_nullable_false': expected non-nullable type"
|
|
in str(e.value)
|
|
)
|
|
|
|
# Verify document was inserted
|
|
stats = basic_collection.stats
|
|
assert stats is not None
|
|
if nullable:
|
|
assert stats.doc_count == 2
|
|
else:
|
|
assert stats.doc_count == 1
|
|
|
|
@pytest.mark.parametrize(
|
|
"expression",
|
|
[
|
|
"1", # Constant integer
|
|
"1.5", # Constant float
|
|
"'test'", # Constant string
|
|
"id", # Reference to existing field
|
|
"weight * 2", # Simple arithmetic
|
|
"weight + id", # Complex arithmetic
|
|
"CASE WHEN weight > 50 THEN 1 ELSE 0 END", # Conditional expression
|
|
],
|
|
)
|
|
def test_add_column_with_different_expressions(
|
|
self, basic_collection: Collection, expression
|
|
):
|
|
field_name = f"test_field_expr_{abs(hash(expression)) % 1000}"
|
|
|
|
# Add a new column with specific expression
|
|
basic_collection.add_column(
|
|
field_schema=FieldSchema(field_name, DataType.INT32),
|
|
expression=expression,
|
|
)
|
|
|
|
# Verify column was added by inserting data
|
|
doc = Doc(
|
|
id="1",
|
|
fields={"id": 1, "name": "test", "weight": 80.5, field_name: 200},
|
|
vectors={
|
|
"dense": np.random.random(128).tolist(),
|
|
"sparse": {1: 1.0, 2: 2.0},
|
|
},
|
|
)
|
|
result = basic_collection.insert(doc)
|
|
assert bool(result), f"Expected 1 result, but got {len(result)}"
|
|
assert result.ok(), (
|
|
f"result={result},Insert operation failed with code = {result.code()}"
|
|
)
|
|
|
|
# Verify document was inserted
|
|
stats = basic_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 1
|
|
|
|
def test_add_column_with_index_param(self, basic_collection: Collection):
|
|
basic_collection.add_column(
|
|
field_schema=FieldSchema(
|
|
"indexed_field",
|
|
DataType.INT32,
|
|
index_param=InvertIndexParam(enable_range_optimization=True),
|
|
),
|
|
expression="id * 2",
|
|
)
|
|
|
|
# Verify column was added by inserting data
|
|
doc = Doc(
|
|
id="1",
|
|
fields={"id": 1, "name": "test", "weight": 80.5, "indexed_field": 200},
|
|
vectors={
|
|
"dense": np.random.random(128).tolist(),
|
|
"sparse": {1: 1.0, 2: 2.0},
|
|
},
|
|
)
|
|
result = basic_collection.insert(doc)
|
|
assert bool(result), f"Expected 1 result, but got {len(result)}"
|
|
assert result.ok(), (
|
|
f"result={result},Insert operation failed with code = {result.code()}"
|
|
)
|
|
|
|
# Verify document was inserted
|
|
stats = basic_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 1
|
|
|
|
@pytest.mark.parametrize(
|
|
"field_name",
|
|
[
|
|
"a", # Minimum length
|
|
"a" * 32, # Maximum length (32 characters)
|
|
"valid_field_name_123", # Alphanumeric with underscore
|
|
"Valid-Field-Name", # With hyphens
|
|
"_underscore_start", # Starting with underscore
|
|
"field_name_with_123_numbers", # Numbers in middle
|
|
"FIELD_NAME_UPPERCASE", # Uppercase
|
|
# "field_with_nums_123_and_hyphens-456", # Complex valid name within limit
|
|
],
|
|
)
|
|
def test_add_column_with_valid_field_names(
|
|
self, basic_collection: Collection, field_name
|
|
):
|
|
basic_collection.add_column(
|
|
field_schema=FieldSchema(field_name, DataType.INT32), expression="200"
|
|
)
|
|
|
|
doc = Doc(
|
|
id="1",
|
|
fields={"id": 1, "name": "test", "weight": 80.5, field_name: 300},
|
|
vectors={
|
|
"dense": np.random.random(128).tolist(),
|
|
"sparse": {1: 1.0, 2: 2.0},
|
|
},
|
|
)
|
|
result = basic_collection.insert(doc)
|
|
assert bool(result), f"Expected 1 result, but got {len(result)}"
|
|
assert result.ok(), (
|
|
f"result={result},Insert operation failed with code = {result.code()}"
|
|
)
|
|
|
|
stats = basic_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 1
|
|
|
|
@pytest.mark.parametrize(
|
|
"invalid_field_name",
|
|
[
|
|
"", # Empty string
|
|
" ", # Space only
|
|
"a" * 33, # Too long (33 characters, exceeds 32)
|
|
"field name", # Contains space
|
|
"field.name", # Contains dot
|
|
"field@name", # Contains special character
|
|
"field/name", # Contains slash
|
|
"field\\name", # Contains backslash
|
|
"field$name", # Contains dollar sign
|
|
"field+name", # Contains plus sign
|
|
"field=name", # Contains equals sign
|
|
None, # None value
|
|
],
|
|
)
|
|
def test_add_column_with_invalid_field_names(
|
|
self, basic_collection: Collection, invalid_field_name
|
|
):
|
|
with pytest.raises(Exception) as exc_info:
|
|
basic_collection.add_column(
|
|
field_schema=FieldSchema(invalid_field_name, DataType.INT32),
|
|
expression="100",
|
|
)
|
|
|
|
if invalid_field_name is None:
|
|
assert "validate failed" in str(exc_info.value), (
|
|
"Error message is unreasonable: e=" + str(exc_info.value)
|
|
)
|
|
else:
|
|
assert (
|
|
"invalid" in str(exc_info.value).lower()
|
|
or "name" in str(exc_info.value).lower()
|
|
)
|
|
|
|
def test_alter_column_rename(self, basic_collection: Collection):
|
|
basic_collection.alter_column(
|
|
old_name="weight",
|
|
new_name="mass",
|
|
option=AlterColumnOption(),
|
|
)
|
|
doc = Doc(
|
|
id="1",
|
|
fields={"id": 1, "name": "test", "mass": 80.5}, # Use new name
|
|
vectors={
|
|
"dense": np.random.random(128).tolist(),
|
|
"sparse": {1: 1.0, 2: 2.0},
|
|
},
|
|
)
|
|
result = basic_collection.insert(doc)
|
|
assert bool(result), f"Expected 1 result, but got {len(result)}"
|
|
assert result.ok(), (
|
|
f"result={result},Insert operation failed with code = {result.code()}"
|
|
)
|
|
|
|
stats = basic_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 1
|
|
|
|
def test_alter_column_non_exist(self, basic_collection: Collection):
|
|
with pytest.raises(Exception) as exc_info:
|
|
basic_collection.alter_column(
|
|
old_name="non_existing",
|
|
new_name="new_name",
|
|
field_schema=FieldSchema("new_name", DataType.STRING),
|
|
)
|
|
assert "column non_existing not found" in str(exc_info.value), (
|
|
"Error message is unreasonable: e=" + str(exc_info.value)
|
|
)
|
|
|
|
def test_alter_column_with_default_option(self, basic_collection: Collection):
|
|
basic_collection.add_column(
|
|
field_schema=FieldSchema("original_field", DataType.INT32), expression="100"
|
|
)
|
|
|
|
basic_collection.alter_column(
|
|
old_name="original_field",
|
|
new_name="renamed_field",
|
|
option=AlterColumnOption(),
|
|
)
|
|
|
|
doc = Doc(
|
|
id="1",
|
|
fields={"id": 1, "name": "test", "weight": 80.5, "renamed_field": 200},
|
|
vectors={
|
|
"dense": np.random.random(128).tolist(),
|
|
"sparse": {1: 1.0, 2: 2.0},
|
|
},
|
|
)
|
|
result = basic_collection.insert(doc)
|
|
assert bool(result), f"Expected 1 result, but got {len(result)}"
|
|
assert result.ok(), (
|
|
f"result={result},Insert operation failed with code = {result.code()}"
|
|
)
|
|
|
|
stats = basic_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 1
|
|
|
|
@pytest.mark.parametrize("concurrency", [0, 1, 4, 8])
|
|
def test_alter_column_with_various_concurrency_options(
|
|
self, basic_collection: Collection, concurrency
|
|
):
|
|
old_field_name = f"orig_field_{concurrency}"
|
|
new_field_name = f"modified_field_{concurrency}"
|
|
|
|
basic_collection.add_column(
|
|
field_schema=FieldSchema(old_field_name, DataType.INT32),
|
|
expression="100",
|
|
)
|
|
|
|
basic_collection.alter_column(
|
|
old_name=old_field_name,
|
|
new_name=new_field_name,
|
|
option=AlterColumnOption(concurrency=concurrency),
|
|
)
|
|
|
|
doc = Doc(
|
|
id="1",
|
|
fields={"id": 1, "name": "test", "weight": 80.5, new_field_name: 200},
|
|
vectors={
|
|
"dense": np.random.random(128).tolist(),
|
|
"sparse": {1: 1.0, 2: 2.0},
|
|
},
|
|
)
|
|
|
|
result = basic_collection.insert(doc)
|
|
assert bool(result), f"Expected 1 result, but got {len(result)}"
|
|
assert result.ok(), (
|
|
f"result={result},Insert operation failed with code = {result.code()}"
|
|
)
|
|
|
|
stats = basic_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 1
|
|
|
|
@pytest.mark.parametrize(
|
|
"old_field_name,new_field_name",
|
|
[
|
|
("a", "new_a"), # Minimum length
|
|
(
|
|
"abcdefghijklmnopqrstuvwxyz123456",
|
|
"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
|
|
), # Maximum length (32 characters)
|
|
("valid_field_name_123", "new_valid_field"), # Alphanumeric with underscore
|
|
("Valid-Field-Name", "New-Field-Name"), # With hyphens
|
|
("_underscore_start", "new_underscore"), # Starting with underscore
|
|
("field_name_with_123_numbers", "new_with_nums"), # Numbers in middle
|
|
("FIELD_NAME_UPPERCASE", "new_uppercase"), # Uppercase
|
|
(
|
|
"field_with_nums_3_and_hyphens-6",
|
|
"new_field_hyphens",
|
|
), # Complex valid name
|
|
],
|
|
)
|
|
def test_alter_column_field_name_valid(
|
|
self, basic_collection: Collection, old_field_name, new_field_name
|
|
):
|
|
basic_collection.add_column(
|
|
field_schema=FieldSchema(old_field_name, DataType.INT32),
|
|
expression="100",
|
|
)
|
|
basic_collection.alter_column(
|
|
old_name=old_field_name,
|
|
new_name=new_field_name,
|
|
option=AlterColumnOption(),
|
|
)
|
|
doc = Doc(
|
|
id="1",
|
|
fields={"id": 1, "name": "test", "weight": 80.5, new_field_name: 200},
|
|
vectors={
|
|
"dense": np.random.random(128).tolist(),
|
|
"sparse": {1: 1.0, 2: 2.0},
|
|
},
|
|
)
|
|
|
|
result = basic_collection.insert(doc)
|
|
assert bool(result), f"Expected 1 result, but got {len(result)}"
|
|
assert result.ok(), (
|
|
f"result={result},Insert operation failed with code = {result.code()}"
|
|
)
|
|
|
|
stats = basic_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 1
|
|
|
|
@pytest.mark.parametrize(
|
|
"valid_old_name,invalid_new_name",
|
|
[
|
|
("temp_field", ""), # Empty new name
|
|
("temp_field", "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"), # Too long new name
|
|
("temp_field", "field name"), # New name with space
|
|
("temp_field", "field.name"), # New name with dot
|
|
("temp_field", "field@name"), # New name with special character
|
|
("temp_field", "field/name"), # New name with slash
|
|
("temp_field", "field\\name"), # New name with backslash
|
|
("temp_field", "field$name"), # New name with dollar sign
|
|
("temp_field", "field+name"), # New name with plus sign
|
|
("temp_field", "field=name"), # New name with equals sign
|
|
("temp_field", None), # None new name
|
|
],
|
|
)
|
|
def test_alter_column_with_invalid_field_names(
|
|
self, basic_collection: Collection, valid_old_name, invalid_new_name
|
|
):
|
|
basic_collection.add_column(
|
|
field_schema=FieldSchema("temp_field", DataType.INT32), expression="100"
|
|
)
|
|
with pytest.raises(Exception) as exc_info:
|
|
basic_collection.alter_column(
|
|
old_name=valid_old_name,
|
|
new_name=invalid_new_name if invalid_new_name is not None else "",
|
|
field_schema=FieldSchema(
|
|
invalid_new_name if invalid_new_name is not None else "",
|
|
DataType.INT32,
|
|
),
|
|
)
|
|
|
|
assert (
|
|
"invalid" in str(exc_info.value).lower()
|
|
or "name" in str(exc_info.value).lower()
|
|
or "incompatible" in str(exc_info.value).lower()
|
|
)
|
|
|
|
def test_drop_column_exist(self, basic_collection: Collection):
|
|
basic_collection.add_column(
|
|
field_schema=FieldSchema("temp_field", DataType.INT32), expression="100"
|
|
)
|
|
doc = Doc(
|
|
id="1",
|
|
fields={"id": 1, "name": "test", "weight": 80.5, "temp_field": 1},
|
|
vectors={
|
|
"dense": np.random.random(128).tolist(),
|
|
"sparse": {1: 1.0, 2: 2.0},
|
|
},
|
|
)
|
|
|
|
result = basic_collection.insert(doc)
|
|
assert bool(result), f"Expected 1 result, but got {len(result)}"
|
|
assert result.ok(), (
|
|
f"result={result},Insert operation failed with code = {result.code()}"
|
|
)
|
|
|
|
stats = basic_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 1
|
|
|
|
basic_collection.drop_column("temp_field")
|
|
doc = Doc(
|
|
id="2",
|
|
fields={"id": 2, "name": "test", "weight": 80.5, "temp_field": 1},
|
|
vectors={
|
|
"dense": np.random.random(128).tolist(),
|
|
"sparse": {1: 1.0, 2: 2.0},
|
|
},
|
|
)
|
|
with pytest.raises(Exception) as exc_info:
|
|
result = basic_collection.insert(doc)
|
|
|
|
assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value)
|
|
|
|
def test_drop_column_non_exist(self, basic_collection: Collection):
|
|
with pytest.raises(Exception) as exc_info:
|
|
basic_collection.drop_column("non_existing_column")
|
|
assert NOT_EXIST_COLUMN_TO_DROP_ERROR_MSG in str(exc_info.value)
|
|
|
|
@pytest.mark.parametrize(
|
|
"field_name",
|
|
[
|
|
"a", # Minimum length
|
|
"a" * 32, # Maximum length (32 characters)
|
|
"valid_field_name_123", # Alphanumeric with underscore
|
|
"Valid-Field-Name", # With hyphens
|
|
"_underscore_start", # Starting with underscore
|
|
"field_name_with_123_numbers", # Numbers in middle
|
|
"FIELD_NAME_UPPERCASE", # Uppercase
|
|
"field_with_nums_3_and_hyphens-6", # Complex valid name within limit
|
|
],
|
|
)
|
|
def test_drop_column_field_name_valid(
|
|
self, basic_collection: Collection, field_name
|
|
):
|
|
basic_collection.add_column(
|
|
field_schema=FieldSchema(field_name, DataType.INT32), expression="100"
|
|
)
|
|
doc = Doc(
|
|
id="1",
|
|
fields={"id": 1, "name": "test", "weight": 80.5, field_name: 200},
|
|
vectors={
|
|
"dense": np.random.random(128).tolist(),
|
|
"sparse": {1: 1.0, 2: 2.0},
|
|
},
|
|
)
|
|
|
|
result = basic_collection.insert(doc)
|
|
assert bool(result), f"Expected 1 result, but got {len(result)}"
|
|
assert result.ok(), (
|
|
f"result={result},Insert operation failed with code = {result.code()}"
|
|
)
|
|
|
|
stats = basic_collection.stats
|
|
assert stats is not None
|
|
assert stats.doc_count == 1
|
|
|
|
basic_collection.drop_column(field_name)
|
|
|
|
doc = Doc(
|
|
id="2",
|
|
fields={"id": 2, "name": "test", "weight": 80.5, field_name: 200},
|
|
vectors={
|
|
"dense": np.random.random(128).tolist(),
|
|
"sparse": {1: 1.0, 2: 2.0},
|
|
},
|
|
)
|
|
with pytest.raises(Exception) as exc_info:
|
|
result = basic_collection.insert(doc)
|
|
|
|
assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value)
|
|
|
|
def test_add_column_then_query_returns_new_field(
|
|
self, basic_collection: Collection
|
|
):
|
|
"""Regression test for issue #426: query() should return fields added via add_column()."""
|
|
basic_collection.add_column(
|
|
field_schema=FieldSchema("score", DataType.INT64, nullable=True),
|
|
)
|
|
|
|
docs = [
|
|
Doc(
|
|
id="1",
|
|
fields={"id": 1, "name": "alice", "weight": 60.0, "score": 100},
|
|
vectors={
|
|
"dense": generate_constant_vector(1, 128),
|
|
"sparse": generate_sparse_vector(1),
|
|
},
|
|
),
|
|
Doc(
|
|
id="2",
|
|
fields={"id": 2, "name": "bob", "weight": 70.0, "score": 200},
|
|
vectors={
|
|
"dense": generate_constant_vector(2, 128),
|
|
"sparse": generate_sparse_vector(2),
|
|
},
|
|
),
|
|
]
|
|
result = basic_collection.insert(docs)
|
|
assert all(r.ok() for r in result)
|
|
|
|
# Query with explicit output_fields
|
|
query_result = basic_collection.query(
|
|
Query(field_name="dense", vector=generate_constant_vector(1, 128)),
|
|
topk=2,
|
|
output_fields=["score"],
|
|
)
|
|
assert len(query_result) == 2
|
|
for doc in query_result:
|
|
assert "score" in doc.fields, (
|
|
f"Doc {doc.id} missing 'score' field after add_column (output_fields explicit)"
|
|
)
|
|
assert doc.fields["score"] in (100, 200)
|
|
|
|
# Query with select-all (no output_fields)
|
|
query_result_all = basic_collection.query(
|
|
Query(field_name="dense", vector=generate_constant_vector(1, 128)),
|
|
topk=2,
|
|
)
|
|
assert len(query_result_all) == 2
|
|
for doc in query_result_all:
|
|
assert "score" in doc.fields, (
|
|
f"Doc {doc.id} missing 'score' field after add_column (select all)"
|
|
)
|
|
|
|
def test_alter_column_rename_then_query_returns_new_name(
|
|
self, basic_collection: Collection
|
|
):
|
|
"""Regression test: query() should use the new field name after alter_column rename."""
|
|
docs = [
|
|
Doc(
|
|
id="1",
|
|
fields={"id": 1, "name": "alice", "weight": 60.0},
|
|
vectors={
|
|
"dense": generate_constant_vector(1, 128),
|
|
"sparse": generate_sparse_vector(1),
|
|
},
|
|
),
|
|
Doc(
|
|
id="2",
|
|
fields={"id": 2, "name": "bob", "weight": 70.0},
|
|
vectors={
|
|
"dense": generate_constant_vector(2, 128),
|
|
"sparse": generate_sparse_vector(2),
|
|
},
|
|
),
|
|
]
|
|
result = basic_collection.insert(docs)
|
|
assert all(r.ok() for r in result)
|
|
|
|
# Rename 'weight' -> 'mass'
|
|
basic_collection.alter_column("weight", new_name="mass")
|
|
|
|
# Query with explicit output_fields using new name
|
|
query_result = basic_collection.query(
|
|
Query(field_name="dense", vector=generate_constant_vector(1, 128)),
|
|
topk=2,
|
|
output_fields=["mass"],
|
|
)
|
|
assert len(query_result) == 2
|
|
for doc in query_result:
|
|
assert "mass" in doc.fields, (
|
|
f"Doc {doc.id} missing 'mass' field after alter_column rename"
|
|
)
|
|
assert "weight" not in doc.fields, (
|
|
f"Doc {doc.id} still has old name 'weight' after rename"
|
|
)
|
|
|
|
# Query with select-all
|
|
query_result_all = basic_collection.query(
|
|
Query(field_name="dense", vector=generate_constant_vector(1, 128)),
|
|
topk=2,
|
|
)
|
|
assert len(query_result_all) == 2
|
|
for doc in query_result_all:
|
|
assert "mass" in doc.fields, (
|
|
f"Doc {doc.id} missing 'mass' in select-all after alter_column rename"
|
|
)
|
|
assert "weight" not in doc.fields, (
|
|
f"Doc {doc.id} still has old name 'weight' in select-all after rename"
|
|
)
|
|
|
|
def test_drop_column_then_query_excludes_dropped_field(
|
|
self, basic_collection: Collection
|
|
):
|
|
"""Regression test: query() should not return fields removed via drop_column()."""
|
|
basic_collection.add_column(
|
|
field_schema=FieldSchema("score", DataType.INT64, nullable=True),
|
|
)
|
|
|
|
docs = [
|
|
Doc(
|
|
id="1",
|
|
fields={"id": 1, "name": "alice", "weight": 60.0, "score": 100},
|
|
vectors={
|
|
"dense": generate_constant_vector(1, 128),
|
|
"sparse": generate_sparse_vector(1),
|
|
},
|
|
),
|
|
Doc(
|
|
id="2",
|
|
fields={"id": 2, "name": "bob", "weight": 70.0, "score": 200},
|
|
vectors={
|
|
"dense": generate_constant_vector(2, 128),
|
|
"sparse": generate_sparse_vector(2),
|
|
},
|
|
),
|
|
]
|
|
result = basic_collection.insert(docs)
|
|
assert all(r.ok() for r in result)
|
|
|
|
# Verify field exists before drop
|
|
query_before = basic_collection.query(
|
|
Query(field_name="dense", vector=generate_constant_vector(1, 128)),
|
|
topk=2,
|
|
)
|
|
assert all("score" in doc.fields for doc in query_before)
|
|
|
|
# Drop the column
|
|
basic_collection.drop_column("score")
|
|
|
|
# Query after drop - 'score' should not appear
|
|
query_after = basic_collection.query(
|
|
Query(field_name="dense", vector=generate_constant_vector(1, 128)),
|
|
topk=2,
|
|
)
|
|
assert len(query_after) == 2
|
|
for doc in query_after:
|
|
assert "score" not in doc.fields, (
|
|
f"Doc {doc.id} still has 'score' after drop_column"
|
|
)
|
|
assert "name" in doc.fields, (
|
|
f"Doc {doc.id} missing 'name' - other fields should still be present"
|
|
)
|