mirror of
https://github.com/qdrant/qdrant-client.git
synced 2026-07-23 11:11:01 -05:00
519 lines
16 KiB
Python
519 lines
16 KiB
Python
from copy import deepcopy
|
|
|
|
import pytest
|
|
|
|
import qdrant_client.http.exceptions as qdrant_exceptions
|
|
from qdrant_client import QdrantClient, models
|
|
from tests.congruence_tests.test_common import (
|
|
compare_collections,
|
|
generate_fixtures,
|
|
generate_sparse_fixtures,
|
|
sparse_vectors_config,
|
|
generate_multivector_fixtures,
|
|
multi_vector_config,
|
|
initialize_fixture_collection,
|
|
)
|
|
from tests.utils import read_version
|
|
|
|
|
|
VECTOR_NUMBER = 1000
|
|
|
|
|
|
@pytest.fixture
|
|
def local_client() -> QdrantClient:
|
|
client = QdrantClient(":memory:")
|
|
delete_collections(client)
|
|
yield client
|
|
delete_collections(client)
|
|
client.close()
|
|
|
|
|
|
second_local_client = deepcopy(local_client)
|
|
|
|
|
|
@pytest.fixture
|
|
def remote_client() -> QdrantClient:
|
|
client = QdrantClient()
|
|
delete_collections(client)
|
|
yield client
|
|
delete_collections(client)
|
|
client.close()
|
|
|
|
|
|
def delete_collections(client: QdrantClient) -> None:
|
|
collection_names = [collection.name for collection in client.get_collections().collections]
|
|
for collection_name in collection_names:
|
|
client.delete_collection(collection_name)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"source_client,dest_client",
|
|
[
|
|
("local_client", "remote_client"),
|
|
("remote_client", "local_client"),
|
|
("local_client", "second_local_client"),
|
|
],
|
|
)
|
|
def test_single_vector_collection(source_client, dest_client, request) -> None:
|
|
"""
|
|
Args:
|
|
source_client: fixture
|
|
dest_client: fixture
|
|
request: pytest internal object to get launch fixtures from parametrize
|
|
"""
|
|
source_client: QdrantClient = request.getfixturevalue(source_client)
|
|
dest_client: QdrantClient = request.getfixturevalue(dest_client)
|
|
vectors_config = models.VectorParams(size=10, distance=models.Distance.COSINE)
|
|
collection_name = "single_vector_collection"
|
|
initialize_fixture_collection(
|
|
source_client, collection_name=collection_name, vectors_config=vectors_config
|
|
)
|
|
dense_points = generate_fixtures(VECTOR_NUMBER, vectors_sizes=vectors_config.size)
|
|
source_client.upload_points(collection_name, dense_points, wait=True)
|
|
source_client.migrate(dest_client)
|
|
compare_collections(
|
|
source_client,
|
|
dest_client,
|
|
num_vectors=VECTOR_NUMBER,
|
|
collection_name=collection_name,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"source_client,dest_client",
|
|
[
|
|
("local_client", "remote_client"),
|
|
("remote_client", "local_client"),
|
|
("local_client", "second_local_client"),
|
|
],
|
|
)
|
|
def test_multiple_vectors_collection(source_client, dest_client, request) -> None:
|
|
"""
|
|
Args:
|
|
source_client: fixture
|
|
dest_client: fixture
|
|
request: pytest internal object to get launch fixtures from parametrize
|
|
"""
|
|
major, minor, patch, dev = read_version()
|
|
if not dev and None not in (major, minor, patch) and (major, minor, patch) < (1, 10, 0):
|
|
pytest.skip("Multivectors are supported as of v1.10.0")
|
|
|
|
source_client: QdrantClient = request.getfixturevalue(source_client)
|
|
dest_client: QdrantClient = request.getfixturevalue(dest_client)
|
|
collection_name = "multiple_vectors_collection"
|
|
initialize_fixture_collection(source_client, collection_name="multiple_vectors_collection")
|
|
points = generate_fixtures(VECTOR_NUMBER)
|
|
|
|
source_client.upload_points(collection_name, points, wait=True)
|
|
source_client.migrate(dest_client)
|
|
compare_collections(
|
|
source_client,
|
|
dest_client,
|
|
num_vectors=VECTOR_NUMBER,
|
|
collection_name=collection_name,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"source_client,dest_client",
|
|
[
|
|
("local_client", "remote_client"),
|
|
("remote_client", "local_client"),
|
|
("local_client", "second_local_client"),
|
|
],
|
|
)
|
|
def test_single_multivector_collection(source_client, dest_client, request) -> None:
|
|
"""
|
|
Args:
|
|
source_client: fixture
|
|
dest_client: fixture
|
|
request: pytest internal object to get launch fixtures from parametrize
|
|
"""
|
|
source_client: QdrantClient = request.getfixturevalue(source_client)
|
|
dest_client: QdrantClient = request.getfixturevalue(dest_client)
|
|
|
|
vector_size = 10
|
|
|
|
vectors_config = models.VectorParams(
|
|
size=vector_size,
|
|
distance=models.Distance.COSINE,
|
|
multivector_config=models.MultiVectorConfig(
|
|
comparator=models.MultiVectorComparator.MAX_SIM,
|
|
),
|
|
)
|
|
collection_name = "single_multivector_collection"
|
|
initialize_fixture_collection(
|
|
source_client, collection_name=collection_name, vectors_config=vectors_config
|
|
)
|
|
multi_vector_points = generate_multivector_fixtures(vectors_sizes=vector_size)
|
|
source_client.upload_points(collection_name, multi_vector_points, wait=True)
|
|
source_client.migrate(dest_client)
|
|
compare_collections(
|
|
source_client,
|
|
dest_client,
|
|
num_vectors=VECTOR_NUMBER,
|
|
collection_name=collection_name,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"source_client,dest_client",
|
|
[
|
|
("local_client", "remote_client"),
|
|
("remote_client", "local_client"),
|
|
("local_client", "second_local_client"),
|
|
],
|
|
)
|
|
def test_multivectors_collection(source_client, dest_client, request) -> None:
|
|
"""
|
|
Args:
|
|
source_client: fixture
|
|
dest_client: fixture
|
|
request: pytest internal object to get launch fixtures from parametrize
|
|
"""
|
|
major, minor, patch, dev = read_version()
|
|
if not dev and None not in (major, minor, patch) and (major, minor, patch) < (1, 10, 0):
|
|
pytest.skip("Multivectors are supported as of v1.10.0")
|
|
|
|
source_client: QdrantClient = request.getfixturevalue(source_client)
|
|
dest_client: QdrantClient = request.getfixturevalue(dest_client)
|
|
collection_name = "multivectors_collection"
|
|
initialize_fixture_collection(
|
|
source_client, collection_name=collection_name, vectors_config=multi_vector_config
|
|
)
|
|
multi_vectors_points = generate_multivector_fixtures()
|
|
source_client.upload_points(collection_name, multi_vectors_points, wait=True)
|
|
source_client.migrate(dest_client)
|
|
compare_collections(
|
|
source_client,
|
|
dest_client,
|
|
num_vectors=VECTOR_NUMBER,
|
|
collection_name=collection_name,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"source_client,dest_client",
|
|
[
|
|
("local_client", "remote_client"),
|
|
("remote_client", "local_client"),
|
|
("local_client", "second_local_client"),
|
|
],
|
|
)
|
|
def test_sparse_vector_collection(source_client, dest_client, request) -> None:
|
|
"""
|
|
Args:
|
|
source_client: fixture
|
|
dest_client: fixture
|
|
request: pytest internal object to get launch fixtures from parametrize
|
|
"""
|
|
source_client: QdrantClient = request.getfixturevalue(source_client)
|
|
dest_client: QdrantClient = request.getfixturevalue(dest_client)
|
|
collection_name = "sparse_vector_collection"
|
|
initialize_fixture_collection(
|
|
source_client, collection_name=collection_name, sparse_vectors_config=sparse_vectors_config
|
|
)
|
|
sparse_points = generate_sparse_fixtures()
|
|
source_client.upload_points(collection_name, sparse_points, wait=True)
|
|
source_client.migrate(dest_client)
|
|
compare_collections(
|
|
source_client,
|
|
dest_client,
|
|
num_vectors=VECTOR_NUMBER,
|
|
collection_name=collection_name,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"source_client,dest_client",
|
|
[
|
|
("local_client", "remote_client"),
|
|
("remote_client", "local_client"),
|
|
("local_client", "second_local_client"),
|
|
],
|
|
)
|
|
def test_migrate_all_collections(source_client, dest_client, request) -> None:
|
|
"""
|
|
Args:
|
|
source_client: fixture
|
|
dest_client: fixture
|
|
request: pytest internal object to get launch fixtures from parametrize
|
|
"""
|
|
vector_number = 100
|
|
collection_names = ["collection_1", "collection_2", "collection_3"]
|
|
source_client: QdrantClient = request.getfixturevalue(source_client)
|
|
dest_client: QdrantClient = request.getfixturevalue(dest_client)
|
|
for collection_name in collection_names:
|
|
initialize_fixture_collection(source_client, collection_name=collection_name)
|
|
points = generate_fixtures(vector_number)
|
|
source_client.upload_points(
|
|
collection_name,
|
|
points,
|
|
wait=True,
|
|
)
|
|
|
|
source_client.migrate(dest_client)
|
|
|
|
for collection_name in collection_names:
|
|
compare_collections(
|
|
source_client,
|
|
dest_client,
|
|
num_vectors=vector_number,
|
|
collection_name=collection_name,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"source_client,dest_client",
|
|
[
|
|
("local_client", "remote_client"),
|
|
("remote_client", "local_client"),
|
|
("local_client", "second_local_client"),
|
|
],
|
|
)
|
|
def test_migrate_particular_collections(source_client, dest_client, request) -> None:
|
|
"""
|
|
Args:
|
|
source_client: fixture
|
|
dest_client: fixture
|
|
request: pytest internal object to get launch fixtures from parametrize
|
|
"""
|
|
vector_number = 100
|
|
collection_names = ["collection_1", "collection_2", "collection_3"]
|
|
source_client: QdrantClient = request.getfixturevalue(source_client)
|
|
dest_client: QdrantClient = request.getfixturevalue(dest_client)
|
|
for collection_name in collection_names:
|
|
initialize_fixture_collection(source_client, collection_name=collection_name)
|
|
points = generate_fixtures(vector_number)
|
|
source_client.upload_points(
|
|
collection_name,
|
|
points,
|
|
wait=True,
|
|
)
|
|
|
|
source_client.migrate(dest_client, collection_names=collection_names[:2])
|
|
|
|
for collection_name in collection_names[:2]:
|
|
compare_collections(
|
|
source_client,
|
|
dest_client,
|
|
num_vectors=vector_number,
|
|
collection_name=collection_name,
|
|
)
|
|
|
|
for collection_name in collection_names[2:]:
|
|
with pytest.raises((qdrant_exceptions.UnexpectedResponse, ValueError)): # type: ignore
|
|
dest_client.get_collection(collection_name)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"source_client,dest_client",
|
|
[
|
|
("local_client", "remote_client"),
|
|
("remote_client", "local_client"),
|
|
("local_client", "second_local_client"),
|
|
],
|
|
)
|
|
def test_action_on_collision(source_client, dest_client, request) -> None:
|
|
"""
|
|
Args:
|
|
source_client: fixture
|
|
dest_client: fixture
|
|
request: pytest internal object to get launch fixtures from parametrize
|
|
"""
|
|
collection_name = "test_collection"
|
|
source_client: QdrantClient = request.getfixturevalue(source_client)
|
|
dest_client: QdrantClient = request.getfixturevalue(dest_client)
|
|
initialize_fixture_collection(source_client, collection_name=collection_name)
|
|
initialize_fixture_collection(dest_client, collection_name=collection_name)
|
|
|
|
with pytest.raises(ValueError):
|
|
source_client.migrate(dest_client, recreate_on_collision=False)
|
|
|
|
points = generate_fixtures(VECTOR_NUMBER)
|
|
source_client.upload_points(
|
|
collection_name,
|
|
points,
|
|
wait=True,
|
|
)
|
|
source_client.migrate(dest_client, recreate_on_collision=True)
|
|
compare_collections(
|
|
source_client,
|
|
dest_client,
|
|
num_vectors=VECTOR_NUMBER,
|
|
collection_name=collection_name,
|
|
)
|
|
|
|
|
|
def test_vector_params(
|
|
local_client: QdrantClient,
|
|
second_local_client: QdrantClient,
|
|
remote_client: QdrantClient,
|
|
):
|
|
collection_name = "test_collection"
|
|
|
|
image_hnsw_config = models.HnswConfigDiff(
|
|
m=9,
|
|
ef_construct=99,
|
|
full_scan_threshold=42,
|
|
max_indexing_threads=4,
|
|
on_disk=True,
|
|
payload_m=5,
|
|
)
|
|
image_quantization_config = models.ScalarQuantization(
|
|
scalar=models.ScalarQuantizationConfig(
|
|
type=models.ScalarType.INT8, quantile=0.69, always_ram=False
|
|
)
|
|
)
|
|
|
|
image_on_disk = True
|
|
|
|
vectors_config = {
|
|
"text": models.VectorParams(size=10, distance=models.Distance.COSINE),
|
|
"image": models.VectorParams(
|
|
size=20,
|
|
distance=models.Distance.DOT,
|
|
hnsw_config=image_hnsw_config,
|
|
quantization_config=image_quantization_config,
|
|
on_disk=image_on_disk,
|
|
),
|
|
}
|
|
|
|
local_client.create_collection(collection_name=collection_name, vectors_config=vectors_config)
|
|
|
|
local_client.migrate(second_local_client)
|
|
|
|
assert local_client.get_collection(collection_name) == second_local_client.get_collection(
|
|
collection_name
|
|
)
|
|
|
|
local_client.migrate(remote_client)
|
|
|
|
local_collection_vector_params = local_client.get_collection(
|
|
collection_name
|
|
).config.params.vectors
|
|
remote_collection_vector_params = remote_client.get_collection(
|
|
collection_name
|
|
).config.params.vectors
|
|
|
|
assert local_collection_vector_params == remote_collection_vector_params
|
|
|
|
local_client.delete_collection(collection_name)
|
|
|
|
remote_client.migrate(local_client)
|
|
local_collection_vector_params = local_client.get_collection(
|
|
collection_name
|
|
).config.params.vectors
|
|
|
|
assert local_collection_vector_params == remote_collection_vector_params
|
|
|
|
|
|
def test_migrate_missing_collections(
|
|
local_client: QdrantClient, second_local_client: QdrantClient
|
|
):
|
|
collection_name = "test_collection"
|
|
with pytest.raises(AssertionError):
|
|
local_client.migrate(second_local_client, collection_names=[collection_name])
|
|
|
|
|
|
def test_recreate_collection(remote_client: QdrantClient):
|
|
collection_name = "test_collection"
|
|
initialize_fixture_collection(remote_client, collection_name=collection_name)
|
|
collection_before_migrate = remote_client.get_collection(collection_name)
|
|
remote_client.migrate(remote_client, recreate_on_collision=True)
|
|
assert collection_before_migrate == remote_client.get_collection(collection_name)
|
|
|
|
remote_client.delete_collection(collection_name)
|
|
|
|
image_hnsw_config = models.HnswConfigDiff(
|
|
m=9,
|
|
ef_construct=99,
|
|
full_scan_threshold=4200,
|
|
max_indexing_threads=2,
|
|
on_disk=True,
|
|
payload_m=5,
|
|
)
|
|
image_quantization_config = models.ScalarQuantization(
|
|
scalar=models.ScalarQuantizationConfig(
|
|
type=models.ScalarType.INT8, quantile=0.89, always_ram=False
|
|
)
|
|
)
|
|
|
|
image_on_disk = True
|
|
|
|
vectors_config = {
|
|
"text": models.VectorParams(size=10, distance=models.Distance.COSINE),
|
|
"image": models.VectorParams(
|
|
size=20,
|
|
distance=models.Distance.DOT,
|
|
hnsw_config=image_hnsw_config,
|
|
quantization_config=image_quantization_config,
|
|
on_disk=image_on_disk,
|
|
),
|
|
}
|
|
|
|
general_hnsw_config = models.HnswConfigDiff(
|
|
m=13,
|
|
ef_construct=101,
|
|
full_scan_threshold=10_001,
|
|
max_indexing_threads=1,
|
|
on_disk=True,
|
|
payload_m=16,
|
|
)
|
|
optimizers_config = models.OptimizersConfigDiff(
|
|
deleted_threshold=0.21,
|
|
vacuum_min_vector_number=1001,
|
|
default_segment_number=2,
|
|
max_segment_size=42_000,
|
|
memmap_threshold=42_000,
|
|
indexing_threshold=42_000,
|
|
flush_interval_sec=6,
|
|
max_optimization_threads=2,
|
|
)
|
|
|
|
wal_config = models.WalConfigDiff(wal_capacity_mb=42, wal_segments_ahead=3)
|
|
|
|
general_quantization_config = models.ProductQuantization(
|
|
product=models.ProductQuantizationConfig(
|
|
compression=models.CompressionRatio.X4, always_ram=False
|
|
)
|
|
)
|
|
strict_mode_config = models.StrictModeConfig(
|
|
read_rate_limit=1000000
|
|
)
|
|
remote_client.create_collection(
|
|
collection_name,
|
|
vectors_config=vectors_config,
|
|
shard_number=3,
|
|
replication_factor=3,
|
|
write_consistency_factor=2,
|
|
on_disk_payload=True,
|
|
hnsw_config=general_hnsw_config,
|
|
optimizers_config=optimizers_config,
|
|
wal_config=wal_config,
|
|
quantization_config=general_quantization_config,
|
|
strict_mode_config=strict_mode_config
|
|
)
|
|
|
|
remote_client.create_payload_index(
|
|
collection_name,
|
|
field_name="title",
|
|
field_schema=models.PayloadSchemaType.KEYWORD,
|
|
)
|
|
|
|
remote_client.create_payload_index(
|
|
collection_name,
|
|
field_name="description",
|
|
field_schema=models.TextIndexParams(
|
|
type=models.TextIndexType.TEXT,
|
|
tokenizer=models.TokenizerType.PREFIX,
|
|
min_token_len=3,
|
|
max_token_len=5,
|
|
lowercase=False,
|
|
),
|
|
)
|
|
|
|
collection_before_migrate = remote_client.get_collection(collection_name)
|
|
remote_client.migrate(remote_client, recreate_on_collision=True)
|
|
assert collection_before_migrate == remote_client.get_collection(collection_name)
|