Add a drop_documents() method to search index classes

Author: abrookins

redisvl/index/index.py

@@ -592,6 +592,27 @@ def drop_keys(self, keys: Union[str, List[str]]) -> int:
         else:
             return self._redis_client.delete(keys)  # type: ignore
 
+    def drop_documents(self, ids: Union[str, List[str]]) -> int:
+        """Remove documents from the index by their document IDs.
+
+        This method converts document IDs to Redis keys automatically by applying
+        the index's key prefix and separator configuration.
+
+        Args:
+            ids (Union[str, List[str]]): The document ID or IDs to remove from the index.
+
+        Returns:
+            int: Count of documents deleted from Redis.
+        """
+        if isinstance(ids, list):
+            if not ids:
+                return 0
+            keys = [self.key(id) for id in ids]
+            return self._redis_client.delete(*keys)  # type: ignore
+        else:
+            key = self.key(ids)
+            return self._redis_client.delete(key)  # type: ignore
+
     def expire_keys(
         self, keys: Union[str, List[str]], ttl: int
     ) -> Union[int, List[int]]:
@@ -1236,6 +1257,28 @@ async def drop_keys(self, keys: Union[str, List[str]]) -> int:
         else:
             return await client.delete(keys)
 
+    async def drop_documents(self, ids: Union[str, List[str]]) -> int:
+        """Remove documents from the index by their document IDs.
+
+        This method converts document IDs to Redis keys automatically by applying
+        the index's key prefix and separator configuration.
+
+        Args:
+            ids (Union[str, List[str]]): The document ID or IDs to remove from the index.
+
+        Returns:
+            int: Count of documents deleted from Redis.
+        """
+        client = await self._get_client()
+        if isinstance(ids, list):
+            if not ids:
+                return 0
+            keys = [self.key(id) for id in ids]
+            return await client.delete(*keys)
+        else:
+            key = self.key(ids)
+            return await client.delete(key)
+
     async def expire_keys(
         self, keys: Union[str, List[str]], ttl: int
     ) -> Union[int, List[int]]:
@@ -1356,7 +1399,7 @@ async def fetch(self, id: str) -> Optional[Dict[str, Any]]:
     async def _aggregate(
         self, aggregation_query: AggregationQuery
     ) -> List[Dict[str, Any]]:
-        """Execute an aggretation query and processes the results."""
+        """Execute an aggregation query and processes the results."""
         results = await self.aggregate(
             aggregation_query, query_params=aggregation_query.params  # type: ignore[attr-defined]
         )

tests/integration/test_async_search_index.py

@@ -246,6 +246,37 @@ async def test_search_index_drop_keys(async_index):
     assert await async_index.exists()
 
 
+@pytest.mark.asyncio
+async def test_search_index_drop_documents(async_index):
+    await async_index.create(overwrite=True, drop=True)
+    data = [
+        {"id": "1", "test": "foo"},
+        {"id": "2", "test": "bar"},
+        {"id": "3", "test": "baz"},
+    ]
+    await async_index.load(data, id_field="id")
+
+    # Test dropping a single document by ID
+    dropped = await async_index.drop_documents("1")
+    assert dropped == 1
+    assert not await async_index.fetch("1")
+    assert await async_index.fetch("2") is not None
+    assert await async_index.fetch("3") is not None
+
+    # Test dropping multiple documents by ID
+    dropped = await async_index.drop_documents(["2", "3"])
+    assert dropped == 2
+    assert not await async_index.fetch("2")
+    assert not await async_index.fetch("3")
+
+    # Test dropping with an empty list
+    dropped = await async_index.drop_documents([])
+    assert dropped == 0
+
+    # Ensure the index still exists
+    assert await async_index.exists()
+
+
 @pytest.mark.asyncio
 async def test_search_index_load_and_fetch(async_index):
     await async_index.create(overwrite=True, drop=True)

tests/integration/test_search_index.py

@@ -240,6 +240,36 @@ def test_search_index_drop_keys(index):
     assert index.exists()
 
 
+def test_search_index_drop_documents(index):
+    index.create(overwrite=True, drop=True)
+    data = [
+        {"id": "1", "test": "foo"},
+        {"id": "2", "test": "bar"},
+        {"id": "3", "test": "baz"},
+    ]
+    index.load(data, id_field="id")
+
+    # Test dropping a single document by ID
+    dropped = index.drop_documents("1")
+    assert dropped == 1
+    assert not index.fetch("1")
+    assert index.fetch("2") is not None
+    assert index.fetch("3") is not None
+
+    # Test dropping multiple documents by ID
+    dropped = index.drop_documents(["2", "3"])
+    assert dropped == 2
+    assert not index.fetch("2")
+    assert not index.fetch("3")
+
+    # Test dropping with an empty list
+    dropped = index.drop_documents([])
+    assert dropped == 0
+
+    # Ensure the index still exists
+    assert index.exists()
+
+
 def test_search_index_load_and_fetch(index):
     index.create(overwrite=True, drop=True)
     data = [{"id": "1", "test": "foo"}]