feat: redis supports geo filter (#579)

AnneYang720 · dong xiang · commit 31da66bd63dd · 2022-10-20T14:47:10.000+08:00
diff --git a/docarray/array/storage/qdrant/find.py b/docarray/array/storage/qdrant/find.py
@@ -4,10 +4,10 @@
     TypeVar,
     Sequence,
     List,
-    Dict,
+    Union,
     Optional,
+    Dict,
 )
-
 from qdrant_client.http.models.models import Distance
 
 from docarray import Document, DocumentArray
@@ -103,3 +103,32 @@ def _find(
                 da = self._find_similar_vectors(q, limit=limit, filter=filter)
                 closest_docs.append(da)
             return closest_docs
+
+    def _find_with_filter(
+        self, filter: Optional[Dict], limit: Optional[Union[int, float]] = 10
+    ):
+        list_of_points, _offset = self.client.scroll(
+            collection_name=self.collection_name,
+            scroll_filter=filter,
+            with_payload=True,
+            limit=limit,
+        )
+        da = DocumentArray()
+        for result in list_of_points[:limit]:
+            doc = Document.from_base64(
+                result.payload['_serialized'], **self.serialize_config
+            )
+            da.append(doc)
+        return da
+
+    def _filter(
+        self, filter: Optional[Dict], limit: Optional[Union[int, float]] = 10
+    ) -> 'DocumentArray':
+        """Returns a subset of documents by filtering by the given filter (`Qdrant` filter)..
+        :param limit: number of retrieved items
+        :param filter: filter query used for filtering.
+        For more information: https://docarray.jina.ai/advanced/document-store/qdrant/#qdrant
+        :return: a `DocumentArray` containing the `Document` objects that verify the filter.
+        """
+
+        return self._find_with_filter(filter, limit=limit)
diff --git a/docarray/array/storage/redis/backend.py b/docarray/array/storage/redis/backend.py
@@ -8,7 +8,7 @@
 from docarray.helper import dataclass_from_dict, random_identity, filter_dict
 
 from redis import Redis
-from redis.commands.search.field import NumericField, TextField, VectorField, GeoField
+from redis.commands.search.field import NumericField, TextField, VectorField
 from redis.commands.search.indexDefinition import IndexDefinition
 
 if TYPE_CHECKING:
@@ -46,7 +46,7 @@ class BackendMixin(BaseBackendMixin):
         'float': TypeMap(type='float', converter=NumericField),
         'double': TypeMap(type='double', converter=NumericField),
         'long': TypeMap(type='long', converter=NumericField),
-        'geo': TypeMap(type='geo', converter=GeoField),
+        'bool': TypeMap(type='long', converter=NumericField),
     }
 
     def _init_storage(
diff --git a/docarray/array/storage/redis/find.py b/docarray/array/storage/redis/find.py
@@ -1,4 +1,3 @@
-import warnings
 from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, TypeVar, Union
 
 import numpy as np
@@ -39,13 +38,14 @@ class FindMixin(BaseFindMixin):
     def _find_similar_vectors(
         self,
         query: 'RedisArrayType',
-        filter: Optional[Union[str, Dict]] = None,
+        filter: Optional[Dict] = None,
         limit: Union[int, float] = 20,
         **kwargs,
     ):
 
         if filter:
-            query_str = _get_redis_filter_query(filter)
+            nodes = _build_query_nodes(filter)
+            query_str = intersect(*nodes).to_string()
         else:
             query_str = '*'
 
@@ -74,7 +74,7 @@ def _find(
         self,
         query: 'RedisArrayType',
         limit: Union[int, float] = 20,
-        filter: Optional[Union[str, Dict]] = None,
+        filter: Optional[Dict] = None,
         **kwargs,
     ) -> List['DocumentArray']:
 
@@ -90,10 +90,11 @@ def _find(
 
     def _find_with_filter(
         self,
-        filter: Union[str, Dict],
+        filter: Dict,
         limit: Union[int, float] = 20,
     ):
-        query_str = _get_redis_filter_query(filter)
+        nodes = _build_query_nodes(filter)
+        query_str = intersect(*nodes).to_string()
         q = Query(query_str)
         q.paging(0, limit)
 
@@ -107,7 +108,7 @@ def _find_with_filter(
 
     def _filter(
         self,
-        filter: Union[str, Dict],
+        filter: Dict,
         limit: Union[int, float] = 20,
     ) -> 'DocumentArray':
 
@@ -217,19 +218,3 @@ def _build_query_nodes(filter):
 def _build_query_str(query):
     query_str = '|'.join(query.split(' '))
     return query_str
-
-
-def _get_redis_filter_query(filter: Union[str, Dict]):
-    if isinstance(filter, dict):
-        warnings.warn(
-            "Dict syntax for redis filter will be deprecated, use string literals instead",
-            DeprecationWarning,
-        )
-        nodes = _build_query_nodes(filter)
-        query_str = intersect(*nodes).to_string()
-    elif isinstance(filter, str):
-        query_str = filter
-    else:
-        raise ValueError(f'Unexpected type of filter: {type(filter)}, expected str')
-
-    return query_str
diff --git a/docs/advanced/document-store/qdrant.md b/docs/advanced/document-store/qdrant.md
@@ -61,6 +61,7 @@ da = DocumentArray(
         'host': 'localhost',
         'port': '6333',
         'n_dim': 10,
+        'distance': "cosine",
     },
 )
 
@@ -98,7 +99,7 @@ Create `docker-compose.yml`:
 version: '3.4'
 services:
   qdrant:
-    image: qdrant/qdrant:v0.7.0
+    image: qdrant/qdrant:v0.8.0
     ports:
       - "6333:6333"
     ulimits: # Only required for tests, as there are a lot of collections created
@@ -121,7 +122,9 @@ from docarray import DocumentArray
 
 N, D = 100, 128
 
-da = DocumentArray.empty(N, storage='qdrant', config={'n_dim': D})  # init
+da = DocumentArray.empty(
+    N, storage='qdrant', config={'n_dim': D, 'distance': 'cosine'}
+)  # init
 
 da.embeddings = np.random.random([N, D])
 
@@ -146,7 +149,7 @@ in [Qdrant's Documentation](https://qdrant.tech/documentation/filtering/)
 Consider Documents with embeddings `[0,0,0]` up to ` [9,9,9]` where the document with embedding `[i,i,i]`
 has as tag `price` with value `i`. We can create such example with the following code:
 
-```python
+```pyt
 from docarray import Document, DocumentArray
 import numpy as np
 
@@ -175,7 +178,7 @@ for embedding, price in zip(da.embeddings, da[:, 'tags__price']):
 
 Consider we want the nearest vectors to the embedding `[8. 8. 8.]`, with the restriction that
 prices must follow a filter. As an example, let's consider that retrieved documents must have `price` value lower
-or equal than `max_price`. We can encode this information in annlite using `filter = {'price': {'$lte': max_price}}`.
+or equal than `max_price`. We can encode this information in qdrant using `filter = {'price': {'$lte': max_price}}`.
 
 Then the search with the proposed filter can be implemented and used with the following code:
 
@@ -206,3 +209,58 @@ Embeddings Nearest Neighbours with "price" at most 7:
 	embedding=[5. 5. 5.],	 price=5
 	embedding=[4. 4. 4.],	 price=4
 ```
+### Example of `.filter` with a filter
+Consider Documents have tags `price` with value `i`. We can create such example with the following code:
+```python
+from docarray import Document, DocumentArray
+import numpy as np
+
+n_dim = 3
+distance = 'euclidean'
+
+da = DocumentArray(
+    storage='qdrant',
+    config={'n_dim': n_dim, 'columns': {'price': 'float'}, 'distance': distance},
+)
+
+print(f'\nDocumentArray distance: {distance}')
+
+with da:
+    da.extend(
+        [
+            Document(id=f'r{i}', embedding=i * np.ones(n_dim), tags={'price': i})
+            for i in range(10)
+        ]
+    )
+
+print('\nIndexed Prices:\n')
+for embedding, price in zip(da.embeddings, da[:, 'tags__price']):
+    print(f'\tembedding={embedding},\t price={price}')
+```
+Consider we want the vectors with the restriction that prices must follow a filter. As an example, 
+let's consider that retrieved documents must have `price` value lower or equal than `max_price`. We can encode 
+this information in qdrant using `filter = {'price': {'$lte': max_price}}`.
+
+Then the search with the proposed filter can be implemented and used with the following code:
+```python
+max_price = 7
+n_limit = 4
+
+filter = {'must': [{'key': 'price', 'range': {'lte': max_price}}]}
+results = da.filter(filter=filter, limit=n_limit)
+
+print('\nPoints with "price" at most 7:\n')
+for embedding, price in zip(results.embeddings, results[:, 'tags__price']):
+    print(f'\tembedding={embedding},\t price={price}')
+```
+This would print:
+
+```
+
+Points with "price" at most 7:
+
+	embedding=[6. 6. 6.],	 price=6
+	embedding=[7. 7. 7.],	 price=7
+	embedding=[1. 1. 1.],	 price=1
+	embedding=[2. 2. 2.],	 price=2
+```
diff --git a/tests/unit/array/mixins/test_find.py b/tests/unit/array/mixins/test_find.py
@@ -545,6 +545,42 @@ def test_filtering(
         )
 
 
+@pytest.mark.parametrize(
+    'storage,filter_gen,numeric_operators,operator',
+    [
+        *[
+            tuple(
+                [
+                    'qdrant',
+                    lambda operator, threshold: {
+                        'must': [{'key': 'price', 'match': {'value': threshold}}]
+                    },
+                    numeric_operators_qdrant,
+                    'eq',
+                ]
+            )
+        ],
+    ],
+)
+@pytest.mark.parametrize('columns', [[('price', 'int')], {'price': 'int'}])
+def test_qdrant_filter_function(
+    storage, filter_gen, operator, numeric_operators, start_storage, columns
+):
+    n_dim = 128
+    da = DocumentArray(storage='qdrant', config={'n_dim': n_dim, 'columns': columns})
+    da.extend([Document(id=f'r{i}', tags={'price': i}) for i in range(50)])
+    thresholds = [10, 20, 30]
+    for threshold in thresholds:
+        filter = filter_gen(operator, threshold)
+        results = da._filter(filter=filter)
+
+        assert len(results) > 0
+
+        assert all(
+            [numeric_operators[operator](r.tags['price'], threshold) for r in results]
+        )
+
+
 @pytest.mark.parametrize('columns', [[('price', 'int')], {'price': 'int'}])
 def test_weaviate_filter_query(start_storage, columns):
     n_dim = 128
