forked from docarray/docarray
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathseqlike.py
More file actions
78 lines (65 loc) · 2.65 KB
/
seqlike.py
File metadata and controls
78 lines (65 loc) · 2.65 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
from typing import Union, Optional, Iterable
from docarray.array.storage.base.seqlike import BaseSequenceLikeMixin
from docarray import Document
class SequenceLikeMixin(BaseSequenceLikeMixin):
"""Implement sequence-like methods"""
def _insert_doc_at_idx(self, doc, idx: Optional[int] = None):
if idx is None:
idx = len(self)
self._sql(
f'INSERT INTO {self._table_name} (doc_id, serialized_value, item_order) VALUES (?, ?, ?)',
(doc.id, doc, idx),
)
self._offset2ids.insert(idx, doc.id)
def _shift_index_right_backward(self, start: int):
idx = len(self) - 1
while idx >= start:
self._sql(
f'UPDATE {self._table_name} SET item_order = ? WHERE item_order = ?',
(idx + 1, idx),
)
idx -= 1
def insert(self, index: int, value: 'Document'):
"""Insert `doc` at `index`.
:param index: Position of the insertion.
:param value: The doc needs to be inserted.
"""
length = len(self)
if index < 0:
index = length + index
index = max(0, min(length, index))
self._shift_index_right_backward(index)
self._insert_doc_at_idx(doc=value, idx=index)
self._commit()
def _append(self, doc: 'Document', commit: bool = True, **kwargs) -> None:
self._sql(
f'INSERT INTO {self._table_name} (doc_id, serialized_value, item_order) VALUES (?, ?, ?)',
(doc.id, doc, len(self)),
)
self._offset2ids.append(doc.id)
if commit:
self._commit()
def __contains__(self, item: Union[str, 'Document']):
if isinstance(item, str):
r = self._sql(f'SELECT 1 FROM {self._table_name} WHERE doc_id=?', (item,))
return len(list(r)) > 0
elif isinstance(item, Document):
return item.id in self # fall back to str check
else:
return False
def __len__(self) -> int:
request = self._sql(f'SELECT COUNT(*) FROM {self._table_name}')
return request.fetchone()[0]
def __repr__(self):
return f'<DocumentArray[SQLite] (length={len(self)}) at {id(self)}>'
def __eq__(self, other):
"""In sqlite backend, data are considered as identical if configs point to the same database source"""
return (
type(self) is type(other)
and type(self._config) is type(other._config)
and self._config == other._config
)
def _extend(self, docs: Iterable['Document'], **kwargs) -> None:
for doc in docs:
self._append(doc, commit=False)
self._commit()