Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Appearance settings

Commit ea2a7a8

Browse filesBrowse files
AnneYang720AnneYang720
authored
feat(redis): add full-text search and io (#535)
* feat: redis support full-text search * docs: add redis full-text search doc * feat: add tag_indices to redis getsetdel * refactor: black redis files * feat: redis supports io * fix: fix test_push_pull_io for redis * refactor: code minor adjustments * docs: default scorer function in redis text search * feat: make redis scoer parameter configurable * docs: redis doc error fix * test: add test for scorer of redis text search
1 parent 638f362 commit ea2a7a8
Copy full SHA for ea2a7a8

6 files changed

+239-10Lines changed: 239 additions & 10 deletions

File tree

Expand file treeCollapse file tree
Open diff view settings
Filter options
Expand file treeCollapse file tree
Open diff view settings
Collapse file

‎docarray/array/storage/redis/backend.py‎

Copy file name to clipboardExpand all lines: docarray/array/storage/redis/backend.py
+21Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@ class RedisConfig:
2424
update_schema: bool = field(default=True)
2525
distance: str = field(default='COSINE')
2626
redis_config: Dict[str, Any] = field(default_factory=dict)
27+
index_text: bool = field(default=False)
28+
tag_indices: List[str] = field(default_factory=list)
2729
batch_size: int = field(default=64)
2830
method: str = field(default='HNSW')
2931
ef_construction: int = field(default=200)
@@ -146,6 +148,12 @@ def _build_schema_from_redis_config(self):
146148
index_param['INITIAL_CAP'] = self._config.initial_cap
147149
schema = [VectorField('embedding', self._config.method, index_param)]
148150

151+
if self._config.index_text:
152+
schema.append(TextField('text'))
153+
154+
for index in self._config.tag_indices:
155+
schema.append(TextField(index))
156+
149157
for col, coltype in self._config.columns.items():
150158
schema.append(self._map_column(col, coltype))
151159

@@ -178,3 +186,16 @@ def _update_offset2ids_meta(self):
178186
self._client.delete(self._offset2id_key)
179187
if len(self._offset2ids.ids) > 0:
180188
self._client.rpush(self._offset2id_key, *self._offset2ids.ids)
189+
190+
def __getstate__(self):
191+
d = dict(self.__dict__)
192+
del d['_client']
193+
return d
194+
195+
def __setstate__(self, state):
196+
self.__dict__ = state
197+
self._client = Redis(
198+
host=self._config.host,
199+
port=self._config.port,
200+
**self._config.redis_config,
201+
)
Collapse file

‎docarray/array/storage/redis/find.py‎

Copy file name to clipboardExpand all lines: docarray/array/storage/redis/find.py
+68-4Lines changed: 68 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ def _find_similar_vectors(
3939
self,
4040
query: 'RedisArrayType',
4141
filter: Optional[Dict] = None,
42-
limit: int = 20,
42+
limit: Union[int, float] = 20,
4343
**kwargs,
4444
):
4545

@@ -73,7 +73,7 @@ def _find_similar_vectors(
7373
def _find(
7474
self,
7575
query: 'RedisArrayType',
76-
limit: int = 20,
76+
limit: Union[int, float] = 20,
7777
filter: Optional[Dict] = None,
7878
**kwargs,
7979
) -> List['DocumentArray']:
@@ -88,7 +88,11 @@ def _find(
8888
for q in query
8989
]
9090

91-
def _find_with_filter(self, filter: Dict, limit: int = 20):
91+
def _find_with_filter(
92+
self,
93+
filter: Dict,
94+
limit: Union[int, float] = 20,
95+
):
9296
nodes = _build_query_nodes(filter)
9397
query_str = intersect(*nodes).to_string()
9498
q = Query(query_str)
@@ -102,10 +106,65 @@ def _find_with_filter(self, filter: Dict, limit: int = 20):
102106
da.append(doc)
103107
return da
104108

105-
def _filter(self, filter: Dict, limit: int = 20) -> 'DocumentArray':
109+
def _filter(
110+
self,
111+
filter: Dict,
112+
limit: Union[int, float] = 20,
113+
) -> 'DocumentArray':
106114

107115
return self._find_with_filter(filter, limit=limit)
108116

117+
def _find_by_text(
118+
self,
119+
query: Union[str, List[str]],
120+
index: str = 'text',
121+
limit: Union[int, float] = 20,
122+
**kwargs,
123+
):
124+
if isinstance(query, str):
125+
query = [query]
126+
127+
return [
128+
self._find_similar_documents_from_text(
129+
q,
130+
index=index,
131+
limit=limit,
132+
**kwargs,
133+
)
134+
for q in query
135+
]
136+
137+
def _find_similar_documents_from_text(
138+
self,
139+
query: str,
140+
index: str = 'text',
141+
limit: Union[int, float] = 20,
142+
**kwargs,
143+
):
144+
query_str = _build_query_str(query)
145+
scorer = kwargs.get('scorer', 'BM25')
146+
if scorer not in [
147+
'BM25',
148+
'TFIDF',
149+
'TFIDF.DOCNORM',
150+
'DISMAX',
151+
'DOCSCORE',
152+
'HAMMING',
153+
]:
154+
raise ValueError(
155+
f'Expecting a valid text similarity ranking algorithm, got {scorer} instead'
156+
)
157+
158+
q = Query(f'@{index}:{query_str}').scorer(scorer).paging(0, limit)
159+
160+
results = self._client.ft(index_name=self._config.index_name).search(q).docs
161+
162+
da = DocumentArray()
163+
for res in results:
164+
doc = Document.from_base64(res.blob.encode())
165+
da.append(doc)
166+
return da
167+
109168

110169
def _build_query_node(key, condition):
111170
operator = list(condition.keys())[0]
@@ -154,3 +213,8 @@ def _build_query_nodes(filter):
154213
nodes.append(child)
155214

156215
return nodes
216+
217+
218+
def _build_query_str(query):
219+
query_str = '|'.join(query.split(' '))
220+
return query_str
Collapse file

‎docarray/array/storage/redis/getsetdel.py‎

Copy file name to clipboardExpand all lines: docarray/array/storage/redis/getsetdel.py
+6Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,12 @@ def _document_to_redis(self, doc: 'Document') -> Dict:
9595
if tag is not None:
9696
extra_columns[col] = int(tag) if isinstance(tag, bool) else tag
9797

98+
if self._config.tag_indices:
99+
for index in self._config.tag_indices:
100+
text = doc.tags.get(index)
101+
if text is not None:
102+
extra_columns[index] = text
103+
98104
payload = {
99105
'id': doc.id,
100106
'embedding': self._map_embedding(doc.embedding),
Collapse file

‎docs/advanced/document-store/redis.md‎

Copy file name to clipboardExpand all lines: docs/advanced/document-store/redis.md
+113-2Lines changed: 113 additions & 2 deletions
  • Display the source diff
  • Display the rich diff
Original file line numberDiff line numberDiff line change
@@ -157,7 +157,7 @@ da.extend(
157157
Document(
158158
id=f'{i}',
159159
embedding=i * np.ones(n_dim),
160-
tags={'price': i, 'color': 'blue', 'stock': i%2==0},
160+
tags={'price': i, 'color': 'blue', 'stock': i % 2 == 0},
161161
)
162162
for i in range(10)
163163
]
@@ -167,7 +167,7 @@ da.extend(
167167
Document(
168168
id=f'{i+10}',
169169
embedding=i * np.ones(n_dim),
170-
tags={'price': i, 'color': 'red', 'stock': i%2==0},
170+
tags={'price': i, 'color': 'red', 'stock': i % 2 == 0},
171171
)
172172
for i in range(10)
173173
]
@@ -281,6 +281,117 @@ More example filter expresses
281281
}
282282
```
283283

284+
### Search by `.text` field
285+
286+
You can perform full-text search in a `DocumentArray` with `storage='redis'`.
287+
To do this, text needs to be indexed using the boolean flag `'index_text'` which is set when the `DocumentArray` is created with `config={'index_text': True, ...}`.
288+
The following example builds a `DocumentArray` with several documents containing text and searches for those that have `token1` in their text description.
289+
290+
```python
291+
from docarray import Document, DocumentArray
292+
293+
da = DocumentArray(
294+
storage='redis', config={'n_dim': 2, 'index_text': True, 'flush': True}
295+
)
296+
da.extend(
297+
[
298+
Document(id='1', text='token1 token2 token3'),
299+
Document(id='2', text='token1 token2'),
300+
Document(id='3', text='token2 token3 token4'),
301+
]
302+
)
303+
304+
results = da.find('token1')
305+
print(results[:, 'text'])
306+
```
307+
308+
This will print:
309+
310+
```console
311+
['token1 token2 token3', 'token1 token2']
312+
```
313+
314+
The default similarity ranking algorithm is `BM25`. Besides, `TFIDF`, `TFIDF.DOCNORM`, `DISMAX`, `DOCSCORE` and `HAMMING` are also supported by [RediSearch](https://redis.io/docs/stack/search/reference/scoring/). You can change it by specifying `scorer` in function `find`:
315+
316+
```python
317+
results = da.find('token1 token3', scorer='TFIDF.DOCNORM')
318+
print('scorer=TFIDF.DOCNORM:')
319+
print(results[:, 'text'])
320+
321+
results = da.find('token1 token3')
322+
print('scorer=BM25:')
323+
print(results[:, 'text'])
324+
```
325+
326+
This will print:
327+
328+
```console
329+
scorer=TFIDF.DOCNORM:
330+
['token1 token2', 'token1 token2 token3', 'token2 token3 token4']
331+
scorer=BM25:
332+
['token1 token2 token3', 'token1 token2', 'token2 token3 token4']
333+
```
334+
335+
### Search by `.tags` field
336+
337+
Text can also be indexed when it is part of `tags`.
338+
This is mostly useful in applications where text data can be split into groups and applications might require retrieving items based on a text search in an specific tag.
339+
340+
For example:
341+
342+
```python
343+
from docarray import Document, DocumentArray
344+
345+
da = DocumentArray(
346+
storage='redis',
347+
config={'n_dim': 32, 'flush': True, 'tag_indices': ['food_type', 'price']},
348+
)
349+
da.extend(
350+
[
351+
Document(
352+
tags={
353+
'food_type': 'Italian and Spanish food',
354+
'price': 'cheap but not that cheap',
355+
},
356+
),
357+
Document(
358+
tags={
359+
'food_type': 'French and Italian food',
360+
'price': 'on the expensive side',
361+
},
362+
),
363+
Document(
364+
tags={
365+
'food_type': 'chinese noddles',
366+
'price': 'quite cheap for what you get!',
367+
},
368+
),
369+
]
370+
)
371+
372+
results_cheap = da.find('cheap', index='price')
373+
print('searching "cheap" in <price>:\n\t', results_cheap[:, 'tags__price'])
374+
375+
results_italian = da.find('italian', index='food_type')
376+
print('searching "italian" in <food_type>:\n\t', results_italian[:, 'tags__food_type'])
377+
```
378+
379+
This will print:
380+
381+
```console
382+
searching "cheap" in <price>:
383+
['cheap but not that cheap', 'quite cheap for what you get!']
384+
searching "italian" in <food_type>:
385+
['French and Italian food', 'Italian and Spanish food']
386+
```
387+
388+
```{note}
389+
By default, if you don't specify the parameter `index` in the `find` method, the Document attribute `text` will be used for search. If you want to use a specific tags field, make sure to specify it with parameter `index`:
390+
```python
391+
results = da.find('cheap', index='price')
392+
```
393+
394+
284395
(vector-search-index)=
285396
### Update Vector Search Indexing Schema
286397

Collapse file

‎tests/unit/array/mixins/test_find.py‎

Copy file name to clipboardExpand all lines: tests/unit/array/mixins/test_find.py
+10-3Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,7 @@ def test_find(storage, config, limit, query, start_storage):
9999
'storage, config',
100100
[
101101
('elasticsearch', {'n_dim': 32, 'index_text': True}),
102+
('redis', {'n_dim': 32, 'flush': True, 'index_text': True}),
102103
],
103104
)
104105
def test_find_by_text(storage, config, start_storage):
@@ -111,7 +112,10 @@ def test_find_by_text(storage, config, start_storage):
111112
]
112113
)
113114

114-
results = da.find('token1')
115+
if storage == 'redis':
116+
results = da.find('token1', scorer='TFIDF')
117+
else:
118+
results = da.find('token1')
115119
assert isinstance(results, DocumentArray)
116120
assert len(results) == 2
117121
assert set(results[:, 'id']) == {'1', '2'}
@@ -140,6 +144,10 @@ def test_find_by_text(storage, config, start_storage):
140144
'storage, config',
141145
[
142146
('elasticsearch', {'n_dim': 32, 'tag_indices': ['attr1', 'attr2', 'attr3']}),
147+
(
148+
'redis',
149+
{'n_dim': 32, 'flush': True, 'tag_indices': ['attr1', 'attr2', 'attr3']},
150+
),
143151
],
144152
)
145153
def test_find_by_tag(storage, config, start_storage):
@@ -193,8 +201,7 @@ def test_find_by_tag(storage, config, start_storage):
193201

194202
results = da.find('token6', index='attr3')
195203
assert len(results) == 2
196-
assert results[0].id == '2'
197-
assert results[1].id == '1'
204+
assert set(results[:, 'id']) == {'1', '2'}
198205

199206
results = da.find('token6', index='attr3', limit=1)
200207
assert len(results) == 1

0 commit comments

Comments
0 (0)
Morty Proxy This is a proxified and sanitized view of the page, visit original site.