import numpy as np
from qdrant_client.http import models
from qdrant_client.local.distances import calculate_distance
from qdrant_client.local.multi_distances import calculate_multi_distance
from qdrant_client.local.sparse_distances import calculate_distance_sparse
def test_distances() -> None:
query = np.array([1.0, 2.0, 3.0])
vectors = np.array([[1.0, 2.0, 3.0], [1.0, 2.0, 3.0]])
assert np.allclose(calculate_distance(query, vectors, models.Distance.DOT), [14.0, 14.0])
assert np.allclose(calculate_distance(query, vectors, models.Distance.EUCLID), [0.0, 0.0])
assert np.allclose(calculate_distance(query, vectors, models.Distance.MANHATTAN), [0.0, 0.0])
# cosine modifies vectors inplace
assert np.allclose(calculate_distance(query, vectors, models.Distance.COSINE), [1.0, 1.0])
query = np.array([1.0, 0.0, 1.0])
vectors = np.array([[1.0, 2.0, 3.0], [0.0, 1.0, 0.0]])
assert np.allclose(
calculate_distance(query, vectors, models.Distance.DOT), [4.0, 0.0], atol=0.0001
)
assert np.allclose(
calculate_distance(query, vectors, models.Distance.EUCLID),
[2.82842712, 1.7320508],
atol=0.0001,
)
assert np.allclose(
calculate_distance(query, vectors, models.Distance.MANHATTAN),
[4.0, 3.0],
atol=0.0001,
)
# cosine modifies vectors inplace
assert np.allclose(
calculate_distance(query, vectors, models.Distance.COSINE),
[0.75592895, 0.0],
atol=0.0001,
)
sparse_query = models.SparseVector(indices=[1, 2], values=[1, 2])
sparse_vectors = [models.SparseVector(indices=[10, 20], values=[1, 2])]
assert calculate_distance_sparse(sparse_query, sparse_vectors) == [np.float32("-inf")]
sparse_vectors = [
models.SparseVector(indices=[1, 2], values=[3, 4]),
models.SparseVector(indices=[1, 2, 3], values=[1, 2, 3]),
]
assert np.allclose(
calculate_distance_sparse(sparse_query, sparse_vectors), [11.0, 5], atol=0.0001
)
multivector_query = np.array([[1, 2, 3], [3, 4, 5]])
docs = [np.array([[1, 2, 3], [0, 1, 2]])]
assert calculate_multi_distance(multivector_query, docs, models.Distance.DOT)[0] == 40.0