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

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 39 additions & 3 deletions docarray/array/storage/weaviate/find.py
Original file line number Diff line number Diff line change
Expand Up @@ -58,9 +58,7 @@ def _find_similar_vectors(
f'find failed, please check your filter query. Errors: \n{errors}'
)

found_results = (
results.get('data', {}).get('Get', {}).get(self._class_name, []) or []
)
found_results = results.get('data', {}).get('Get', {}).get(self._class_name, [])

# The serialized document is stored in results['data']['Get'][self._class_name]
for result in found_results:
Expand All @@ -80,6 +78,44 @@ def _find_similar_vectors(

return DocumentArray(docs)

def _filter(
self,
filter: Dict,
) -> 'DocumentArray':
"""Returns a subset of documents by filtering by the given filter (Weaviate `where` filter).

:param filter: the input filter to apply in each stored document
:return: a `DocumentArray` containing the `Document` objects that verify the filter.
"""
if not filter:
return self

results = (
self._client.query.get(self._class_name, '_serialized')
.with_additional('id')
.with_where(filter)
.do()
)

docs = []
if 'errors' in results:
errors = '\n'.join(map(lambda error: error['message'], results['errors']))
raise ValueError(
f'filter failed, please check your filter query. Errors: \n{errors}'
)

found_results = results.get('data', {}).get('Get', {}).get(self._class_name, [])

# The serialized document is stored in results['data']['Get'][self._class_name]
for result in found_results:
doc = Document.from_base64(result['_serialized'], **self._serialize_config)

doc.tags['wid'] = result['_additional']['id']

docs.append(doc)

return DocumentArray(docs)

def _find(
self,
query: 'WeaviateArrayType',
Expand Down
59 changes: 57 additions & 2 deletions docs/advanced/document-store/weaviate.md
Original file line number Diff line number Diff line change
Expand Up @@ -169,13 +169,66 @@ print(results[0].text)
Persist Documents with Weaviate.
```

(weaviate-filter)=
## Vector search with filter

Search with `.find` can be restricted by user-defined filters. Such filters can be constructed following the guidelines
in [Weaviate's Documentation](https://weaviate.io/developers/weaviate/current/graphql-references/filters.html).

### Example of `.find` with a filter
### Example of `.find` with a filter only

Consider you store Documents with a certain tag `price` into weaviate and you want to retrieve all Documents
with `price` lower or equal to some `max_price` value.


You can index such Documents as follows:
```python
from docarray import Document, DocumentArray
import numpy as np

n_dim = 3
da = DocumentArray(
storage='weaviate',
config={
'n_dim': n_dim,
'columns': [('price', 'float')],
},
)

with da:
da.extend([Document(id=f'r{i}', tags={'price': i}) for i in range(10)])

print('\nIndexed Embeddings:\n')
for price in da[:, 'tags__price']:
print(f'\t price={price}')
```

Then you can retrieve all documents whose price is lower than or equal to `max_price` by applying the following
filter:

```python
max_price = 3
n_limit = 4

filter = {'path': 'price', 'operator': 'LessThanEqual', 'valueNumber': max_price}
results = da.find(filter=filter)

print('\n Returned examples that verify filter "price at most 7":\n')
for price in results[:, 'tags__price']:
print(f'\t price={price}')
```

This would print

```
Returned examples that satisfy condition "price at most 3":

price=0
price=1
price=2
price=3
```

### Example of `.find` with query vector and filter

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:
Expand Down Expand Up @@ -236,3 +289,5 @@ Embeddings Nearest Neighbours with "price" at most 7:
embedding=[5. 5. 5.], price=5
embedding=[4. 4. 4.], price=4
```


44 changes: 44 additions & 0 deletions tests/unit/array/mixins/test_find.py
Original file line number Diff line number Diff line change
Expand Up @@ -310,6 +310,45 @@ def test_search_pre_filtering(
)


@pytest.mark.parametrize(
'storage,filter_gen,numeric_operators,operator',
[
*[
tuple(
[
'weaviate',
lambda operator, threshold: {
'path': ['price'],
'operator': operator,
'valueNumber': threshold,
},
numeric_operators_weaviate,
operator,
]
)
for operator in numeric_operators_weaviate.keys()
]
],
)
def test_filtering(storage, filter_gen, operator, numeric_operators, start_storage):
n_dim = 128
da = DocumentArray(
storage=storage, config={'n_dim': n_dim, 'columns': [('price', 'float')]}
)

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.find(filter=filter)

assert all(
[numeric_operators[operator](r.tags['price'], threshold) for r in results]
)


def test_weaviate_filter_query(start_storage):
n_dim = 128
da = DocumentArray(
Expand All @@ -326,6 +365,11 @@ def test_weaviate_filter_query(start_storage):
with pytest.raises(ValueError):
da.find(np.random.rand(n_dim), filter={'wrong': 'filter'})

with pytest.raises(ValueError):
da._filter(filter={'wrong': 'filter'})

assert isinstance(da._filter(filter={}), type(da))


@pytest.mark.parametrize('storage', ['memory', 'elasticsearch'])
def test_unsupported_pre_filtering(storage, start_storage):
Expand Down