WIP on async batch_query

Author: abrookins

redisvl/index/index.py

@@ -1283,8 +1283,8 @@ async def aggregate(self, *args, **kwargs) -> "AggregateResult":
             raise RedisSearchError(f"Error while aggregating: {str(e)}") from e
 
     async def batch_search(
-        self, queries: List[str], batch_size: int = 100, **query_params
-    ) -> List[List[Dict[str, Any]]]:
+        self, queries: List[BaseQuery], batch_size: int = 100, **query_params
+    ) -> List["Result"]:
         """Perform a search against the index for multiple queries.
 
         This method takes a list of queries and returns a list of search results.
@@ -1298,9 +1298,9 @@ async def batch_search(
                 for each query.
 
         Returns:
-            List[List[Dict[str, Any]]]: The search results.
+            List[Result]: The search results.
         """
-        all_parsed = []
+        all_results = []
         client = await self._get_client()
         search = client.ft(self.schema.index.name)
         options = {}
@@ -1340,16 +1340,8 @@ async def batch_search(
                         query=_built_query,
                         duration=duration,
                     )
-                    parsed = process_results(
-                        parsed_raw,
-                        query=_built_query,
-                        storage_type=self.schema.index.storage_type,
-                    )
-                    # Create separate lists of parsed results for each query
-                    # passed in to the batch_search method, so that callers can
-                    # access the results for each query individually
-                    all_parsed.append(parsed)
-        return all_parsed
+                    all_results.append(parsed_raw)
+        return all_results
 
     async def search(self, *args, **kwargs) -> "Result":
         """Perform a search on this index.
@@ -1367,6 +1359,34 @@ async def search(self, *args, **kwargs) -> "Result":
         except Exception as e:
             raise RedisSearchError(f"Error while searching: {str(e)}") from e
 
+    async def _batch_query(
+        self, queries: List[BaseQuery], batch_size: int = 100
+    ) -> List[List[Dict[str, Any]]]:
+        """Asynchronously execute a batch of queries and process results."""
+        results = await self.batch_search(queries, batch_size=batch_size)
+        all_parsed = []
+        for query, batch_results in zip(queries, results):
+            parsed = process_results(
+                batch_results,
+                query=query,
+                storage_type=self.schema.index.storage_type,
+            )
+            # Create separate lists of parsed results for each query
+            # passed in to the batch_search method, so that callers can
+            # access the results for each query individually
+            all_parsed.append(parsed)
+
+        return all_parsed
+
+    async def batch_query(
+        self, queries: List[BaseQuery], batch_size: int = 100
+    ) -> List[List[Dict[str, Any]]]:
+        """Asynchronously execute a batch of queries and process results."""
+        return await self._batch_query(
+            [query.query for query in queries],
+            batch_size=batch_size,
+        )
+
     async def _query(self, query: BaseQuery) -> List[Dict[str, Any]]:
         """Asynchronously execute a query and process results."""
         results = await self.search(query.query, query_params=query.params)

tests/integration/test_async_search_index.py

@@ -8,6 +8,7 @@
 from redisvl.exceptions import RedisModuleVersionError, RedisSearchError
 from redisvl.index import AsyncSearchIndex
 from redisvl.query import VectorQuery
+from redisvl.query.query import FilterQuery
 from redisvl.redis.utils import convert_bytes
 from redisvl.schema import IndexSchema, StorageType
 
@@ -487,3 +488,15 @@ async def test_batch_search_with_multiple_batches(async_index):
     assert results[3][0]["id"] == "rvl:1"
     assert results[4][0]["id"] == "rvl:2"
     assert len(results[5]) == 0
+
+
+@pytest.mark.asyncio
+async def test_batch_query(async_index):
+    await async_index.create(overwrite=True, drop=True)
+    data = [{"id": "1", "test": "foo"}, {"id": "2", "test": "bar"}]
+    await async_index.load(data, id_field="id")
+
+    query = FilterQuery(filter_expression="@test:{foo}")
+    results = await async_index.batch_query([query])
+
+    assert len(results) == 1