docarray · hanxiao · Aug 26, 2022 · Jul 25, 2022 · Jul 26, 2022 · Jul 26, 2022
diff --git a/docarray/array/document.py b/docarray/array/document.py
@@ -10,10 +10,12 @@
     from docarray.array.annlite import DocumentArrayAnnlite
     from docarray.array.weaviate import DocumentArrayWeaviate
     from docarray.array.elastic import DocumentArrayElastic
+    from docarray.array.redis import DocumentArrayRedis
     from docarray.array.storage.sqlite import SqliteConfig
     from docarray.array.storage.annlite import AnnliteConfig
     from docarray.array.storage.weaviate import WeaviateConfig
     from docarray.array.storage.elastic import ElasticConfig
+    from docarray.array.storage.redis import RedisConfig
 
 
 class DocumentArray(AllMixins, BaseDocumentArray):
@@ -127,6 +129,16 @@ def __new__(
         """Create a Elastic-powered DocumentArray object."""
         ...
 
+    @overload
+    def __new__(
+        cls,
+        _docs: Optional['DocumentArraySourceType'] = None,
+        storage: str = 'redis',
+        config: Optional[Union['RedisConfig', Dict]] = None,
+    ) -> 'DocumentArrayRedis':
+        """Create a Redis-powered DocumentArray object."""
+        ...
+
     def __enter__(self):
         return self
 
@@ -163,6 +175,10 @@ def __new__(cls, *args, storage: str = 'memory', **kwargs):
                 from docarray.array.elastic import DocumentArrayElastic
 
                 instance = super().__new__(DocumentArrayElastic)
+            elif storage == 'redis':
+                from .redis import DocumentArrayRedis
+
+                instance = super().__new__(DocumentArrayRedis)
 
             else:
                 raise ValueError(f'storage=`{storage}` is not supported.')

diff --git a/docarray/array/redis.py b/docarray/array/redis.py
@@ -0,0 +1,19 @@
+from .document import DocumentArray
+from .storage.redis import RedisConfig, StorageMixins
+
+__all__ = ['DocumentArrayRedis', 'RedisConfig']
+
+
+class DocumentArrayRedis(StorageMixins, DocumentArray):
+    """This is a :class:`DocumentArray` that uses Redis as
+    vector search engine and storage.
+    """
+
+    def __new__(cls, *args, **kwargs):
+        """``__new__`` method for :class:`DocumentArrayRedis`
+
+        :param *args: list of args to instantiate the object
+        :param **kwargs: dict of args to instantiate the object
+        :return: the instantiated :class:`DocumentArrayRedis` object
+        """
+        return super().__new__(cls)
diff --git a/docarray/array/storage/redis/__init__.py b/docarray/array/storage/redis/__init__.py
@@ -0,0 +1,12 @@
+from abc import ABC
+
+from .backend import BackendMixin, RedisConfig
+from .find import FindMixin
+from .getsetdel import GetSetDelMixin
+from .seqlike import SequenceLikeMixin
+
+__all__ = ['StorageMixins', 'RedisConfig']
+
+
+class StorageMixins(FindMixin, BackendMixin, GetSetDelMixin, SequenceLikeMixin, ABC):
+    ...
diff --git a/docarray/array/storage/redis/backend.py b/docarray/array/storage/redis/backend.py
@@ -0,0 +1,180 @@
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, Union
+
+import numpy as np
+from docarray import Document
+from docarray.array.storage.base.backend import BaseBackendMixin, TypeMap
+from docarray.helper import dataclass_from_dict
+
+from redis import Redis
+from redis.commands.search.field import NumericField, TextField, VectorField
+from redis.commands.search.indexDefinition import IndexDefinition
+
+if TYPE_CHECKING:
+    from docarray.typing import ArrayType, DocumentArraySourceType
+
+
+@dataclass
+class RedisConfig:
+    n_dim: int
+    host: str = field(default='localhost')
+    port: int = field(default=6379)
+    index_name: str = field(default='idx')
+    flush: bool = field(default=False)
+    update_schema: bool = field(default=True)
+    distance: str = field(default='COSINE')
+    redis_config: Dict[str, Any] = field(default_factory=dict)
+    batch_size: int = field(default=64)
+    method: str = field(default='HNSW')
+    ef_construction: int = field(default=200)
+    m: int = field(default=16)
+    ef_runtime: int = field(default=10)
+    block_size: int = field(default=1048576)
+    initial_cap: Optional[int] = None
+    columns: Optional[List[Tuple[str, str]]] = None
+
+
+class BackendMixin(BaseBackendMixin):
+    """Provide necessary functions to enable this storage backend."""
+
+    TYPE_MAP = {
+        'str': TypeMap(type='text', converter=TextField),
+        'bytes': TypeMap(type='text', converter=TextField),
+        'int': TypeMap(type='integer', converter=NumericField),
+        'float': TypeMap(type='float', converter=NumericField),
+        'double': TypeMap(type='double', converter=NumericField),
+        'long': TypeMap(type='long', converter=NumericField),
+        'bool': TypeMap(type='long', converter=NumericField),
+    }
+
+    def _init_storage(
+        self,
+        _docs: Optional['DocumentArraySourceType'] = None,
+        config: Optional[Union[RedisConfig, Dict]] = None,
+        **kwargs,
+    ):
+        if not config:
+            raise ValueError('Empty config is not allowed for Redis storage')
+        elif isinstance(config, dict):
+            config = dataclass_from_dict(RedisConfig, config)
+
+        if config.distance not in ['L2', 'IP', 'COSINE']:
+            raise ValueError(
+                f'Expecting distance metric one of COSINE, L2 OR IP, got {config.distance} instead'
+            )
+        if config.method not in ['HNSW', 'FLAT']:
+            raise ValueError(
+                f'Expecting search method one of HNSW OR FLAT, got {config.method} instead'
+            )
+
+        if config.redis_config.get('decode_responses'):
+            config.redis_config['decode_responses'] = False
+
+        self._offset2id_key = config.index_name + '__offset2id'
+        self._config = config
+        self.n_dim = self._config.n_dim
+        self._doc_prefix = config.index_name + ':'
+        self._config.columns = self._normalize_columns(self._config.columns)
+
+        self._client = self._build_client()
+        super()._init_storage()
+
+        if _docs is None:
+            return
+        elif isinstance(_docs, Iterable):
+            self.extend(_docs)
+        elif isinstance(_docs, Document):
+            self.append(_docs)
+
+    def _build_client(self):
+        client = Redis(
+            host=self._config.host,
+            port=self._config.port,
+            **self._config.redis_config,
+        )
+
+        if self._config.flush:
+            client.flushdb()
+
+        if self._config.update_schema:
+            if self._config.index_name.encode() in client.execute_command('FT._LIST'):
+                client.ft(index_name=self._config.index_name).dropindex()
+
+        if self._config.flush or self._config.update_schema:
+            schema = self._build_schema_from_redis_config()
+            idef = IndexDefinition(prefix=[self._doc_prefix])
+            client.ft(index_name=self._config.index_name).create_index(
+                schema, definition=idef
+            )
+
+        return client
+
+    def _ensure_unique_config(
+        self,
+        config_root: dict,
+        config_subindex: dict,
+        config_joined: dict,
+        subindex_name: str,
+    ) -> dict:
+        if 'index_name' not in config_subindex:
+            config_joined['index_name'] = (
+                config_joined['index_name'] + '_subindex_' + subindex_name
+            )
+        config_joined['flush'] = False
+        return config_joined
+
+    def _build_schema_from_redis_config(self):
+        index_param = {
+            'TYPE': 'FLOAT32',
+            'DIM': self.n_dim,
+            'DISTANCE_METRIC': self._config.distance,
+        }
+
+        if self._config.method == 'HNSW':
+            index_options = {
+                'M': self._config.m,
+                'EF_CONSTRUCTION': self._config.ef_construction,
+                'EF_RUNTIME': self._config.ef_runtime,
+            }
+            index_param.update(index_options)
+
+        if self._config.method == 'FLAT':
+            index_options = {'BLOCK_SIZE': self._config.block_size}
+            index_param.update(index_options)
+
+        if self._config.initial_cap:
+            index_param['INITIAL_CAP'] = self._config.initial_cap
+        schema = [VectorField('embedding', self._config.method, index_param)]
+
+        for col, coltype in self._config.columns:
+            schema.append(self._map_column(col, coltype))
+
+        return schema
+
+    def _doc_id_exists(self, doc_id):
+        return self._client.exists(self._doc_prefix + doc_id)
+
+    def _map_embedding(self, embedding: 'ArrayType') -> bytes:
+        if embedding is not None:
+            from docarray.math.ndarray import to_numpy_array
+
+            embedding = to_numpy_array(embedding)
+
+            if embedding.ndim > 1:
+                embedding = np.asarray(embedding).squeeze()
+        else:
+            embedding = np.zeros(self.n_dim)
+        return embedding.astype(np.float32).tobytes()
+
+    def _get_offset2ids_meta(self) -> List[str]:
+        if not self._client.exists(self._offset2id_key):
+            return []
+        ids = self._client.lrange(self._offset2id_key, 0, -1)
+        return [id.decode() for id in ids]
+
+    def _update_offset2ids_meta(self):
+        """Update the offset2ids in redis"""
+        if self._client.exists(self._offset2id_key):
+            self._client.delete(self._offset2id_key)
+        if len(self._offset2ids.ids) > 0:
+            self._client.rpush(self._offset2id_key, *self._offset2ids.ids)
diff --git a/docarray/array/storage/redis/find.py b/docarray/array/storage/redis/find.py
@@ -0,0 +1,127 @@
+from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, TypeVar, Union
+
+import numpy as np
+from docarray import Document, DocumentArray
+from docarray.array.mixins.find import FindMixin as BaseFindMixin
+from docarray.math import ndarray
+from docarray.math.ndarray import to_numpy_array
+from docarray.score import NamedScore
+
+from redis.commands.search.query import NumericFilter, Query
+
+if TYPE_CHECKING:
+    import tensorflow
+    import torch
+
+    RedisArrayType = TypeVar(
+        'RedisArrayType',
+        np.ndarray,
+        tensorflow.Tensor,
+        torch.Tensor,
+        Sequence[float],
+        Dict,
+    )
+
+
+class FindMixin(BaseFindMixin):
+    def _find_similar_vectors(
+        self,
+        query: 'RedisArrayType',
+        filter: Optional[Dict] = None,
+        limit: Optional[Union[int, float]] = 20,
+        **kwargs,
+    ):
+
+        query_str = self._build_query_str(filter) if filter else "*"
+
+        q = (
+            Query(f'{query_str}=>[KNN {limit} @embedding $vec AS vector_score]')
+            .sort_by('vector_score')
+            .paging(0, limit)
+            .dialect(2)
+        )
+
+        query_params = {'vec': to_numpy_array(query).astype(np.float32).tobytes()}
+        results = (
+            self._client.ft(index_name=self._config.index_name)
+            .search(q, query_params)
+            .docs
+        )
+
+        da = DocumentArray()
+        for res in results:
+            doc = Document.from_base64(res.blob.encode())
+            doc.scores['score'] = NamedScore(value=res.vector_score)
+            da.append(doc)
+        return da
+
+    def _find(
+        self,
+        query: 'RedisArrayType',
+        limit: Optional[Union[int, float]] = 20,
+        filter: Optional[Dict] = None,
+        **kwargs,
+    ) -> List['DocumentArray']:
+
+        query = np.array(query)
+        num_rows, n_dim = ndarray.get_array_rows(query)
+        if n_dim != 2:
+            query = query.reshape((num_rows, -1))
+
+        return [
+            self._find_similar_vectors(q, filter=filter, limit=limit, **kwargs)
+            for q in query
+        ]
+
+    def _find_with_filter(self, filter: Dict, limit: Optional[Union[int, float]] = 20):
+        s = self._build_query_str(filter)
+        q = Query(s)
+        q.paging(0, limit)
+
+        results = self._client.ft(index_name=self._config.index_name).search(q).docs
+
+        da = DocumentArray()
+        for res in results:
+            doc = Document.from_base64(res.blob.encode())
+            da.append(doc)
+        return da
+
+    def _filter(
+        self, filter: Dict, limit: Optional[Union[int, float]] = 20
+    ) -> 'DocumentArray':
+
+        return self._find_with_filter(filter, limit=limit)
+
+    def _build_query_str(self, filter: Dict) -> str:
+        INF = "+inf"
+        NEG_INF = "-inf"
+        s = "("
+
+        for key in filter:
+            operator = list(filter[key].keys())[0]
+            value = filter[key][operator]
+            if operator == '$gt':
+                s += f"@{key}:[({value} {INF}] "
+            elif operator == '$gte':
+                s += f"@{key}:[{value} {INF}] "
+            elif operator == '$lt':
+                s += f"@{key}:[{NEG_INF} ({value}] "
+            elif operator == '$lte':
+                s += f"@{key}:[{NEG_INF} {value}] "
+            elif operator == '$eq':
+                if type(value) is int:
+                    s += f"@{key}:[{value} {value}] "
+                elif type(value) is bool:
+                    s += f"@{key}:[{int(value)} {int(value)}] "
+                else:
+                    s += f"@{key}:{value} "
+            elif operator == '$ne':
+                if type(value) is int:
+                    s += f"-@{key}:[{value} {value}] "
+                elif type(value) is bool:
+                    s += f"-@{key}:[{int(value)} {int(value)}] "
+                else:
+                    s += f"-@{key}:{value} "
+        s += ")"
+
+        return s