feat: InMemoryExactNNIndex pre filtering (#1713)

jupyterjazz · web-flow · commit c96707a133e2 · 2023-07-19T13:20:10.000+02:00
Signed-off-by: jupyterjazz &lt;saba.sturua@jina.ai&gt;
diff --git a/docarray/index/backends/in_memory.py b/docarray/index/backends/in_memory.py
@@ -22,10 +22,7 @@
 from docarray.array.any_array import AnyDocArray
 from docarray.helper import _shallow_copy_doc
 from docarray.index.abstract import BaseDocIndex, _raise_not_supported
-from docarray.index.backends.helper import (
-    _collect_query_args,
-    _execute_find_and_filter_query,
-)
+from docarray.index.backends.helper import _collect_query_args
 from docarray.typing import AnyTensor, NdArray
 from docarray.typing.tensor.abstract_tensor import AbstractTensor
 from docarray.utils._internal._typing import safe_issubclass
@@ -293,12 +290,44 @@ def execute_query(self, query: List[Tuple[str, Dict]], *args, **kwargs) -> Any:
             raise ValueError(
                 f'args and kwargs not supported for `execute_query` on {type(self)}'
             )
-        find_res = _execute_find_and_filter_query(
-            doc_index=self,
-            query=query,
-            reverse_order=True,
-        )
-        return find_res
+        return self._find_and_filter(query)
+
+    def _find_and_filter(self, query: List[Tuple[str, Dict]]) -> FindResult:
+        """
+        The function executes search operations such as 'find' and 'filter' in the order
+        they appear in the query. The 'find' operation performs a vector similarity search.
+        The 'filter' operation filters out documents based on a filter query.
+        The documents are finally sorted based on their scores.
+
+        :param query: The query to execute.
+        :return: A tuple of retrieved documents and their scores.
+        """
+        out_docs = self._docs
+        doc_to_score: Dict[BaseDoc, Any] = {}
+        for op, op_kwargs in query:
+            if op == 'find':
+                out_docs, scores = find(
+                    index=out_docs,
+                    query=op_kwargs['query'],
+                    search_field=op_kwargs['search_field'],
+                    limit=op_kwargs.get('limit', len(out_docs)),
+                    metric=self._column_infos[op_kwargs['search_field']].config[
+                        'space'
+                    ],
+                )
+                doc_to_score.update(zip(out_docs.id, scores))
+            elif op == 'filter':
+                out_docs = filter_docs(out_docs, op_kwargs['filter_query'])
+                if 'limit' in op_kwargs:
+                    out_docs = out_docs[: op_kwargs['limit']]
+            else:
+                raise ValueError(f'Query operation is not supported: {op}')
+
+        scores_and_docs = zip([doc_to_score[doc.id] for doc in out_docs], out_docs)
+        sorted_lists = sorted(scores_and_docs, reverse=True)
+        out_scores, out_docs = zip(*sorted_lists)
+
+        return FindResult(documents=out_docs, scores=out_scores)
 
     def find(
         self,
diff --git a/tests/index/in_memory/test_in_memory.py b/tests/index/in_memory/test_in_memory.py
@@ -15,7 +15,6 @@
 tf_available = is_tf_available()
 if tf_available:
     import tensorflow as tf
-    from docarray.typing import TensorFlowTensor
 
 
 class SchemaDoc(BaseDoc):
@@ -165,37 +164,49 @@ def test_with_text_doc_torch():
         assert len(r) == 5
 
 
-def test_concatenated_queries(doc_index):
-    query = SchemaDoc(text='query', price=0, tensor=np.ones(10))
-
+def test_query_builder_pre_filtering(doc_index):
     q = (
         doc_index.build_query()
-        .find(query=query, search_field='tensor', limit=5)
-        .filter(filter_query={'price': {'$neq': 5}})
+        .filter(filter_query={'price': {'$lte': 3}})
+        .find(query=np.ones(10), search_field='tensor', limit=5)
         .build()
     )
 
     docs, scores = doc_index.execute_query(q)
 
     assert len(docs) == 4
+    for doc in docs:
+        assert doc.price <= 3
 
 
-@pytest.mark.parametrize(
-    'find_limit, filter_limit, expected_docs', [(10, 3, 3), (5, None, 1)]
-)
-def test_query_builder_limits(doc_index, find_limit, filter_limit, expected_docs):
-    query = SchemaDoc(text='query', price=3, tensor=np.array([3] * 10))
+def test_query_builder_post_filtering(doc_index):
+    q = (
+        doc_index.build_query()
+        .find(query=np.ones(10), search_field='tensor')
+        .filter(filter_query={'price': {'$gt': 3}}, limit=5)
+        .build()
+    )
 
+    docs, scores = doc_index.execute_query(q)
+
+    assert len(docs) == 5
+    for doc in docs:
+        assert doc.price > 3
+
+
+def test_query_builder_pre_post_filtering(doc_index):
     q = (
         doc_index.build_query()
-        .find(query=query, search_field='tensor', limit=find_limit)
-        .filter(filter_query={'price': {'$lte': 5}}, limit=filter_limit)
+        .filter(filter_query={'price': {'$lte': 3}})
+        .find(query=np.ones(10), search_field='tensor')
+        .filter(filter_query={'text': {'$eq': 'hello 1'}})
         .build()
     )
 
     docs, scores = doc_index.execute_query(q)
 
-    assert len(docs) == expected_docs
+    assert len(docs) == 1
+    assert docs[0].text == 'hello 1' and docs[0].price <= 3
 
 
 def test_filter(doc_index):
