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

Commit 77ce8b6

Browse files
winstonwwwinstonww
committed
feat: add setitem and delitem initial commit
1 parent 1fff548 commit 77ce8b6

2 files changed

Lines changed: 305 additions & 171 deletions

File tree

‎docarray/array/weaviate.py‎

Lines changed: 101 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import itertools
22
import uuid
3+
from collections.abc import Iterable
34
from typing import (
45
Optional,
56
TYPE_CHECKING,
@@ -17,6 +18,7 @@
1718

1819
import numpy as np
1920
import weaviate
21+
import scipy.sparse
2022

2123
from .mixins import AllMixins
2224
from .. import Document, DocumentArray
@@ -33,6 +35,15 @@
3335
)
3436

3537

38+
def wmap(doc_id: str):
39+
# TODO: although rare, it's possible for key collision to occur with SHA-1
40+
return str(uuid.uuid5(uuid.NAMESPACE_URL, doc_id))
41+
42+
43+
def find(l, item):
44+
return l.index(item) if item in l else -1
45+
46+
3647
class DocumentArrayWeaviate(AllMixins, MutableSequence[Document]):
3748
def __init__(
3849
self,
@@ -71,6 +82,7 @@ def _upload_weaviate_schema(self):
7182
}
7283
]
7384
}
85+
self._client.schema.delete_all()
7486
self._client.schema.create(doc_schema)
7587

7688
def insert(self, index: int, value: 'Document'):
@@ -79,15 +91,22 @@ def insert(self, index: int, value: 'Document'):
7991
:param index: Position of the insertion.
8092
:param value: The doc needs to be inserted.
8193
"""
82-
self._offset2ids.insert(index, value.id)
94+
self._offset2ids.insert(index, wmap(value.id))
8395
self._client.data_object.create(**self._doc2weaviate_create_payload(value))
8496

8597
def _doc2weaviate_create_payload(self, value: 'Document'):
98+
if value.embedding is None:
99+
embedding = [0]
100+
elif isinstance(value.embedding, scipy.sparse.spmatrix):
101+
embedding = value.embedding.toarray()
102+
else:
103+
embedding = value.embedding
104+
86105
return dict(
87106
data_object={'_serialized': value.to_base64()},
88107
class_name=self._class_name,
89-
uuid=str(uuid.UUID(value.id)),
90-
vector=value.embedding or [0],
108+
uuid=wmap(value.id),
109+
vector=embedding,
91110
)
92111

93112
def __eq__(self, other):
@@ -105,14 +124,14 @@ def __len__(self):
105124
)
106125

107126
def __iter__(self) -> Iterator['Document']:
108-
for _id in self._offset2ids:
109-
yield self[_id]
127+
for int_id in range(len(self._offset2ids)):
128+
yield self[int_id]
110129

111130
def __contains__(self, x: Union[str, 'Document']):
112131
if isinstance(x, str):
113-
return self._client.data_object.exists(x)
132+
return self._client.data_object.exists(wmap(x))
114133
elif isinstance(x, Document):
115-
return self._client.data_object.exists(x.id)
134+
return self._client.data_object.exists(wmap(x.id))
116135
else:
117136
return False
118137

@@ -134,25 +153,27 @@ def __getitem__(
134153
) -> List[List[Any]]:
135154
...
136155

156+
def _getitem(self, wid: str):
157+
resp = self._client.data_object.get_by_id(wid, with_vector=True)
158+
if not resp:
159+
raise KeyError(wid)
160+
return Document.from_base64(resp['properties']['_serialized'])
161+
137162
def __getitem__(
138163
self, index: 'DocumentArrayIndexType'
139164
) -> Union['Document', 'DocumentArray']:
140165
if isinstance(index, (int, np.generic)) and not isinstance(index, bool):
141166
# convert into integer, map to string key, delegate the string __getitem__
142-
return self[self._offset2ids[int(index)]]
167+
return self._getitem(self._offset2ids[index])
143168
elif isinstance(index, str):
144169
if index.startswith('@'):
145170
return self.traverse_flat(index[1:])
146171
else:
147-
return Document.from_base64(
148-
self._client.data_object.get_by_id(index, with_vector=True)[
149-
'properties'
150-
]['_serialized']
151-
)
172+
return self._getitem(wmap(index))
152173
elif isinstance(index, slice):
153174
# convert it to a sequence of strings.
154175
_ids = self._offset2ids[index]
155-
return DocumentArray(self[_id] for _id in _ids)
176+
return DocumentArray(self._getitem(_id) for _id in _ids)
156177
elif index is Ellipsis:
157178
return self.flatten()
158179
elif isinstance(index, Sequence):
@@ -163,7 +184,7 @@ def __getitem__(
163184
):
164185
if isinstance(index[0], str) and isinstance(index[1], str):
165186
# ambiguity only comes from the second string
166-
if index[1] in self._offset2ids:
187+
if wmap(index[1]) in self._offset2ids:
167188
return DocumentArray([self[index[0]], self[index[1]]])
168189
else:
169190
return getattr(self[index[0]], index[1])
@@ -174,8 +195,8 @@ def __getitem__(
174195
_attrs = (index[1],)
175196
return _docs._get_attributes(*_attrs)
176197
elif isinstance(index[0], bool):
177-
_ids = itertools.compress(self._offset2ids, index)
178-
return self[_ids]
198+
_ids = list(itertools.compress(self._offset2ids, index))
199+
return DocumentArray(self._getitem(_id) for _id in _ids)
179200
elif isinstance(index[0], int):
180201
return DocumentArray(self[t] for t in index)
181202
elif isinstance(index[0], str):
@@ -222,31 +243,39 @@ def __setitem__(
222243
):
223244
...
224245

246+
def _setitem(self, wid: str, value: Document):
247+
exist = self._client.data_object.exists(wid)
248+
self._client.data_object.delete(wid)
249+
self._client.data_object.create(**self._doc2weaviate_create_payload(value))
250+
self._offset2ids[find(self._offset2ids, wid)] = wmap(value.id)
251+
225252
def __setitem__(
226253
self,
227254
index: 'DocumentArrayIndexType',
228255
value: Union['Document', Sequence['Document']],
229256
):
230257
if isinstance(index, (int, np.generic)) and not isinstance(index, bool):
231-
index = int(index)
232-
self._data[index] = value
233-
self._id2offset[value.id] = index
258+
wid = self._offset2ids[int(index)]
259+
self._setitem(wid, value)
260+
# update the weaviatei
261+
self._offset2ids[index] = wmap(value.id)
262+
234263
elif isinstance(index, str):
235264
if index.startswith('@'):
236-
for _d, _v in zip(self.traverse_flat(index[1:]), value):
237-
_d._data = _v._data
238-
self._rebuild_id2offset()
265+
raise NotImplementedError(
266+
'set along traversal paths is not implemented'
267+
)
239268
else:
240-
old_idx = self._id2offset.pop(index)
241-
self._data[old_idx] = value
242-
self._id2offset[value.id] = old_idx
269+
self._setitem(wmap(index), value)
243270
elif isinstance(index, slice):
244-
self._data[index] = value
245-
self._rebuild_id2offset()
271+
if not isinstance(value, Iterable):
272+
raise TypeError('can only assign an iterable')
273+
index = self._offset2ids[index]
274+
for _i, _v in zip(index, value):
275+
self._setitem(_i, _v)
246276
elif index is Ellipsis:
247277
for _d, _v in zip(self.flatten(), value):
248-
_d._data = _v._data
249-
self._rebuild_id2offset()
278+
self._setitem(wmap(d.id), _v)
250279
elif isinstance(index, Sequence):
251280
if (
252281
isinstance(index, tuple)
@@ -255,12 +284,14 @@ def __setitem__(
255284
):
256285
if isinstance(index[0], str) and isinstance(index[1], str):
257286
# ambiguity only comes from the second string
258-
if index[1] in self._id2offset:
287+
i = find(self._offset2ids, wmap(index[1]))
288+
if i >= 0:
259289
for _d, _v in zip((self[index[0]], self[index[1]]), value):
260-
_d._data = _v._data
261-
self._rebuild_id2offset()
290+
self._setitem(wmap(_d.id), _v)
262291
elif hasattr(self[index[0]], index[1]):
263-
setattr(self[index[0]], index[1], value)
292+
doc = self[index[0]]
293+
setattr(doc, index[1], value)
294+
self._setitem(wmap(doc.id), doc)
264295
else:
265296
# to avoid accidentally add new unsupport attribute
266297
raise ValueError(
@@ -295,15 +326,19 @@ def __setitem__(
295326
else:
296327
for _d, _vv in zip(_docs, _v):
297328
setattr(_d, _a, _vv)
329+
for d in _docs:
330+
# TODO: refactor this to batch update
331+
self._setitem(wmap(d.id), d)
332+
298333
elif isinstance(index[0], bool):
299-
if len(index) != len(self._data):
334+
if len(index) != len(self):
300335
raise IndexError(
301-
f'Boolean mask index is required to have the same length as {len(self._data)}, '
336+
f'Boolean mask index is required to have the same length as {len(self)}, '
302337
f'but receiving {len(index)}'
303338
)
304-
_selected = itertools.compress(self._data, index)
339+
_selected = itertools.compress(self._offset2ids, index)
305340
for _idx, _val in zip(_selected, value):
306-
self[_idx.id] = _val
341+
self._setitem(_idx, _val)
307342
elif isinstance(index[0], (int, str)):
308343
if not isinstance(value, Sequence) or len(index) != len(value):
309344
raise ValueError(
@@ -327,25 +362,29 @@ def __setitem__(
327362
else:
328363
raise IndexError(f'Unsupported index type {typename(index)}: {index}')
329364

365+
def _delitem(self, key):
366+
self._offset2ids.pop(find(self._offset2ids, key))
367+
self._client.data_object.delete(key)
368+
330369
def __delitem__(self, index: 'DocumentArrayIndexType'):
331370
if isinstance(index, (int, np.generic)) and not isinstance(index, bool):
332371
index = int(index)
333-
self._id2offset.pop(self._data[index].id)
334-
del self._data[index]
372+
self._delitem(self._offset2ids[index])
335373
elif isinstance(index, str):
336374
if index.startswith('@'):
337375
raise NotImplementedError(
338376
'Delete elements along traversal paths is not implemented'
339377
)
340378
else:
341-
del self._data[self._id2offset[index]]
342-
self._id2offset.pop(index)
379+
self._delitem(wmap(index))
343380
elif isinstance(index, slice):
344-
del self._data[index]
345-
self._rebuild_id2offset()
381+
start = index.start or 0
382+
stop = index.stop or len(self)
383+
step = index.step or 1
384+
del self[list(range(start, stop, step))]
346385
elif index is Ellipsis:
347-
self._data.clear()
348-
self._id2offset.clear()
386+
self._client.schema.delete_all()
387+
self._offset2ids.clear()
349388
elif isinstance(index, Sequence):
350389
if (
351390
isinstance(index, tuple)
@@ -354,23 +393,32 @@ def __delitem__(self, index: 'DocumentArrayIndexType'):
354393
):
355394
if isinstance(index[0], str) and isinstance(index[1], str):
356395
# ambiguity only comes from the second string
357-
if index[1] in self._id2offset:
396+
i = find(self._offset2ids, wmap(index[1]))
397+
if i >= 0:
358398
del self[index[0]]
359399
del self[index[1]]
400+
elif index[1] != 'id':
401+
doc = self[index[0]]
402+
doc.pop(index[1])
403+
self._setitem(wmap(index[0]), doc)
360404
else:
361-
self[index[0]].pop(index[1])
405+
raise ValueError('cannot pop id from DocumentArrayWeaviate')
362406
elif isinstance(index[0], (slice, Sequence)):
363407
_docs = self[index[0]]
364408
_attrs = index[1]
409+
if 'id' in _attrs:
410+
raise ValueError('cannot pop id from DocumentArrayWeaviate')
365411
if isinstance(_attrs, str):
366412
_attrs = (index[1],)
367413
for _d in _docs:
368414
_d.pop(*_attrs)
415+
self._setitem(wmap(_d.id), _d)
369416
elif isinstance(index[0], bool):
370-
self._data = list(
371-
itertools.compress(self._data, (not _i for _i in index))
417+
idx = list(
418+
itertools.compress(self._offset2ids, (not _i for _i in index))
372419
)
373-
self._rebuild_id2offset()
420+
for _idx in reversed(idx):
421+
self._delitem(_idx)
374422
elif isinstance(index[0], int):
375423
for t in sorted(index, reverse=True):
376424
del self[t]
@@ -390,7 +438,7 @@ def __delitem__(self, index: 'DocumentArrayIndexType'):
390438

391439
def clear(self):
392440
"""Clear the data of :class:`DocumentArray`"""
393-
self._data.clear()
441+
self._client.schema.delete_all()
394442
self._offset2ids.clear()
395443

396444
def __bool__(self):
@@ -415,4 +463,4 @@ def extend(self, values: Iterable['Document']) -> None:
415463
with self._client.batch as _b:
416464
for d in values:
417465
_b.add_data_object(**self._doc2weaviate_create_payload(d))
418-
self._offset2ids.append(d.id)
466+
self._offset2ids.append(wmap(d.id))

0 commit comments

Comments
 (0)