-
Notifications
You must be signed in to change notification settings - Fork 243
Expand file tree
/
Copy pathgetsetdel.py
More file actions
341 lines (274 loc) · 12.5 KB
/
Copy pathgetsetdel.py
File metadata and controls
341 lines (274 loc) · 12.5 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
import itertools
from abc import abstractmethod, ABC
from typing import (
Sequence,
Any,
Iterable,
Dict,
)
from docarray.array.storage.base.helper import Offset2ID
from docarray import Document, DocumentArray
def _check_valid_values_nested_set(docs, values):
docs, values = DocumentArray(docs), DocumentArray(values)
if len(docs) != len(values):
raise ValueError(
f'length of docs to set({len(docs)}) does not match '
f'length of values({len(values)})'
)
if docs[:, 'id'] != values[:, 'id']:
raise ValueError(
'Setting Documents by traversal paths with different IDs is not supported'
)
class BaseGetSetDelMixin(ABC):
"""Provide abstract methods and derived methods for ``__getitem__``, ``__setitem__`` and ``__delitem__``
.. note::
The following methods must be implemented:
- :meth:`._get_doc_by_id`
- :meth:`._set_doc_by_id`
- :meth:`._del_doc_by_id`
Keep in mind that these methods above ** must not ** handle offset2id of the DocumentArray.
These methods are actually wrapped by the following methods which handle the offset2id:
- :meth:`._set_doc`
- :meth:`._del_doc`
- :meth:`._del_all_docs`
Therefore, you should make sure to use the wrapper methods in case you expect offset2id to be updated, and use
the inner methods in case you don't want to handle offset2id (for example, if you want to handle it in a
later step)
Other methods implemented a generic-but-slow version that leverage the methods above.
Please override those methods in the subclass whenever a more efficient implementation is available.
Mainly, if the backend storage supports operations in batches, you can implement the following methods:
- :meth:`._get_docs_by_ids`
- :meth:`._set_docs_by_ids`
- :meth:`._del_docs_by_ids`
- :meth:`._clear_storage`
Likewise, the methods above do not handle offset2id. They are wrapped by the following methods that update the
offset2id in a single step:
- :meth:`._set_docs`
- :meth:`._del_docs`
- :meth:`._del_all_docs`
"""
# Getitem APIs
def _get_doc_by_offset(self, offset: int) -> 'Document':
return self._get_doc_by_id(self._offset2ids.get_id(offset))
@abstractmethod
def _get_doc_by_id(self, _id: str) -> 'Document':
...
def _get_docs_by_slice(self, _slice: slice) -> Iterable['Document']:
"""This function is derived from :meth:`_get_doc_by_offset`
Override this function if there is a more efficient logic
:param _slice: the slice used for indexing
:return: an iterable of document
"""
return self._get_docs_by_ids(self._offset2ids.get_id(_slice))
def _get_docs_by_offsets(self, offsets: Sequence[int]) -> Iterable['Document']:
"""This function is derived from :meth:`_get_doc_by_offset`
Override this function if there is a more efficient logic
:param offsets: the offsets used for indexing
:return: an iterable of document
"""
return (self._get_doc_by_offset(o) for o in offsets)
def _get_docs_by_ids(self, ids: Sequence[str]) -> Iterable['Document']:
"""This function is derived from :meth:`_get_doc_by_id`
Override this function if there is a more efficient logic
:param ids: the ids used for indexing
:return: an iterable of document
"""
return (self._get_doc_by_id(_id) for _id in ids)
# Delitem APIs
def _del_doc_by_offset(self, offset: int):
self._del_doc_by_id(self._offset2ids.get_id(offset))
self._offset2ids.delete_by_offset(offset)
def _del_doc(self, _id: str):
self._offset2ids.delete_by_id(_id)
self._del_doc_by_id(_id)
@abstractmethod
def _del_doc_by_id(self, _id: str):
...
def _del_docs_by_slice(self, _slice: slice):
"""This function is derived and may not have the most efficient implementation.
Override this function if there is a more efficient logic
:param _slice: the slice used for indexing
"""
ids = self._offset2ids.get_id(_slice)
self._del_docs(ids)
def _del_docs_by_mask(self, mask: Sequence[bool]):
"""This function is derived and may not have the most efficient implementation.
Override this function if there is a more efficient logic
:param mask: the boolean mask used for indexing
"""
ids = list(itertools.compress(self._offset2ids, (_i for _i in mask)))
self._del_docs(ids)
def _del_all_docs(self):
self._clear_subindices()
self._clear_storage()
self._offset2ids = Offset2ID(list_like=getattr(self, '_list_like', True))
def _del_docs_by_ids(self, ids):
"""This function is derived from :meth:`_del_doc_by_id`
Override this function if there is a more efficient logic
:param ids: the ids used for indexing
"""
for _id in ids:
self._del_doc_by_id(_id)
def _update_subindices_del(self, ids):
if isinstance(ids, str) and ids.startswith('@'):
return # deleting via access path is not supported
if getattr(self, '_subindices', None):
for selector, da in self._subindices.items():
ids_subindex = DocumentArray(self[ids])[selector, 'id']
del da[ids_subindex]
def _del_docs(self, ids):
self._del_docs_by_ids(ids)
self._offset2ids.delete_by_ids(ids)
def _clear_storage(self):
"""This function is derived and may not have the most efficient implementation.
Override this function if there is a more efficient logic.
If you override this method, you should only take care of clearing the storage backend."""
for doc in self:
self._del_doc_by_id(doc.id)
def _clear_subindices(self):
if getattr(self, '_subindices', None):
for selector, da in self._subindices.items():
da._del_all_docs()
# Setitem API
def _set_doc_by_offset(self, offset: int, value: 'Document'):
self._set_doc(self._offset2ids.get_id(offset), value)
def _set_doc(self, _id: str, value: 'Document'):
if _id != value.id:
self._offset2ids.update(self._offset2ids.index(_id), value.id)
self._set_doc_by_id(_id, value)
@abstractmethod
def _set_doc_by_id(self, _id: str, value: 'Document'):
...
def _set_docs_by_ids(self, ids, docs: Iterable['Document'], mismatch_ids: Dict):
"""This function is derived from :meth:`_set_doc_by_id`
Override this function if there is a more efficient logic
:param ids: the ids used for indexing
"""
for _id, doc in zip(ids, docs):
self._set_doc_by_id(_id, doc)
def _update_subindices_set(self, set_index, docs):
subindices = getattr(self, '_subindices', None)
if not subindices:
return
if isinstance(set_index, tuple): # handled later in recursive call
return
if isinstance(set_index, str) and set_index.startswith('@'):
# 'nested' (non root-level) set, update entire subindex directly
_check_valid_values_nested_set(self[set_index], docs)
if set_index in subindices:
subindex_da = subindices[set_index]
subindex_da.clear()
subindex_da.extend(docs)
else: # root level set, update subindices iteratively
for subindex_selector, subindex_da in subindices.items():
old_ids = DocumentArray(self[set_index])[subindex_selector, 'id']
del subindex_da[old_ids]
value = DocumentArray(docs)
if (
getattr(subindex_da, '_config', None) # checks if in-memory da
and subindex_da._config.root_id
):
for v in value:
for doc in DocumentArray(v)[subindex_selector]:
doc.tags['_root_id_'] = v.id
subindex_da.extend(value[subindex_selector])
def _set_docs(self, ids, docs: Iterable['Document']):
docs = list(docs)
mismatch_ids = {_id: doc.id for _id, doc in zip(ids, docs) if _id != doc.id}
self._set_docs_by_ids(ids, docs, mismatch_ids)
self._offset2ids.update_ids(mismatch_ids)
def _set_docs_by_slice(self, _slice: slice, value: Sequence['Document']):
"""This function is derived and may not have the most efficient implementation.
Override this function if there is a more efficient logic
:param _slice: the slice used for indexing
:param value: the value docs will be updated to
:raises TypeError: error raised when right-hand assignment is not an iterable
"""
if not isinstance(value, Iterable):
raise TypeError(
f'You right-hand assignment must be an iterable, receiving {type(value)}'
)
ids = self._offset2ids.get_id(_slice)
self._set_docs(ids, value)
def _set_doc_value_pairs(
self, docs: Iterable['Document'], values: Sequence['Document']
):
docs = list(docs)
if len(docs) != len(values):
raise ValueError(
f'length of docs to set({len(docs)}) does not match '
f'length of values({len(values)})'
)
for _d, _v in zip(docs, values):
self._set_doc(_d.id, _v)
def _set_doc_value_pairs_nested(
self, docs: Iterable['Document'], values: Sequence['Document']
):
"""This function is derived and may not have the most efficient implementation.
Override this function if there is a more efficient logic
:param docs: the docs to update
:param values: the value docs will be updated to
"""
docs = list(docs)
_check_valid_values_nested_set(docs, values)
for _d, _v in zip(docs, values):
_d._data = _v._data
if _d not in self:
root_d = self._find_root_doc_and_modify(_d)
else:
# _d is already on the root-level
root_d = _d
if root_d:
self._set_doc(root_d.id, root_d)
def _set_doc_attr_by_offset(self, offset: int, attr: str, value: Any):
"""This function is derived and may not have the most efficient implementation.
Override this function if there is a more efficient logic
:param offset: the offset used for indexing
:param attr: the attribute of document to update
:param value: the value doc's attr will be updated to
"""
if attr == 'id' and value is None:
raise ValueError(
'setting the ID of a Document stored in a DocumentArray to None is not allowed'
)
_id = self._offset2ids.get_id(offset)
d = self._get_doc_by_id(_id)
if hasattr(d, attr):
setattr(d, attr, value)
self._set_doc(_id, d)
def _set_doc_attr_by_id(self, _id: str, attr: str, value: Any):
"""This function is derived and may not have the most efficient implementation.
Override this function if there is a more efficient logic
:param _id: the id used for indexing
:param attr: the attribute of document to update
:param value: the value doc's attr will be updated to
"""
if attr == 'id' and value is None:
raise ValueError(
'setting the ID of a Document stored in a DocumentArray to None is not allowed'
)
d = self._get_doc_by_id(_id)
if hasattr(d, attr):
setattr(d, attr, value)
self._set_doc(_id, d)
def _find_root_doc_and_modify(self, d: Document) -> 'Document':
"""Find `d`'s root Document in an exhaustive manner
:param: d: the input document
:return: the root of the input document
"""
from docarray import DocumentArray
for _d in self:
da = DocumentArray(_d)[...]
_all_ids = set(da[:, 'id'])
if d.id in _all_ids:
da[d.id].copy_from(d)
return _d
@abstractmethod
def _load_offset2ids(self):
...
@abstractmethod
def _save_offset2ids(self):
...
def sync(self):
if hasattr(self, '_offset2ids'):
self._save_offset2ids()