Skip to content
This repository was archived by the owner on Sep 9, 2026. It is now read-only.

Commit ecf5d23

Browse files
AnneYang720AnneYang720
authored
feat: support root_id for storage backends (#808)
1 parent 5fce6b6 commit ecf5d23

24 files changed

Lines changed: 280 additions & 36 deletions

File tree

‎docarray/array/mixins/find.py‎

Lines changed: 28 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,13 @@
11
import abc
2-
from typing import overload, Optional, Union, Dict, List, Tuple, Callable, TYPE_CHECKING
2+
from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Union, overload
33

44
import numpy as np
5-
65
from docarray.math import ndarray
76
from docarray.score import NamedScore
87

98
if TYPE_CHECKING: # pragma: no cover
10-
from docarray.typing import T, ArrayType
11-
129
from docarray import Document, DocumentArray
10+
from docarray.typing import ArrayType, T
1311

1412

1513
class FindMixin:
@@ -99,6 +97,7 @@ def find(
9997
filter: Union[Dict, str, None] = None,
10098
only_id: bool = False,
10199
index: str = 'text',
100+
return_root: Optional[bool] = False,
102101
on: Optional[str] = None,
103102
**kwargs,
104103
) -> Union['DocumentArray', List['DocumentArray']]:
@@ -126,14 +125,17 @@ def find(
126125
parameter is ignored. By default, the Document `text` attribute will be used for search,
127126
otherwise the tag field specified by `index` will be used. You can only use this parameter if the
128127
storage backend supports searching by text.
128+
:param return_root: if set, then the root-level DocumentArray will be returned
129129
:param on: specifies a subindex to search on. If set, the returned DocumentArray will be retrieved from the given subindex.
130130
:param kwargs: other kwargs.
131131
132132
:return: a list of DocumentArrays containing the closest Document objects for each of the queries in `query`.
133133
"""
134+
from docarray import Document, DocumentArray
135+
134136
index_da = self._get_index(subindex_name=on)
135137
if index_da is not self:
136-
return index_da.find(
138+
results = index_da.find(
137139
query,
138140
metric,
139141
limit,
@@ -144,7 +146,15 @@ def find(
144146
index,
145147
on=None,
146148
)
147-
from docarray import Document, DocumentArray
149+
150+
if return_root:
151+
da = self._get_root_docs(results)
152+
for d, s in zip(da, results[:, 'scores']):
153+
d.scores = s
154+
155+
return da
156+
157+
return results
148158

149159
if isinstance(query, dict):
150160
if filter is None:
@@ -301,3 +311,15 @@ def _find_by_text(self, *args, **kwargs):
301311
raise NotImplementedError(
302312
f'Search by text is not supported with this backend {self.__class__.__name__}'
303313
)
314+
315+
def _get_root_docs(self, docs: 'DocumentArray') -> 'DocumentArray':
316+
"""Get the root documents of the current DocumentArray.
317+
318+
:return: a `DocumentArray` containing the root documents.
319+
"""
320+
321+
if not all(docs[:, 'tags___root_id_']):
322+
raise ValueError(
323+
f'Not all Documents in this subindex have the "_root_id_" attribute set in all `tags`.'
324+
)
325+
return self[docs[:, 'tags___root_id_']]

‎docarray/array/mixins/setitem.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,10 @@ def __setitem__(
6363
index: 'DocumentArrayIndexType',
6464
value: Union['Document', Sequence['Document']],
6565
):
66+
from docarray.helper import check_root_id
67+
68+
if self._is_subindex:
69+
check_root_id(self, value)
6670

6771
self._update_subindices_set(index, value)
6872
# set by offset

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ class AnnliteConfig:
3131
max_connection: Optional[int] = None
3232
n_components: Optional[int] = None
3333
columns: Optional[Union[List[Tuple[str, str]], Dict[str, str]]] = None
34+
root_id: bool = True
3435

3536

3637
class BackendMixin(BaseBackendMixin):
@@ -104,7 +105,7 @@ def _init_storage(
104105

105106
self._annlite = AnnLite(self.n_dim, lock=False, **filter_dict(config))
106107

107-
super()._init_storage()
108+
super()._init_storage(**kwargs)
108109

109110
if _docs is None:
110111
return

‎docarray/array/storage/annlite/seqlike.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ def _extend(self, values: Iterable['Document']) -> None:
2020
self._offset2ids.extend([doc.id for doc in docs])
2121

2222
def _append(self, value: 'Document'):
23-
self.extend([value])
23+
self._extend([value])
2424

2525
def __eq__(self, other):
2626
"""In annlite backend, data are considered as identical if configs point to the same database source"""

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,11 @@ def _init_storage(
1717
self,
1818
_docs: Optional['DocumentArraySourceType'] = None,
1919
copy: bool = False,
20+
_is_subindex: bool = False,
2021
*args,
2122
**kwargs,
2223
):
24+
self._is_subindex = _is_subindex
2325
self._load_offset2ids()
2426

2527
def _init_subindices(
@@ -40,7 +42,9 @@ def _init_subindices(
4042
config_joined = self._ensure_unique_config(
4143
config, config_subindex, config_joined, name
4244
)
43-
self._subindices[name] = self.__class__(config=config_joined)
45+
self._subindices[name] = self.__class__(
46+
config=config_joined, _is_subindex=True
47+
)
4448
if _docs:
4549
from docarray import DocumentArray
4650

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

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -200,13 +200,25 @@ def _update_subindices_set(self, set_index, docs):
200200
_check_valid_values_nested_set(self[set_index], docs)
201201
if set_index in subindices:
202202
subindex_da = subindices[set_index]
203+
203204
subindex_da.clear()
204205
subindex_da.extend(docs)
205206
else: # root level set, update subindices iteratively
206207
for subindex_selector, subindex_da in subindices.items():
207208
old_ids = DocumentArray(self[set_index])[subindex_selector, 'id']
208209
del subindex_da[old_ids]
209-
subindex_da.extend(DocumentArray(docs)[subindex_selector])
210+
211+
value = DocumentArray(docs)
212+
213+
if (
214+
getattr(subindex_da, '_config', None) # checks if in-memory da
215+
and subindex_da._config.root_id
216+
):
217+
for v in value:
218+
for doc in DocumentArray(v)[subindex_selector]:
219+
doc.tags['_root_id_'] = v.id
220+
221+
subindex_da.extend(value[subindex_selector])
210222

211223
def _set_docs(self, ids, docs: Iterable['Document']):
212224
docs = list(docs)

‎docarray/array/storage/base/seqlike.py‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
1+
import warnings
12
from abc import abstractmethod
2-
from typing import Iterator, Iterable, MutableSequence
3+
from typing import Iterable, Iterator, MutableSequence
34

45
from docarray import Document, DocumentArray
56

@@ -10,7 +11,15 @@ class BaseSequenceLikeMixin(MutableSequence[Document]):
1011
def _update_subindices_append_extend(self, value):
1112
if getattr(self, '_subindices', None):
1213
for selector, da in self._subindices.items():
13-
docs_selector = DocumentArray(value)[selector]
14+
15+
value = DocumentArray(value)
16+
17+
if getattr(da, '_config', None) and da._config.root_id:
18+
for v in value:
19+
for doc in DocumentArray(v)[selector]:
20+
doc.tags['_root_id_'] = v.id
21+
22+
docs_selector = value[selector]
1423
if len(docs_selector) > 0:
1524
da.extend(docs_selector)
1625

@@ -63,6 +72,12 @@ def __bool__(self):
6372
return len(self) > 0
6473

6574
def extend(self, values: Iterable['Document'], **kwargs) -> None:
75+
76+
from docarray.helper import check_root_id
77+
78+
if self._is_subindex:
79+
check_root_id(self, values)
80+
6681
self._extend(values, **kwargs)
6782
self._update_subindices_append_extend(values)
6883

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ class ElasticConfig:
4646
ef_construction: Optional[int] = None
4747
m: Optional[int] = None
4848
columns: Optional[Union[List[Tuple[str, str]], Dict[str, str]]] = None
49+
root_id: bool = True
4950

5051

5152
_banned_indexname_chars = ['[', ' ', '"', '*', '\\', '<', '|', ',', '>', '/', '?', ']']
@@ -100,7 +101,7 @@ def _init_storage(
100101
self._build_offset2id_index()
101102

102103
# Note super()._init_storage() calls _load_offset2ids which calls _get_offset2ids_meta
103-
super()._init_storage()
104+
super()._init_storage(**kwargs)
104105

105106
if _docs is None:
106107
return

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -180,3 +180,19 @@ def _get_dist(da: 'DocumentArray'):
180180
idx = np.take_along_axis(top_inds, permutation, axis=1)
181181

182182
return dist, idx
183+
184+
def _get_root_docs(self, docs: 'DocumentArray') -> 'DocumentArray':
185+
"""Get the root documents of the current DocumentArray.
186+
187+
:return: a `DocumentArray` containing the root documents.
188+
"""
189+
from docarray import DocumentArray
190+
191+
root_da_flat = self[...]
192+
da = DocumentArray()
193+
for doc in docs:
194+
result = doc
195+
while getattr(result, 'parent_id', None):
196+
result = root_da_flat[result.parent_id]
197+
da.append(result)
198+
return da

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,7 @@ class MilvusConfig:
9393
batch_size: int = -1
9494
columns: Optional[Union[List[Tuple[str, str]], Dict[str, str]]] = None
9595
list_like: bool = True
96+
root_id: bool = True
9697

9798

9899
class BackendMixin(BaseBackendMixin):
@@ -134,7 +135,7 @@ def _init_storage(
134135
self._collection = self._create_or_reuse_collection()
135136
self._offset2id_collection = self._create_or_reuse_offset2id_collection()
136137
self._build_index()
137-
super()._init_storage()
138+
super()._init_storage(**kwargs)
138139

139140
# To align with Sqlite behavior; if `docs` is not `None` and table name
140141
# is provided, :class:`DocumentArraySqlite` will clear the existing

0 commit comments

Comments
 (0)