diff --git a/qdrant_client/local/local_collection.py b/qdrant_client/local/local_collection.py index b21c18414..55f1601e6 100644 --- a/qdrant_client/local/local_collection.py +++ b/qdrant_client/local/local_collection.py @@ -88,6 +88,18 @@ # https://github.com/qdrant/qdrant/blob/7164ac4a5987d28f1c93f5712aef8e09e7d93555/lib/segment/src/spaces/simple_avx.rs#L99C10-L99C10 +def normalize_dense(vector: types.NumpyArray) -> types.NumpyArray: + """Unit-normalize a dense vector, the form cosine collections store it in.""" + norm = np.linalg.norm(vector) + return vector / norm if norm > EPSILON else vector + + +def normalize_multivector(vector: types.NumpyArray) -> types.NumpyArray: + """Unit-normalize each token vector, the form cosine collections store them in.""" + vector_norm = np.linalg.norm(vector, axis=-1)[:, np.newaxis] + return vector / np.where(vector_norm != 0.0, vector_norm, EPSILON) + + def to_jsonable_python(x: Any) -> Any: try: return json.loads(json.dumps(x, allow_nan=True)) @@ -279,7 +291,10 @@ def load_vectors(self) -> None: for name in all_dense_vector_names: v = loaded_vector.get(name) if v is not None: - vectors[name].append(v) + vector_np = np.array(v, dtype=np.float32) + if self.vectors_config[name].distance == models.Distance.COSINE: + vector_np = normalize_dense(vector_np) + vectors[name].append(vector_np) else: vectors[name].append( np.ones(self.vectors_config[name].size, dtype=np.float32) @@ -301,7 +316,10 @@ def load_vectors(self) -> None: for name in all_multivector_names: v = loaded_vector.get(name) if v is not None: - multivectors[name].append(v) + vector_np = np.array(v, dtype=np.float32) + if self.multivectors_config[name].distance == models.Distance.COSINE: + vector_np = normalize_multivector(vector_np) + multivectors[name].append(vector_np) else: multivectors[name].append(np.array([])) deleted_ids.append((idx, name)) @@ -320,7 +338,7 @@ def load_vectors(self) -> None: # setup multivectors by name for name, named_vectors in multivectors.items(): - self.multivectors[name] = [np.array(vector) for vector in named_vectors] + self.multivectors[name] = named_vectors self.deleted_per_vector[name] = np.zeros(len(self.payload), dtype=bool) # track deleted points by named vector @@ -2408,12 +2426,12 @@ def _update_point(self, point: models.PointStruct) -> None: for vector_name, _named_vectors in self.vectors.items(): vector = vectors.get(vector_name) if vector is not None: + vector_np = np.array(vector, dtype=np.float32) + assert not np.isnan(vector_np).any(), "Vector contains NaN values" params = self.get_vector_params(vector_name) - assert not np.isnan(vector).any(), "Vector contains NaN values" if params.distance == models.Distance.COSINE: - norm = np.linalg.norm(vector) - vector = np.array(vector) / norm if norm > EPSILON else vector - self.vectors[vector_name][idx] = vector + vector_np = normalize_dense(vector_np) + self.vectors[vector_name][idx] = vector_np self.deleted_per_vector[vector_name][idx] = 0 else: self.deleted_per_vector[vector_name][idx] = 1 @@ -2437,13 +2455,13 @@ def _update_point(self, point: models.PointStruct) -> None: for vector_name, _named_vector in self.multivectors.items(): vector = vectors.get(vector_name) if vector is not None: - params = self.get_vector_params(vector_name) - assert not np.isnan(vector).any(), "Vector contains NaN values" + vector_np = np.array(vector, dtype=np.float32) + assert not np.isnan(vector_np).any(), "Vector contains NaN values" + params = self.get_vector_params(vector_name) if params.distance == models.Distance.COSINE: - vector_norm = np.linalg.norm(vector, axis=-1)[:, np.newaxis] - vector /= np.where(vector_norm != 0.0, vector_norm, EPSILON) - self.multivectors[vector_name][idx] = np.array(vector) + vector_np = normalize_multivector(vector_np) + self.multivectors[vector_name][idx] = vector_np self.deleted_per_vector[vector_name][idx] = 0 else: self.deleted_per_vector[vector_name][idx] = 1 @@ -2488,8 +2506,7 @@ def _add_point(self, point: models.PointStruct) -> None: assert not np.isnan(vector_np).any(), "Vector contains NaN values" params = self.get_vector_params(vector_name) if params.distance == models.Distance.COSINE: - norm = np.linalg.norm(vector_np) - vector_np = vector_np / norm if norm > EPSILON else vector_np + vector_np = normalize_dense(vector_np) named_vectors[idx] = vector_np self.deleted_per_vector[vector_name] = np.append( self.deleted_per_vector[vector_name], 0 @@ -2540,8 +2557,7 @@ def _add_point(self, point: models.PointStruct) -> None: assert not np.isnan(vector_np).any(), "Vector contains NaN values" params = self.get_vector_params(vector_name) if params.distance == models.Distance.COSINE: - vector_norm = np.linalg.norm(vector_np, axis=-1)[:, np.newaxis] - vector_np /= np.where(vector_norm != 0.0, vector_norm, EPSILON) + vector_np = normalize_multivector(vector_np) named_vectors[idx] = vector_np self.deleted_per_vector[vector_name] = np.append( self.deleted_per_vector[vector_name], 0 @@ -2680,13 +2696,11 @@ def _update_named_vectors( params = self.get_vector_params(vector_name) if vector_name in self.vectors: if params.distance == models.Distance.COSINE: - norm = np.linalg.norm(vector_np) - vector_np = vector_np / norm if norm > EPSILON else vector_np + vector_np = normalize_dense(vector_np) self.vectors[vector_name][idx] = vector_np else: if params.distance == models.Distance.COSINE: - vector_norm = np.linalg.norm(vector_np, axis=-1)[:, np.newaxis] - vector_np /= np.where(vector_norm != 0.0, vector_norm, EPSILON) + vector_np = normalize_multivector(vector_np) self.multivectors[vector_name][idx] = vector_np def update_vectors( diff --git a/tests/congruence_tests/test_persistence.py b/tests/congruence_tests/test_persistence.py new file mode 100644 index 000000000..2354bff95 --- /dev/null +++ b/tests/congruence_tests/test_persistence.py @@ -0,0 +1,49 @@ +import tempfile + +from tests.congruence_tests.test_common import ( + compare_collections, + generate_fixtures, + generate_multivector_fixtures, + init_client, + init_local, + init_remote, + multi_vector_config, +) + + +def test_reopened_collection_matches_remote(): + """A collection reopened from disk must still hold what the server holds. + + The other persistence tests here reopen a collection but only compare scored + search results, which agree within `rel_tol=1e-4` even when the stored + vectors do not. Comparing the vectors directly catches a reload that fails + to rebuild them the way the write path stored them. + """ + points = generate_fixtures() + with tempfile.TemporaryDirectory() as tmpdir: + local_client = init_local(tmpdir) + init_client(local_client, points) + local_client.close() + + remote_client = init_remote() + init_client(remote_client, points) + + reopened_client = init_local(tmpdir) + compare_collections(reopened_client, remote_client, len(points)) + reopened_client.close() + + +def test_reopened_multivector_collection_matches_remote(): + """Same as above, for multivectors.""" + points = generate_multivector_fixtures() + with tempfile.TemporaryDirectory() as tmpdir: + local_client = init_local(tmpdir) + init_client(local_client, points, vectors_config=multi_vector_config) + local_client.close() + + remote_client = init_remote() + init_client(remote_client, points, vectors_config=multi_vector_config) + + reopened_client = init_local(tmpdir) + compare_collections(reopened_client, remote_client, len(points)) + reopened_client.close() diff --git a/tests/test_local_persistence.py b/tests/test_local_persistence.py index f095c0f13..a34b3008f 100644 --- a/tests/test_local_persistence.py +++ b/tests/test_local_persistence.py @@ -201,3 +201,191 @@ def test_update_persistence(): "not_important": "missing", } client.close() + + +def ingest_multivector_data( + vector_size: int = 32, + tokens: int = 8, + path: str | None = None, + collection_name: str = default_collection_name, +): + client = QdrantClient(path=path) + + if client.collection_exists(collection_name): + client.delete_collection(collection_name) + client.create_collection( + collection_name, + vectors_config=rest.VectorParams( + size=vector_size, + distance=rest.Distance.COSINE, + multivector_config=rest.MultiVectorConfig( + comparator=rest.MultiVectorComparator.MAX_SIM + ), + ), + ) + client.upsert( + collection_name=collection_name, + points=[ + rest.PointStruct(id=i, vector=np.random.randn(tokens, vector_size).tolist()) + for i in range(10) + ], + ) + return client + + +def test_dense_cosine_persistence_is_stable(): + """Reopening a cosine collection must not change its vectors or its scores. + + Cosine vectors are unit-normalized on upsert, but `load_vectors` handed the + raw persisted vectors straight to the collection. A reopened collection + returned un-normalized vectors, scored differently from the same collection + before it was closed, and got silently normalized in place by the first + search that touched it. + """ + vector_size = 32 + with tempfile.TemporaryDirectory() as tmpdir: + client = ingest_dense_vector_data(vector_size=vector_size, path=tmpdir) + query = np.random.randn(vector_size).tolist() + + ids = sorted(point.id for point in client.scroll(default_collection_name, limit=10)[0]) + retrieve = lambda c: [ + point.vector + for point in sorted( + c.retrieve(default_collection_name, ids, with_vectors=True), + key=lambda point: point.id, + ) + ] + search = lambda c: [ + (point.id, point.score) + for point in c.query_points(default_collection_name, query=query, limit=10).points + ] + + before_vectors, before_scores = retrieve(client), search(client) + before_searched_vectors = retrieve(client) + client.close() + + client = QdrantClient(path=tmpdir) + after_vectors = retrieve(client) + after_scores = search(client) + after_searched_vectors = retrieve(client) + + assert after_vectors == before_vectors + assert after_scores == before_scores + assert after_searched_vectors == before_searched_vectors + for vector in after_vectors: + assert np.isclose(np.linalg.norm(vector), 1.0) + client.close() + + +def test_multivector_cosine_persistence_is_stable(): + """Same as the dense case, for multivectors.""" + vector_size, tokens = 32, 8 + with tempfile.TemporaryDirectory() as tmpdir: + client = ingest_multivector_data(vector_size=vector_size, tokens=tokens, path=tmpdir) + query = np.random.randn(3, vector_size).tolist() + + ids = list(range(10)) + retrieve = lambda c: [ + point.vector + for point in sorted( + c.retrieve(default_collection_name, ids, with_vectors=True), + key=lambda point: point.id, + ) + ] + search = lambda c: [ + (point.id, point.score) + for point in c.query_points(default_collection_name, query=query, limit=10).points + ] + + before_vectors, before_scores = retrieve(client), search(client) + before_searched_vectors = retrieve(client) + client.close() + + client = QdrantClient(path=tmpdir) + after_vectors = retrieve(client) + after_scores = search(client) + after_searched_vectors = retrieve(client) + + assert after_vectors == before_vectors + assert after_scores == before_scores + assert after_searched_vectors == before_searched_vectors + for multivector in after_vectors: + assert np.allclose(np.linalg.norm(multivector, axis=-1), 1.0) + client.close() + + +def test_upsert_over_an_existing_point_matches_a_fresh_insert(): + """Overwriting a point must store exactly what inserting it fresh would store. + + `_add_point` casts to float32 before normalizing, `_update_point` normalized + first and cast after, so the same vector landed in the collection with + different values depending on which path wrote it. + """ + vector_size, tokens = 32, 8 + vector = { + "dense": np.random.randn(vector_size).tolist(), + "multi": np.random.randn(tokens, vector_size).tolist(), + } + with tempfile.TemporaryDirectory() as tmpdir: + client = QdrantClient(path=tmpdir) + client.create_collection( + default_collection_name, + vectors_config={ + "dense": rest.VectorParams(size=vector_size, distance=rest.Distance.COSINE), + "multi": rest.VectorParams( + size=vector_size, + distance=rest.Distance.COSINE, + multivector_config=rest.MultiVectorConfig( + comparator=rest.MultiVectorComparator.MAX_SIM + ), + ), + }, + ) + # id 1 is inserted once, id 2 is inserted and then overwritten + client.upsert( + default_collection_name, + points=[rest.PointStruct(id=i, vector=vector) for i in (1, 2)], + ) + client.upsert(default_collection_name, points=[rest.PointStruct(id=2, vector=vector)]) + + inserted, updated = sorted( + client.retrieve(default_collection_name, [1, 2], with_vectors=True), + key=lambda point: point.id, + ) + assert inserted.vector == updated.vector + client.close() + + +def test_zero_norm_cosine_vector_survives_reload(): + """A cosine vector of zero norm must come back as it went in. + + Normalizing it would divide by zero, so the write path leaves it alone. The + reload path has to make the same exception. + """ + vector_size = 32 + with tempfile.TemporaryDirectory() as tmpdir: + client = QdrantClient(path=tmpdir) + client.create_collection( + default_collection_name, + vectors_config=rest.VectorParams(size=vector_size, distance=rest.Distance.COSINE), + ) + client.upsert( + default_collection_name, + points=[ + rest.PointStruct(id=1, vector=[0.0] * vector_size), + rest.PointStruct(id=2, vector=np.random.randn(vector_size).tolist()), + ], + ) + before = client.retrieve(default_collection_name, [1], with_vectors=True)[0].vector + client.close() + + client = QdrantClient(path=tmpdir) + after = client.retrieve(default_collection_name, [1], with_vectors=True)[0].vector + assert after == before + assert not np.isnan(after).any() + # the point must still be searchable, not poison the whole collection + assert ( + len(client.query_points(default_collection_name, query=[1.0] * vector_size).points) + == 2 + ) + client.close()