FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

feat: InMemoryExactNNIndex pre filtering by jupyterjazz · Pull Request #1713 · docarray/docarray · GitHub

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension .py  (2) All 1 file type selected
Viewed files
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Unified
Split
Hide whitespace
Diff view
Unified
Split
Hide whitespace
49 changes: 39 additions & 10 deletions docarray/index/backends/in_memory.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
39 changes: 25 additions & 14 deletions tests/index/in_memory/test_in_memory.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@
tf_available = is_tf_available()
if tf_available:
import tensorflow as tf
from docarray.typing import TensorFlowTensor


class SchemaDoc(BaseDoc):
Expand Down Expand Up @@ -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):
Expand Down

Back | FazBrowse Home | New Git URL