Skip to content

Commit 6b4f6fc

Browse files
JohannesMessnerJohannesMessner
authored
fix: keep offset2id updated when deleting in sqlite (#471)
1 parent 505d776 commit 6b4f6fc

2 files changed

Lines changed: 66 additions & 0 deletions

File tree

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ class GetSetDelMixin(BaseGetSetDelMixin):
1313

1414
def _del_doc_by_id(self, _id: str):
1515
self._sql(f'DELETE FROM {self._table_name} WHERE doc_id=?', (_id,))
16+
self._save_offset2ids()
1617
self._commit()
1718

1819
def _set_doc_by_id(self, _id: str, value: 'Document'):
@@ -47,6 +48,7 @@ def _del_docs_by_ids(self, ids: str) -> Iterable['Document']:
4748
f"DELETE FROM {self._table_name} WHERE doc_id in ({','.join(['?'] * len(ids))})",
4849
ids,
4950
)
51+
self._save_offset2ids()
5052
self._commit()
5153

5254
def _load_offset2ids(self):
@@ -64,3 +66,11 @@ def _save_offset2ids(self):
6466
(offset, doc_id),
6567
)
6668
self._commit()
69+
70+
def _del_docs(self, ids):
71+
super()._del_docs(ids)
72+
self._save_offset2ids()
73+
74+
def _del_doc_by_offset(self, offset: int):
75+
super()._del_doc_by_offset(offset)
76+
self._save_offset2ids()

‎tests/unit/array/test_sequence.py‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,3 +103,59 @@ def test_context_manager_from_disk(storage, config, start_storage, tmpdir, tmpfi
103103

104104
del da
105105
del da2
106+
107+
108+
@pytest.mark.parametrize(
109+
'storage, config',
110+
[
111+
('memory', None),
112+
('weaviate', {'n_dim': 3, 'distance': 'l2-squared'}),
113+
('annlite', {'n_dim': 3, 'metric': 'Euclidean'}),
114+
('qdrant', {'n_dim': 3, 'distance': 'euclidean'}),
115+
('elasticsearch', {'n_dim': 3, 'distance': 'l2_norm'}),
116+
('sqlite', dict()),
117+
],
118+
)
119+
@pytest.mark.parametrize(
120+
'index', [1, '1', slice(1, 2), [1], [False, True, False, False, False]]
121+
)
122+
def test_del_and_append(index, storage, config):
123+
da = DocumentArray(storage=storage, config=config)
124+
125+
with da:
126+
da.extend([Document(id=str(i)) for i in range(5)])
127+
with da:
128+
del da[index]
129+
da.append(Document(id='new'))
130+
131+
assert da[:, 'id'] == ['0', '2', '3', '4', 'new']
132+
133+
134+
@pytest.mark.parametrize(
135+
'index', [1, '1', slice(1, 2), [1], [False, True, False, False, False]]
136+
)
137+
@pytest.mark.parametrize(
138+
'storage, config',
139+
[
140+
('memory', None),
141+
('weaviate', {'n_dim': 3, 'distance': 'l2-squared'}),
142+
('annlite', {'n_dim': 3, 'metric': 'Euclidean'}),
143+
('qdrant', {'n_dim': 3, 'distance': 'euclidean'}),
144+
('elasticsearch', {'n_dim': 3, 'distance': 'l2_norm'}),
145+
('sqlite', dict()),
146+
],
147+
)
148+
def test_set_and_append(index, storage, config):
149+
da = DocumentArray(storage=storage, config=config)
150+
151+
with da:
152+
da.extend([Document(id=str(i)) for i in range(5)])
153+
with da:
154+
da[index] = (
155+
Document(id='new')
156+
if isinstance(index, int) or isinstance(index, str)
157+
else [Document(id='new')]
158+
)
159+
da.append(Document(id='new_new'))
160+
161+
assert da[:, 'id'] == ['0', 'new', '2', '3', '4', 'new_new']

0 commit comments

Comments
 (0)