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

Commit 706e884

Browse files
dongxiang123dongxiang123
authored
feat: add flag to disable offsettoid in Offsettoid class (#730)
1 parent a2ce9bb commit 706e884

18 files changed

Lines changed: 216 additions & 56 deletions

File tree

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
class AnnliteConfig:
2424
n_dim: int
2525
metric: str = 'cosine'
26+
list_like: bool = True
2627
serialize_config: Dict = field(default_factory=dict)
2728
data_path: Optional[str] = None
2829
ef_construction: Optional[int] = None
@@ -98,7 +99,7 @@ def _init_storage(
9899
self._config.columns = self._normalize_columns(self._config.columns)
99100
config = asdict(config)
100101
self.n_dim = config.pop('n_dim')
101-
102+
self._list_like = config.pop("list_like")
102103
from annlite import AnnLite
103104

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

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

Lines changed: 20 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -46,15 +46,25 @@ def __del__(self) -> None:
4646
self._annlite.close()
4747

4848
def _load_offset2ids(self):
49-
self._offsetmapping = OffsetMapping(
50-
data_path=self._config.data_path, in_memory=False
51-
)
52-
self._offsetmapping.create_table()
53-
self._offset2ids = Offset2ID(self._offsetmapping.get_all_ids())
49+
if self._list_like:
50+
self._offsetmapping = OffsetMapping(
51+
data_path=self._config.data_path, in_memory=False
52+
)
53+
self._offsetmapping.create_table()
54+
self._offset2ids = Offset2ID(
55+
self._offsetmapping.get_all_ids(),
56+
list_like=self._list_like,
57+
)
58+
else:
59+
self._offset2ids = Offset2ID(
60+
[],
61+
list_like=self._list_like,
62+
)
5463

5564
def _save_offset2ids(self):
56-
self._offsetmapping.drop()
57-
self._offsetmapping.create_table()
58-
self._offsetmapping._insert(
59-
[(i, doc_id) for i, doc_id in enumerate(self._offset2ids.ids)]
60-
)
65+
if self._list_like:
66+
self._offsetmapping.drop()
67+
self._offsetmapping.create_table()
68+
self._offsetmapping._insert(
69+
[(i, doc_id) for i, doc_id in enumerate(self._offset2ids.ids)]
70+
)

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -129,7 +129,7 @@ def _del_docs_by_mask(self, mask: Sequence[bool]):
129129
def _del_all_docs(self):
130130
self._clear_subindices()
131131
self._clear_storage()
132-
self._offset2ids = Offset2ID()
132+
self._offset2ids = Offset2ID(list_like=self._list_like)
133133

134134
def _del_docs_by_ids(self, ids):
135135
"""This function is derived from :meth:`_del_doc_by_id`

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

Lines changed: 33 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -2,44 +2,64 @@
22

33

44
class Offset2ID:
5-
def __init__(self, ids=None):
5+
def __init__(self, ids=None, list_like=True):
66
self.ids = ids or []
7+
self._list_like = list_like
78

89
def get_id(self, idx):
10+
if not self._list_like:
11+
raise ValueError(
12+
"The offset2id is not enabled for list-like indexes. To avoid this error, configure the "
13+
"`list_like` to True"
14+
)
915
return self.ids[idx]
1016

1117
def append(self, data):
12-
self.ids.append(data)
18+
if self._list_like:
19+
self.ids.append(data)
1320

1421
def extend(self, data):
15-
self.ids.extend(data)
22+
if self._list_like:
23+
self.ids.extend(data)
1624

1725
def update(self, position, data_id):
18-
self.ids[position] = data_id
26+
if self._list_like:
27+
self.ids[position] = data_id
1928

2029
def delete_by_id(self, _id):
21-
del self.ids[self.ids.index(_id)]
30+
if self._list_like:
31+
del self.ids[self.ids.index(_id)]
2232

2333
def index(self, _id):
34+
if not self._list_like:
35+
raise ValueError(
36+
"The offset2id is not enabled for list-like indexes. To avoid this error, configure the "
37+
"`list_like` to True"
38+
)
2439
return self.ids.index(_id)
2540

2641
def delete_by_offset(self, position):
27-
del self.ids[position]
42+
if self._list_like:
43+
del self.ids[position]
2844

2945
def insert(self, position, data_id):
30-
self.ids.insert(position, data_id)
46+
if self._list_like:
47+
self.ids.insert(position, data_id)
3148

3249
def clear(self):
33-
self.ids.clear()
50+
if self._list_like:
51+
self.ids.clear()
3452

3553
def delete_by_ids(self, ids):
36-
ids = set(ids)
37-
self.ids = list(filter(lambda _id: _id not in ids, self.ids))
54+
if self._list_like:
55+
ids = set(ids)
56+
self.ids = list(filter(lambda _id: _id not in ids, self.ids))
3857

3958
def update_ids(self, _ids_map: Dict[str, str]):
40-
for i in range(len(self.ids)):
41-
if self.ids[i] in _ids_map:
42-
self.ids[i] = _ids_map[self.ids[i]]
59+
if self._list_like:
60+
for i in range(len(self.ids)):
61+
if self.ids[i] in _ids_map:
62+
self.ids[i] = _ids_map[self.ids[i]]
4363

4464
def save(self):
4565
pass

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ class ElasticConfig:
3838
str, List[Union[str, Mapping[str, Union[str, int]]]], None
3939
] = 'http://localhost:9200'
4040
index_name: Optional[str] = None
41+
list_like: bool = True
4142
es_config: Dict[str, Any] = field(default_factory=dict)
4243
index_text: bool = False
4344
tag_indices: List[str] = field(default_factory=list)
@@ -93,6 +94,7 @@ def _init_storage(
9394

9495
self.n_dim = self._config.n_dim
9596
self._client = self._build_client()
97+
self._list_like = self._config.list_like
9698
self._build_offset2id_index()
9799

98100
# Note super()._init_storage() calls _load_offset2ids which calls _get_offset2ids_meta

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

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -123,8 +123,12 @@ def _clear_storage(self):
123123
self._client.indices.delete(index=self._config.index_name)
124124

125125
def _load_offset2ids(self):
126-
ids = self._get_offset2ids_meta()
127-
self._offset2ids = Offset2ID(ids)
126+
if self._list_like:
127+
ids = self._get_offset2ids_meta()
128+
self._offset2ids = Offset2ID(ids, list_like=self._list_like)
129+
else:
130+
self._offset2ids = Offset2ID([], list_like=self._list_like)
128131

129132
def _save_offset2ids(self):
130-
self._update_offset2ids_meta()
133+
if self._list_like:
134+
self._update_offset2ids_meta()

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

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -87,9 +87,14 @@ def _upload_batch(self, docs: Iterable['Document'], **kwargs) -> List[int]:
8787
def _extend(self, docs: Iterable['Document'], **kwargs):
8888
docs = list(docs)
8989
successful_indexed_ids = self._upload_batch(docs, **kwargs)
90-
self._offset2ids.extend(
91-
[_id for _id in successful_indexed_ids if _id not in self._offset2ids.ids]
92-
)
90+
if self._list_like:
91+
self._offset2ids.extend(
92+
[
93+
_id
94+
for _id in successful_indexed_ids
95+
if _id not in self._offset2ids.ids
96+
]
97+
)
9398

9499
if len(successful_indexed_ids) != len(docs):
95100
doc_ids = [doc.id for doc in docs]

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ class QdrantConfig:
3838
n_dim: int
3939
distance: str = 'cosine'
4040
collection_name: Optional[str] = None
41+
list_like: bool = True
4142
host: Optional[str] = field(default="localhost")
4243
port: Optional[int] = field(default=6333)
4344
grpc_port: Optional[int] = field(default=6334)
@@ -116,7 +117,7 @@ def _init_storage(
116117
)
117118

118119
self._config = config
119-
120+
self._list_like = config.list_like
120121
self._config.columns = self._normalize_columns(self._config.columns)
121122

122123
self._config.collection_name = (

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

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -123,11 +123,15 @@ def scan(self) -> Iterator['Document']:
123123
break
124124

125125
def _load_offset2ids(self):
126-
ids = self._get_offset2ids_meta()
127-
self._offset2ids = Offset2ID(ids)
126+
if self._list_like:
127+
ids = self._get_offset2ids_meta()
128+
self._offset2ids = Offset2ID(ids, list_like=self._list_like)
129+
else:
130+
self._offset2ids = Offset2ID([], list_like=self._list_like)
128131

129132
def _save_offset2ids(self):
130-
self._update_offset2ids_meta()
133+
if self._list_like:
134+
self._update_offset2ids_meta()
131135

132136
def _clear_storage(self):
133137
self.client.recreate_collection(

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ class RedisConfig:
2121
host: str = field(default='localhost')
2222
port: int = field(default=6379)
2323
index_name: Optional[str] = None
24+
list_like: bool = True
2425
update_schema: bool = field(default=True)
2526
distance: str = field(default='COSINE')
2627
redis_config: Dict[str, Any] = field(default_factory=dict)
@@ -79,6 +80,7 @@ def _init_storage(
7980
self._offset2id_key = config.index_name + '__offset2id'
8081
self._config = config
8182
self.n_dim = self._config.n_dim
83+
self._list_like = config.list_like
8284
self._doc_prefix = config.index_name + ':'
8385
self._config.columns = self._normalize_columns(self._config.columns)
8486

0 commit comments

Comments
 (0)