from typing import (
Union,
Optional,
TYPE_CHECKING,
List,
Dict,
)
if TYPE_CHECKING: # pragma: no cover
import numpy as np
from docarray import DocumentArray
[docs]class FindMixin:
def _find(
self,
query: 'np.ndarray',
limit: Optional[Union[int, float]] = 20,
only_id: bool = False,
filter: Optional[Dict] = None,
**kwargs,
) -> List['DocumentArray']:
"""Returns approximate nearest neighbors given an input query.
:param query: the query documents to search.
:param limit: the number of results to get for each query document in search.
:param only_id: if set, then returning matches will only contain ``id``
:param filter: filter query used for pre-filtering
:param kwargs: other kwargs.
:return: a list of DocumentArrays containing the closest Document objects for each of the queries in `query`.
"""
from docarray.math import ndarray
n_rows, _ = ndarray.get_array_rows(query)
if n_rows == 1:
query = query.reshape(1, -1)
_, match_docs = self._annlite.search_by_vectors(
query, limit=limit, filter=filter or {}, include_metadata=not only_id
)
return match_docs
def _filter(
self,
filter: Dict,
limit: Optional[Union[int, float]] = 20,
only_id: bool = False,
) -> 'DocumentArray':
"""Returns a subset of documents by filtering by the given filter (`Annlite` filter).
:param filter: the input filter to apply in each stored document
:param limit: the number of results to get for each query document in search.
:param only_id: if set, then returning matches will only contain ``id``
:return: a `DocumentArray` containing the `Document` objects that verify the filter.
"""
docs = self._annlite.filter(
filter=filter,
limit=limit,
include_metadata=not only_id,
)
return DocumentArray(docs)
