Added routing functionality for bulk write and bulk delete requests

Author: gauravbansal-curefit

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

@@ -43,7 +43,9 @@ class OpenSearchDocumentStore:
 
     Usage example:
     ```python
-    from haystack_integrations.document_stores.opensearch import OpenSearchDocumentStore
+    from haystack_integrations.document_stores.opensearch import (
+        OpenSearchDocumentStore,
+    )
     from haystack import Document
 
     document_store = OpenSearchDocumentStore(hosts="localhost:9200")
@@ -393,6 +395,10 @@ def _prepare_bulk_write_request(
         opensearch_actions = []
         for doc in documents:
             doc_dict = doc.to_dict()
+
+            # Extract routing from document metadata
+            doc_routing = doc_dict.pop("_routing", None)
+
             if "sparse_embedding" in doc_dict:
                 sparse_embedding = doc_dict.pop("sparse_embedding", None)
                 if sparse_embedding:
@@ -402,13 +408,17 @@ def _prepare_bulk_write_request(
                         "The `sparse_embedding` field will be ignored.",
                         id=doc.id,
                     )
-            opensearch_actions.append(
-                {
-                    "_op_type": action,
-                    "_id": doc.id,
-                    "_source": doc_dict,
-                }
-            )
+
+            action_dict = {
+                "_op_type": action,
+                "_id": doc.id,
+                "_source": doc_dict,
+            }
+
+            if doc_routing:
+                action_dict["_routing"] = doc_routing
+
+            opensearch_actions.append(action_dict)
 
         return {
             "client": self._client if not is_async else self._async_client,
@@ -498,36 +508,48 @@ def _deserialize_document(hit: dict[str, Any]) -> Document:
 
         return Document.from_dict(data)
 
-    def _prepare_bulk_delete_request(self, *, document_ids: list[str], is_async: bool) -> dict[str, Any]:
+    def _prepare_bulk_delete_request(
+        self, *, document_ids: list[str], is_async: bool, routing: Optional[dict[str, str]] = None
+    ) -> dict[str, Any]:
+        def action_generator():
+            for id_ in document_ids:
+                action = {"_op_type": "delete", "_id": id_}
+                # Add routing if provided for this document ID
+                if routing and id_ in routing:
+                    action["_routing"] = routing[id_]
+                yield action
+
         return {
             "client": self._client if not is_async else self._async_client,
-            "actions": ({"_op_type": "delete", "_id": id_} for id_ in document_ids),
+            "actions": action_generator(),
             "refresh": "wait_for",
             "index": self._index,
             "raise_on_error": False,
             "max_chunk_bytes": self._max_chunk_bytes,
         }
 
-    def delete_documents(self, document_ids: list[str]) -> None:
+    def delete_documents(self, document_ids: list[str], routing: Optional[dict[str, str]] = None) -> None:
         """
         Deletes documents that match the provided `document_ids` from the document store.
 
         :param document_ids: the document ids to delete
+        :param routing: the routing to use when deleting documents
         """
 
         self._ensure_initialized()
 
-        bulk(**self._prepare_bulk_delete_request(document_ids=document_ids, is_async=False))
+        bulk(**self._prepare_bulk_delete_request(document_ids=document_ids, is_async=False, routing=routing))
 
-    async def delete_documents_async(self, document_ids: list[str]) -> None:
+    async def delete_documents_async(self, document_ids: list[str], routing: Optional[dict[str, str]] = None) -> None:
         """
         Asynchronously deletes documents that match the provided `document_ids` from the document store.
 
         :param document_ids: the document ids to delete
+        :param routing: the routing to use when deleting documents
         """
         self._ensure_initialized()
 
-        await async_bulk(**self._prepare_bulk_delete_request(document_ids=document_ids, is_async=True))
+        await async_bulk(**self._prepare_bulk_delete_request(document_ids=document_ids, is_async=True, routing=routing))
 
     def _prepare_delete_all_request(self, *, is_async: bool) -> dict[str, Any]:
         return {

integrations/opensearch/tests/test_document_store.py

@@ -111,6 +111,48 @@ def test_get_default_mappings(_mock_opensearch_client):
     }
 
 
+@patch("haystack_integrations.document_stores.opensearch.document_store.bulk")
+def test_routing_extracted_from_metadata(mock_bulk, document_store):
+    """Test routing extraction from document metadata"""
+    mock_bulk.return_value = (2, [])
+
+    docs = [
+        Document(id="1", content="Doc", meta={"_routing": "user_a", "other": "data"}),
+        Document(id="2", content="Doc"),
+    ]
+    document_store.write_documents(docs)
+
+    actions = list(mock_bulk.call_args.kwargs["actions"])
+
+    # Routing should be at action level, not in _source
+    assert actions[0]["_routing"] == "user_a"
+    assert "_routing" not in actions[0]["_source"].get("meta", {})
+
+    # Other metadata should be preserved
+    assert (
+        actions[0]["_source"].get("meta", {}).get("other") == "data" or actions[0]["_source"].get("other") == "data"
+    )  # Handle both nested and flat
+
+    # Second doc has no routing
+    assert "_routing" not in actions[1]
+    assert "_routing" not in actions[1]["_source"].get("meta", {})
+
+
+@patch("haystack_integrations.document_stores.opensearch.document_store.bulk")
+def test_routing_in_delete(mock_bulk, document_store):
+    """Test routing parameter in delete operations"""
+    mock_bulk.return_value = (2, [])
+
+    routing_map = {"1": "user_a", "2": "user_b"}
+    document_store.delete_documents(["1", "2", "3"], routing=routing_map)
+
+    actions = list(mock_bulk.call_args.kwargs["actions"])
+
+    assert actions[0]["_routing"] == "user_a"
+    assert actions[1]["_routing"] == "user_b"
+    assert "_routing" not in actions[2]
+
+
 @pytest.mark.integration
 class TestDocumentStore(CountDocumentsTest, WriteDocumentsTest, DeleteDocumentsTest):
     """
@@ -576,3 +618,39 @@ def test_update_by_filter(self, document_store: OpenSearchDocumentStore):
         )
         assert len(draft_docs) == 1
         assert draft_docs[0].meta["category"] == "B"
+
+    @pytest.mark.integration
+    def test_write_with_routing(self, document_store: OpenSearchDocumentStore):
+        """Test writing documents with routing metadata"""
+        docs = [
+            Document(id="1", content="User A doc", meta={"_routing": "user_a", "category": "test"}),
+            Document(id="2", content="User B doc", meta={"_routing": "user_b"}),
+            Document(id="3", content="No routing"),
+        ]
+
+        written = document_store.write_documents(docs)
+        assert written == 3
+        assert document_store.count_documents() == 3
+
+        # Verify _routing not stored in metadata
+        retrieved = document_store.filter_documents()
+        for doc in retrieved:
+            assert "_routing" not in doc.meta
+            if doc.id == "1":
+                assert doc.meta["category"] == "test"
+
+    @pytest.mark.integration
+    def test_delete_with_routing(self, document_store: OpenSearchDocumentStore):
+        """Test deleting documents with routing"""
+        docs = [
+            Document(id="1", content="Doc 1", meta={"_routing": "user_a"}),
+            Document(id="2", content="Doc 2", meta={"_routing": "user_b"}),
+            Document(id="3", content="Doc 3"),
+        ]
+        document_store.write_documents(docs)
+
+        routing_map = {"1": "user_a", "2": "user_b"}
+        document_store.delete_documents(["1", "2"], routing=routing_map)
+
+        time.sleep(1)
+        assert document_store.count_documents() == 1

integrations/opensearch/tests/test_document_store_async.py

@@ -302,6 +302,42 @@ async def test_delete_all_documents_no_index_recreation(self, document_store: Op
         assert len(results) == 1
         assert results[0].content == "New document after delete all"
 
+    @pytest.mark.asyncio
+    async def test_write_with_routing(self, document_store: OpenSearchDocumentStore):
+        """Test async writing documents with routing metadata"""
+        docs = [
+            Document(id="1", content="User A doc", meta={"_routing": "user_a", "category": "test"}),
+            Document(id="2", content="User B doc", meta={"_routing": "user_b"}),
+            Document(id="3", content="No routing"),
+        ]
+
+        written = await document_store.write_documents_async(docs)
+        assert written == 3
+        assert await document_store.count_documents_async() == 3
+
+        # Verify _routing not stored in metadata
+        retrieved = await document_store.filter_documents_async()
+        for doc in retrieved:
+            assert "_routing" not in doc.meta
+            if doc.id == "1":
+                assert doc.meta["category"] == "test"
+
+    @pytest.mark.asyncio
+    async def test_delete_with_routing(self, document_store: OpenSearchDocumentStore):
+        """Test async deleting documents with routing"""
+        docs = [
+            Document(id="1", content="Doc 1", meta={"_routing": "user_a"}),
+            Document(id="2", content="Doc 2", meta={"_routing": "user_b"}),
+            Document(id="3", content="Doc 3"),
+        ]
+        document_store.write_documents(docs)
+
+        routing_map = {"1": "user_a", "2": "user_b"}
+        await document_store.delete_documents_async(["1", "2"], routing=routing_map)
+
+        time.sleep(1)
+        assert await document_store.count_documents_async() == 1
+
     async def test_delete_by_filter_async(self, document_store: OpenSearchDocumentStore):
         docs = [
             Document(content="Doc 1", meta={"category": "A"}),