Add optional data load preprocessor hook (#67)

Another finding from working the arxiv example.. often you need to
unpack or edit a record before writing to the source. If you do this
before invoking redisvl, you add an additional loop, one that is
unnecessary in the end. So this allows devs to optionally add a
preprocessor method on load to call against each record.

Author: tylerhutcherson

redisvl/index.py

@@ -1,5 +1,5 @@
 import asyncio
-from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional, Union
 from uuid import uuid4
 
 if TYPE_CHECKING:
@@ -217,13 +217,17 @@ def delete(self, drop: bool = True):
         Args:
             drop (bool, optional): Delete the documents in the index. Defaults to True.
 
-        raises:
+        Raises:
             redis.exceptions.ResponseError: If the index does not exist.
         """
         raise NotImplementedError
 
     def load(
-        self, data: Iterable[Dict[str, Any]], key_field: Optional[str] = None, **kwargs
+        self,
+        data: Iterable[Dict[str, Any]],
+        key_field: Optional[str] = None,
+        preprocess: Optional[Callable] = None,
+        **kwargs,
     ):
         """Load data into Redis and index using this SearchIndex object.
 
@@ -232,8 +236,10 @@ def load(
                 containing the data to be indexed.
             key_field (Optional[str], optional): A field within the record
                 to use in the Redis hash key.
+            preprocess (Optional[Callabl], optional): An optional preprocessor function
+                that mutates the individual record before writing to redis.
 
-        raises:
+        Raises:
             redis.exceptions.ResponseError: If the index does not exist.
         """
         raise NotImplementedError
@@ -357,7 +363,11 @@ def delete(self, drop: bool = True):
 
     @check_connected("_redis_conn")
     def load(
-        self, data: Iterable[Dict[str, Any]], key_field: Optional[str] = None, **kwargs
+        self,
+        data: Iterable[Dict[str, Any]],
+        key_field: Optional[str] = None,
+        preprocess: Optional[Callable] = None,
+        **kwargs,
     ):
         """Load data into Redis and index using this SearchIndex object.
 
@@ -366,9 +376,16 @@ def load(
                 containing the data to be indexed.
             key_field (Optional[str], optional): A field within the record to
                 use in the Redis hash key.
+            preprocess (Optional[Callable], optional): An optional preprocessor function
+                that mutates the individual record before writing to redis.
 
         raises:
             redis.exceptions.ResponseError: If the index does not exist.
+
+        Example:
+            >>> data = [{"foo": "bar"}, {"test": "values"}]
+            >>> def func(record: dict): record["new"]="value";return record
+            >>> index.load(data, preprocess=func)
         """
         # TODO -- should we return a count of the upserts? or some kind of metadata?
         if data:
@@ -381,6 +398,19 @@ def load(
             with self._redis_conn.pipeline(transaction=False) as pipe:  # type: ignore
                 for record in data:
                     key = self._create_key(record, key_field)
+                    # Optionally preprocess the record and validate type
+                    if preprocess:
+                        try:
+                            record = preprocess(record)
+                        except Exception as e:
+                            raise RuntimeError(
+                                "Error while preprocessing records on load"
+                            ) from e
+                    if not isinstance(record, dict):
+                        raise TypeError(
+                            f"Individual records must be of type dict, got type {type(record)}"
+                        )
+                    # Write the record to Redis
                     pipe.hset(key, mapping=record)  # type: ignore
                     if ttl:
                         pipe.expire(key, ttl)
@@ -406,8 +436,8 @@ class AsyncSearchIndex(SearchIndexBase):
     Example:
         >>> from redisvl.index import AsyncSearchIndex
         >>> index = AsyncSearchIndex.from_yaml("schema.yaml")
-        >>> index.create(overwrite=True)
-        >>> index.load(data) # data is an iterable of dictionaries
+        >>> await index.create(overwrite=True)
+        >>> await index.load(data) # data is an iterable of dictionaries
     """
 
     def __init__(
@@ -502,7 +532,7 @@ async def delete(self, drop: bool = True):
         Args:
             drop (bool, optional): Delete the documents in the index. Defaults to True.
 
-        raises:
+        Raises:
             redis.exceptions.ResponseError: If the index does not exist.
         """
         # Delete the search index
@@ -514,6 +544,7 @@ async def load(
         data: Iterable[Dict[str, Any]],
         concurrency: int = 10,
         key_field: Optional[str] = None,
+        preprocess: Optional[Callable] = None,
         **kwargs,
     ):
         """Load data into Redis and index using this SearchIndex object.
@@ -524,21 +555,41 @@ async def load(
             concurrency (int, optional): Number of concurrent tasks to run. Defaults to 10.
             key_field (Optional[str], optional): A field within the record to
                 use in the Redis hash key.
+            preprocess (Optional[Callable], optional): An optional preprocessor function
+                that mutates the individual record before writing to redis.
 
-        raises:
+        Raises:
             redis.exceptions.ResponseError: If the index does not exist.
+
+        Example:
+            >>> data = [{"foo": "bar"}, {"test": "values"}]
+            >>> def func(record: dict): record["new"]="value";return record
+            >>> await index.load(data, preprocess=func)
         """
         ttl = kwargs.get("ttl")
         semaphore = asyncio.Semaphore(concurrency)
 
         async def _load(record: dict):
             async with semaphore:
                 key = self._create_key(record, key_field)
+                # Optionally preprocess the record and validate type
+                if preprocess:
+                    try:
+                        record = preprocess(record)
+                    except Exception as e:
+                        raise RuntimeError(
+                            "Error while preprocessing records on load"
+                        ) from e
+                if not isinstance(record, dict):
+                    raise TypeError(
+                        f"Individual records must be of type dict, got type {type(record)}"
+                    )
+                # Write the record to Redis
                 await self._redis_conn.hset(key, mapping=record)  # type: ignore
                 if ttl:
                     await self._redis_conn.expire(key, ttl)  # type: ignore
 
-        # gather with concurrency
+        # Gather with concurrency
         await asyncio.gather(*[_load(record) for record in data])
 
     @check_connected("_redis_conn")

tests/test_index.py

@@ -62,6 +62,26 @@ def test_search_index_load(client):
     assert convert_bytes(client.hget("rvl:1", "value")) == "test"
 
 
+def test_search_index_load_preprocess(client):
+    si = SearchIndex("my_index", fields=fields)
+    si.set_client(client)
+    si.create(overwrite=True)
+    data = [{"id": "1", "value": "test"}]
+
+    def preprocess(record):
+        record["test"] = "foo"
+        return record
+
+    si.load(data, key_field="id", preprocess=preprocess)
+    assert convert_bytes(client.hget("rvl:1", "test")) == "foo"
+
+    def bad_preprocess(record):
+        return 1
+
+    with pytest.raises(TypeError):
+        si.load(data, key_field="id", preprocess=bad_preprocess)
+
+
 @pytest.mark.asyncio
 async def test_async_search_index_creation(async_client):
     asi = AsyncSearchIndex("my_index", fields=fields)