adds drop_keys and async drop_keys methods to index

Author: justin-cechmanek

redisvl/index/index.py

@@ -497,6 +497,19 @@ def clear(self) -> int:
 
         return total_records_deleted
 
+    def drop_keys(self, keys: Union[str, List[str]]) -> None:
+        """Remove a specific entry or entries from the index by it's key ID.
+        Args:
+            keys (Union[str, List[str]]): The document ID or IDs to remove from the index.
+        """
+        if isinstance(keys, List):
+            with self._redis_client.pipeline(transaction=False) as pipe:  # type: ignore
+                for key in keys:  # type: ignore
+                    pipe.delete(key)
+                pipe.execute()
+        else:
+            self._redis_client.delete(keys)  # type: ignore
+
     def load(
         self,
         data: Iterable[Any],
@@ -935,6 +948,16 @@ async def clear(self) -> int:
 
         return total_records_deleted
 
+    async def drop_keys(self, keys: Union[str, List[str]]) -> None:
+        """Remove a specific entry or entries from the index by it's key ID.
+        Args:
+            keys (Union[str, List[str]]): The document ID or IDs to remove from the index.
+        """
+        if isinstance(keys, List):
+            await self._redis_client.delete(*keys)  # type: ignore
+        else:
+            await self._redis_client.delete(keys)  # type: ignore
+
     async def load(
         self,
         data: Iterable[Any],

tests/integration/test_async_search_index.py

@@ -184,6 +184,37 @@ async def test_search_index_clear(async_client, async_index):
     assert await async_index.exists()
 
 
+@pytest.mark.asyncio
+async def test_search_index_drop_key(async_client, async_index):
+    async_index.set_client(async_client)
+    await async_index.create(overwrite=True, drop=True)
+    data = [{"id": "1", "test": "foo"}, {"id": "2", "test": "bar"}]
+    keys = await async_index.load(data, id_field="id")
+
+    await async_index.drop_keys(keys[0])
+    assert not await async_index.fetch(keys[0])
+    assert await async_index.fetch(keys[1]) is not None
+
+
+@pytest.mark.asyncio
+async def test_search_index_drop_keys(async_client, async_index):
+    async_index.set_client(async_client)
+    await async_index.create(overwrite=True, drop=True)
+    data = [
+        {"id": "1", "test": "foo"},
+        {"id": "2", "test": "bar"},
+        {"id": "3", "test": "baz"},
+    ]
+    keys = await async_index.load(data, id_field="id")
+
+    await async_index.drop_keys(keys[0:2])
+    assert not await async_index.fetch(keys[0])
+    assert not await async_index.fetch(keys[1])
+    assert await async_index.fetch(keys[2]) is not None
+
+    assert await async_index.exists()
+
+
 @pytest.mark.asyncio
 async def test_search_index_load_and_fetch(async_client, async_index):
     async_index.set_client(async_client)

tests/integration/test_search_index.py

@@ -170,6 +170,37 @@ def test_search_index_clear(client, index):
     assert index.exists()
 
 
+def test_search_index_drop_key(client, index):
+    index.set_client(client)
+    index.create(overwrite=True, drop=True)
+    data = [{"id": "1", "test": "foo"}, {"id": "2", "test": "bar"}]
+    keys = index.load(data, id_field="id")
+
+    # test passing a single string key removes only that key
+    index.drop_keys(keys[0])
+    assert not index.fetch(keys[0])
+    assert index.fetch(keys[1]) is not None  # still have all other entries
+
+
+def test_search_index_drop_keys(client, index):
+    index.set_client(client)
+    index.create(overwrite=True, drop=True)
+    data = [
+        {"id": "1", "test": "foo"},
+        {"id": "2", "test": "bar"},
+        {"id": "3", "test": "baz"},
+    ]
+    keys = index.load(data, id_field="id")
+
+    # test passing a list of keys selectively removes only those keys
+    index.drop_keys(keys[0:2])
+    assert not index.fetch(keys[0])
+    assert not index.fetch(keys[1])
+    assert index.fetch(keys[2]) is not None
+
+    assert index.exists()
+
+
 def test_search_index_load_and_fetch(client, index):
     index.set_client(client)
     index.create(overwrite=True, drop=True)