diff --git a/docarray/index/backends/in_memory.py b/docarray/index/backends/in_memory.py index 8a8132c2394..313f353b1c3 100644 --- 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 index a3b1419baa5..db375e07119 100644 --- 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):