feat(Qdrant): introducing use_sparse_embeddings parameters for document store to make sparse embeddings non breaking change. Need more testing

Author: lambda-science

integrations/qdrant/src/haystack_integrations/document_stores/qdrant/converters.py

@@ -21,22 +21,26 @@ def documents_to_batch(
         documents: List[Document],
         *,
         embedding_field: str,
+        use_sparse_embeddings: bool,
         sparse_embedding_field: str,
     ) -> List[rest.PointStruct]:
         points = []
         for document in documents:
             payload = document.to_dict(flatten=False)
-            vector = {}
+            if use_sparse_embeddings:
+                vector = {}
 
-            dense_vector = payload.pop(embedding_field, None)
-            if dense_vector is not None:
-                vector[DENSE_VECTORS_NAME] = dense_vector
+                dense_vector = payload.pop(embedding_field, None)
+                if dense_vector is not None:
+                    vector[DENSE_VECTORS_NAME] = dense_vector
 
-            sparse_vector = payload.pop(sparse_embedding_field, None)
-            if sparse_vector is not None:
-                sparse_vector_instance = rest.SparseVector(**sparse_vector)
-                vector[SPARSE_VECTORS_NAME] = sparse_vector_instance
+                sparse_vector = payload.pop(sparse_embedding_field, None)
+                if sparse_vector is not None:
+                    sparse_vector_instance = rest.SparseVector(**sparse_vector)
+                    vector[SPARSE_VECTORS_NAME] = sparse_vector_instance
 
+            if not use_sparse_embeddings:
+                vector = payload.pop(embedding_field) or {}
             _id = self.convert_id(payload.get("id"))
 
             point = rest.PointStruct(
@@ -61,25 +65,38 @@ def convert_id(self, _id: str) -> str:
 
 
 class QdrantToHaystack:
-    def __init__(self, content_field: str, name_field: str, embedding_field: str, sparse_embedding_field: str):
+    def __init__(
+        self,
+        content_field: str,
+        name_field: str,
+        embedding_field: str,
+        use_sparse_embeddings: bool,  # noqa: FBT001
+        sparse_embedding_field: str,
+    ):
         self.content_field = content_field
         self.name_field = name_field
         self.embedding_field = embedding_field
+        self.use_sparse_embeddings = use_sparse_embeddings
         self.sparse_embedding_field = sparse_embedding_field
 
     def point_to_document(self, point: QdrantPoint) -> Document:
         payload = {**point.payload}
-        if hasattr(point, "vector") and point.vector is not None and DENSE_VECTORS_NAME in point.vector:
-            payload["embedding"] = point.vector[DENSE_VECTORS_NAME]
-        else:
-            payload["embedding"] = None
         payload["score"] = point.score if hasattr(point, "score") else None
-        if hasattr(point, "vector") and point.vector is not None and SPARSE_VECTORS_NAME in point.vector:
-            parse_vector_dict = {
-                "indices": point.vector[SPARSE_VECTORS_NAME].indices,
-                "values": point.vector[SPARSE_VECTORS_NAME].values,
-            }
-            payload["sparse_embedding"] = parse_vector_dict
-        else:
-            payload["sparse_embedding"] = None
+        if not self.use_sparse_embeddings:
+            payload["embedding"] = point.vector if hasattr(point, "vector") else None
+
+        if self.use_sparse_embeddings:
+            if hasattr(point, "vector") and point.vector is not None and DENSE_VECTORS_NAME in point.vector:
+                payload["embedding"] = point.vector[DENSE_VECTORS_NAME]
+            else:
+                payload["embedding"] = None
+
+            if hasattr(point, "vector") and point.vector is not None and SPARSE_VECTORS_NAME in point.vector:
+                parse_vector_dict = {
+                    "indices": point.vector[SPARSE_VECTORS_NAME].indices,
+                    "values": point.vector[SPARSE_VECTORS_NAME].values,
+                }
+                payload["sparse_embedding"] = parse_vector_dict
+            else:
+                payload["sparse_embedding"] = None
         return Document.from_dict(payload)

integrations/qdrant/src/haystack_integrations/document_stores/qdrant/document_store.py

@@ -71,6 +71,7 @@ def __init__(
         content_field: str = "content",
         name_field: str = "name",
         embedding_field: str = "embedding",
+        use_sparse_embeddings: bool = False,  # noqa: FBT001, FBT002
         sparse_embedding_field: str = "sparse_embedding",
         similarity: str = "cosine",
         return_embedding: bool = False,  # noqa: FBT001, FBT002
@@ -140,13 +141,16 @@ def __init__(
         self.payload_fields_to_index = payload_fields_to_index
 
         # Make sure the collection is properly set up
-        self._set_up_collection(index, embedding_dim, recreate_index, similarity, on_disk, payload_fields_to_index)
+        self._set_up_collection(
+            index, embedding_dim, recreate_index, similarity, use_sparse_embeddings, on_disk, payload_fields_to_index
+        )
 
         self.embedding_dim = embedding_dim
         self.on_disk = on_disk
         self.content_field = content_field
         self.name_field = name_field
         self.embedding_field = embedding_field
+        self.use_sparse_embeddings = use_sparse_embeddings
         self.sparse_embedding_field = sparse_embedding_field
         self.similarity = similarity
         self.index = index
@@ -155,7 +159,9 @@ def __init__(
         self.duplicate_documents = duplicate_documents
         self.qdrant_filter_converter = QdrantFilterConverter()
         self.haystack_to_qdrant_converter = HaystackToQdrant()
-        self.qdrant_to_haystack = QdrantToHaystack(content_field, name_field, embedding_field, sparse_embedding_field)
+        self.qdrant_to_haystack = QdrantToHaystack(
+            content_field, name_field, embedding_field, use_sparse_embeddings, sparse_embedding_field
+        )
         self.write_batch_size = write_batch_size
         self.scroll_size = scroll_size
 
@@ -196,7 +202,7 @@ def write_documents(
             if not isinstance(doc, Document):
                 msg = f"DocumentStore.write_documents() expects a list of Documents but got an element of {type(doc)}."
                 raise ValueError(msg)
-        self._set_up_collection(self.index, self.embedding_dim, False, self.similarity)
+        self._set_up_collection(self.index, self.embedding_dim, False, self.similarity, self.use_sparse_embeddings)
 
         if len(documents) == 0:
             logger.warning("Calling QdrantDocumentStore.write_documents() with empty list")
@@ -214,6 +220,7 @@ def write_documents(
                 batch = self.haystack_to_qdrant_converter.documents_to_batch(
                     document_batch,
                     embedding_field=self.embedding_field,
+                    use_sparse_embeddings=self.use_sparse_embeddings,
                     sparse_embedding_field=self.sparse_embedding_field,
                 )
 
@@ -309,10 +316,17 @@ def query_by_sparse(
         scale_score: bool = True,  # noqa: FBT001, FBT002
         return_embedding: bool = False,  # noqa: FBT001, FBT002
     ) -> List[Document]:
+
+        if not self.use_sparse_embeddings:
+            message = (
+                "Error: tried to query by sparse vector with a Qdrant "
+                "Document Store initialized with use_sparse_embeddings=False"
+            )
+            raise ValueError(message)
+
         qdrant_filters = self.qdrant_filter_converter.convert(filters)
         query_indices = query_sparse_embedding.indices
         query_values = query_sparse_embedding.values
-
         points = self.client.search(
             collection_name=self.index,
             query_vector=rest.NamedSparseVector(
@@ -326,7 +340,6 @@ def query_by_sparse(
             limit=top_k,
             with_vectors=return_embedding,
         )
-
         results = [self.qdrant_to_haystack.point_to_document(point) for point in points]
         if scale_score:
             for document in results:
@@ -345,17 +358,25 @@ def query_by_embedding(
     ) -> List[Document]:
         qdrant_filters = self.qdrant_filter_converter.convert(filters)
 
-        points = self.client.search(
-            collection_name=self.index,
-            query_vector=rest.NamedVector(
-                name=DENSE_VECTORS_NAME,
-                vector=query_embedding,
-            ),
-            query_filter=qdrant_filters,
-            limit=top_k,
-            with_vectors=return_embedding,
-        )
-
+        if self.use_sparse_embeddings:
+            points = self.client.search(
+                collection_name=self.index,
+                query_vector=rest.NamedVector(
+                    name=DENSE_VECTORS_NAME,
+                    vector=query_embedding,
+                ),
+                query_filter=qdrant_filters,
+                limit=top_k,
+                with_vectors=return_embedding,
+            )
+        if not self.use_sparse_embeddings:
+            points = self.client.search(
+                collection_name=self.index,
+                query_vector=query_embedding,
+                query_filter=qdrant_filters,
+                limit=top_k,
+                with_vectors=return_embedding,
+            )
         results = [self.qdrant_to_haystack.point_to_document(point) for point in points]
         if scale_score:
             for document in results:
@@ -397,6 +418,7 @@ def _set_up_collection(
         embedding_dim: int,
         recreate_collection: bool,  # noqa: FBT001
         similarity: str,
+        use_sparse_embeddings: bool,  # noqa: FBT001
         on_disk: bool = False,  # noqa: FBT001, FBT002
         payload_fields_to_index: Optional[List[dict]] = None,
     ):
@@ -405,7 +427,7 @@ def _set_up_collection(
         if recreate_collection:
             # There is no need to verify the current configuration of that
             # collection. It might be just recreated again.
-            self._recreate_collection(collection_name, distance, embedding_dim, on_disk)
+            self._recreate_collection(collection_name, distance, embedding_dim, on_disk, use_sparse_embeddings)
             # Create Payload index if payload_fields_to_index is provided
             self._create_payload_index(collection_name, payload_fields_to_index)
             return
@@ -421,12 +443,33 @@ def _set_up_collection(
             # Qdrant local raises ValueError if the collection is not found, but
             # with the remote server UnexpectedResponse / RpcError is raised.
             # Until that's unified, we need to catch both.
-            self._recreate_collection(collection_name, distance, embedding_dim, on_disk)
+            self._recreate_collection(collection_name, distance, embedding_dim, on_disk, use_sparse_embeddings)
             # Create Payload index if payload_fields_to_index is provided
             self._create_payload_index(collection_name, payload_fields_to_index)
             return
-        current_distance = collection_info.config.params.vectors[DENSE_VECTORS_NAME].distance
-        current_vector_size = collection_info.config.params.vectors[DENSE_VECTORS_NAME].size
+        if self.use_sparse_embeddings:
+            current_distance = collection_info.config.params.vectors[DENSE_VECTORS_NAME].distance
+            current_vector_size = collection_info.config.params.vectors[DENSE_VECTORS_NAME].size
+        if not self.use_sparse_embeddings:
+            current_distance = collection_info.config.params.vectors.distance
+            current_vector_size = collection_info.config.params.vectors.size
+
+        if self.use_sparse_embeddings and not isinstance(collection_info.config.params.vectors, dict):
+            msg = (
+                f"Collection '{collection_name}' already exists in Qdrant, "
+                f"but it has been originaly created without sparse embedding vectors."
+                f"If you want to use that collection, either set `use_sparse_embeddings=False` "
+                f"or run a migration script "
+                f"to use Named Dense Vectors (`text-sparse`) and Named Sparse Vectors (`text-dense`)."
+            )
+            raise ValueError(msg)
+        if not self.use_sparse_embeddings and isinstance(collection_info.config.params.vectors, dict):
+            msg = (
+                f"Collection '{collection_name}' already exists in Qdrant, "
+                f"but it has been originaly created with sparse embedding vectors."
+                f"If you want to use that collection, please set `use_sparse_embeddings=True`"
+            )
+            raise ValueError(msg)
 
         if current_distance != distance:
             msg = (
@@ -446,33 +489,59 @@ def _set_up_collection(
             )
             raise ValueError(msg)
 
-    def _recreate_collection(self, collection_name: str, distance, embedding_dim: int, on_disk: bool):  # noqa: FBT001
-        self.client.recreate_collection(
-            collection_name=collection_name,
-            vectors_config={
-                DENSE_VECTORS_NAME: rest.VectorParams(
+    def _recreate_collection(
+        self,
+        collection_name: str,
+        distance,
+        embedding_dim: int,
+        on_disk: bool,  # noqa: FBT001
+        use_sparse_embeddings: bool,  # noqa: FBT001
+    ):
+        if use_sparse_embeddings:
+            self.client.recreate_collection(
+                collection_name=collection_name,
+                vectors_config={
+                    DENSE_VECTORS_NAME: rest.VectorParams(
+                        size=embedding_dim,
+                        on_disk=on_disk,
+                        distance=distance,
+                    ),
+                },
+                sparse_vectors_config={
+                    SPARSE_VECTORS_NAME: rest.SparseVectorParams(
+                        index=rest.SparseIndexParams(
+                            on_disk=on_disk,
+                        )
+                    )
+                },
+                shard_number=self.shard_number,
+                replication_factor=self.replication_factor,
+                write_consistency_factor=self.write_consistency_factor,
+                on_disk_payload=self.on_disk_payload,
+                hnsw_config=self.hnsw_config,
+                optimizers_config=self.optimizers_config,
+                wal_config=self.wal_config,
+                quantization_config=self.quantization_config,
+                init_from=self.init_from,
+            )
+        if not use_sparse_embeddings:
+            self.client.recreate_collection(
+                collection_name=collection_name,
+                vectors_config=rest.VectorParams(
                     size=embedding_dim,
                     on_disk=on_disk,
                     distance=distance,
                 ),
-            },
-            sparse_vectors_config={
-                SPARSE_VECTORS_NAME: rest.SparseVectorParams(
-                    index=rest.SparseIndexParams(
-                        on_disk=on_disk,
-                    )
-                )
-            },
-            shard_number=self.shard_number,
-            replication_factor=self.replication_factor,
-            write_consistency_factor=self.write_consistency_factor,
-            on_disk_payload=self.on_disk_payload,
-            hnsw_config=self.hnsw_config,
-            optimizers_config=self.optimizers_config,
-            wal_config=self.wal_config,
-            quantization_config=self.quantization_config,
-            init_from=self.init_from,
-        )
+                shard_number=self.shard_number,
+                replication_factor=self.replication_factor,
+                write_consistency_factor=self.write_consistency_factor,
+                on_disk_payload=self.on_disk_payload,
+                hnsw_config=self.hnsw_config,
+                optimizers_config=self.optimizers_config,
+                wal_config=self.wal_config,
+                quantization_config=self.quantization_config,
+                init_from=self.init_from,
+            )
 
     def _handle_duplicate_documents(
         self,

integrations/qdrant/tests/test_converters.py

@@ -15,11 +15,12 @@ def haystack_to_qdrant() -> HaystackToQdrant:
 
 
 @pytest.fixture
-def qdrant_to_haystack() -> QdrantToHaystack:
+def qdrant_to_haystack(request) -> QdrantToHaystack:
     return QdrantToHaystack(
         content_field=CONTENT_FIELD,
         name_field=NAME_FIELD,
         embedding_field=EMBEDDING_FIELD,
+        use_sparse_embeddings=request.param,
         sparse_embedding_field=SPARSE_EMBEDDING_FIELD,
     )
 
@@ -30,7 +31,8 @@ def test_convert_id_is_deterministic(haystack_to_qdrant: HaystackToQdrant):
     assert first_id == second_id
 
 
-def test_point_to_document_reverts_proper_structure_from_record(
+@pytest.mark.parametrize("qdrant_to_haystack", [True], indirect=True)
+def test_point_to_document_reverts_proper_structure_from_record_with_sparse(
     qdrant_to_haystack: QdrantToHaystack,
 ):
     point = rest.Record(
@@ -56,3 +58,29 @@ def test_point_to_document_reverts_proper_structure_from_record(
     assert {"indices": [7, 1024, 367], "values": [0.1, 0.98, 0.33]} == document.sparse_embedding.to_dict()
     assert {"test_field": 1} == document.meta
     assert 0.0 == np.sum(np.array([1.0, 0.0, 0.0, 0.0]) - document.embedding)
+
+
+@pytest.mark.parametrize("qdrant_to_haystack", [False], indirect=True)
+def test_point_to_document_reverts_proper_structure_from_record_without_sparse(
+    qdrant_to_haystack: QdrantToHaystack,
+):
+    point = rest.Record(
+        id="c7c62e8e-02b9-4ec6-9f88-46bd97b628b7",
+        payload={
+            "id": "my-id",
+            "id_hash_keys": ["content"],
+            "content": "Lorem ipsum",
+            "content_type": "text",
+            "meta": {
+                "test_field": 1,
+            },
+        },
+        vector=[1.0, 0.0, 0.0, 0.0],
+    )
+    document = qdrant_to_haystack.point_to_document(point)
+    assert "my-id" == document.id
+    assert "Lorem ipsum" == document.content
+    assert "text" == document.content_type
+    assert document.sparse_embedding is None
+    assert {"test_field": 1} == document.meta
+    assert 0.0 == np.sum(np.array([1.0, 0.0, 0.0, 0.0]) - document.embedding)