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

Commit eeaba13

Browse files
AnneYang720AnneYang720
authored
fix: rebuild index when clear storage (#837)
* fix: rebuild index of redis when clear storage Signed-off-by: AnneY <evangeline-lun@foxmail.com> * refactor: change build index funtion name Signed-off-by: AnneY <evangeline-lun@foxmail.com> * fix: rebuild index of elastic when clear storage Signed-off-by: AnneY <evangeline-lun@foxmail.com> * test: add tests for getset subindex in store Signed-off-by: AnneY <evangeline-lun@foxmail.com> * fix: fix elastic __setstate__ client Signed-off-by: AnneY <evangeline-lun@foxmail.com> * refactor: keep _build_client Signed-off-by: AnneY <evangeline-lun@foxmail.com> Signed-off-by: AnneY <evangeline-lun@foxmail.com>
1 parent e6c419c commit eeaba13

5 files changed

Lines changed: 60 additions & 34 deletions

File tree

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

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -93,8 +93,10 @@ def _init_storage(
9393
self._config.columns = self._normalize_columns(self._config.columns)
9494

9595
self.n_dim = self._config.n_dim
96-
self._client = self._build_client()
9796
self._list_like = self._config.list_like
97+
98+
self._client = self._build_client()
99+
self._build_index()
98100
self._build_offset2id_index()
99101

100102
# Note super()._init_storage() calls _load_offset2ids which calls _get_offset2ids_meta
@@ -167,21 +169,22 @@ def _build_schema_from_elastic_config(self, elastic_config):
167169
return da_schema
168170

169171
def _build_client(self):
170-
171172
client = Elasticsearch(
172173
hosts=self._config.hosts,
173174
**self._config.es_config,
174175
)
175176

177+
return client
178+
179+
def _build_index(self):
176180
schema = self._build_schema_from_elastic_config(self._config)
177181

178-
if not client.indices.exists(index=self._config.index_name):
179-
client.indices.create(
182+
if not self._client.indices.exists(index=self._config.index_name):
183+
self._client.indices.create(
180184
index=self._config.index_name, mappings=schema['mappings']
181185
)
182186

183-
client.indices.refresh(index=self._config.index_name)
184-
return client
187+
self._client.indices.refresh(index=self._config.index_name)
185188

186189
def _send_requests(self, request, **kwargs) -> List[Dict]:
187190
"""Send bulk request to Elastic and gather the successful info"""

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,7 @@ def _del_doc_by_id(self, _id: str):
121121
def _clear_storage(self):
122122
"""Concrete implementation of base class' ``_clear_storage``"""
123123
self._client.indices.delete(index=self._config.index_name)
124+
self._build_index()
124125

125126
def _load_offset2ids(self):
126127
if self._list_like:

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

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,8 @@ def _init_storage(
8585
self._config.columns = self._normalize_columns(self._config.columns)
8686

8787
self._client = self._build_client()
88+
self._build_index()
89+
8890
super()._init_storage()
8991

9092
if _docs is None:
@@ -100,19 +102,21 @@ def _build_client(self):
100102
port=self._config.port,
101103
**self._config.redis_config,
102104
)
105+
return client
103106

104-
if self._config.update_schema:
105-
if self._config.index_name.encode() in client.execute_command('FT._LIST'):
106-
client.ft(index_name=self._config.index_name).dropindex()
107+
def _build_index(self, rebuild: bool = False):
108+
if self._config.update_schema or rebuild:
109+
if self._config.index_name.encode() in self._client.execute_command(
110+
'FT._LIST'
111+
):
112+
self._client.ft(index_name=self._config.index_name).dropindex()
107113

108114
schema = self._build_schema_from_redis_config()
109115
idef = IndexDefinition(prefix=[self._doc_prefix])
110-
client.ft(index_name=self._config.index_name).create_index(
116+
self._client.ft(index_name=self._config.index_name).create_index(
111117
schema, definition=idef
112118
)
113119

114-
return client
115-
116120
def _ensure_unique_config(
117121
self,
118122
config_root: dict,
@@ -195,8 +199,4 @@ def __getstate__(self):
195199

196200
def __setstate__(self, state):
197201
self.__dict__ = state
198-
self._client = Redis(
199-
host=self._config.host,
200-
port=self._config.port,
201-
**self._config.redis_config,
202-
)
202+
self._client = self._build_client()

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,4 +125,5 @@ def _clear_storage(self):
125125
self._client.ft(index_name=self._config.index_name).dropindex(
126126
delete_documents=True
127127
)
128+
self._build_index(rebuild=True)
128129
self._client.delete(self._offset2id_key)

‎tests/unit/array/test_advance_indexing.py‎

Lines changed: 38 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -459,25 +459,10 @@ def test_path_syntax_indexing_set(storage, config, use_subindex, start_storage):
459459
assert da[2].id == 'new_id'
460460

461461

462-
@pytest.mark.parametrize(
463-
'storage,config',
464-
[
465-
('memory', None),
466-
('sqlite', None),
467-
('weaviate', WeaviateConfig(n_dim=123)),
468-
('annlite', AnnliteConfig(n_dim=123)),
469-
('qdrant', QdrantConfig(n_dim=123)),
470-
('qdrant', QdrantConfig(n_dim=123, prefer_grpc=True)),
471-
('elasticsearch', ElasticConfig(n_dim=123)),
472-
('redis', RedisConfig(n_dim=123)),
473-
('milvus', MilvusConfig(n_dim=123)),
474-
],
475-
)
476-
def test_getset_subindex(storage, config, start_storage):
462+
def test_getset_subindex():
477463
da = DocumentArray(
478464
[Document(chunks=[Document() for _ in range(5)]) for _ in range(3)],
479-
config=config,
480-
subindex_configs={'@c': {'n_dim': 123}} if config else {'@c': None},
465+
subindex_configs={'@c': None},
481466
)
482467
with da:
483468
assert len(da['@c']) == 15
@@ -509,6 +494,42 @@ def test_getset_subindex(storage, config, start_storage):
509494
assert collected_chunks == new_chunks
510495

511496

497+
@pytest.mark.parametrize(
498+
'storage,config,subindex_config',
499+
[
500+
('memory', None, None),
501+
('sqlite', None, None),
502+
('weaviate', WeaviateConfig(n_dim=123), {'n_dim': 123}),
503+
('annlite', AnnliteConfig(n_dim=123), {'n_dim': 123}),
504+
('qdrant', QdrantConfig(n_dim=123), {'n_dim': 123}),
505+
('qdrant', QdrantConfig(n_dim=123, prefer_grpc=True), {'n_dim': 123}),
506+
('elasticsearch', ElasticConfig(n_dim=123), {'n_dim': 123}),
507+
('redis', RedisConfig(n_dim=123), {'n_dim': 123}),
508+
('milvus', MilvusConfig(n_dim=123), {'n_dim': 123}),
509+
],
510+
)
511+
def test_getset_subindex_in_store(storage, config, subindex_config, start_storage):
512+
da = DocumentArray(
513+
[Document(chunks=[Document() for _ in range(5)]) for _ in range(3)],
514+
storage=storage,
515+
config=config,
516+
subindex_configs={'@c': subindex_config},
517+
)
518+
with da:
519+
assert len(da['@c']) == 15
520+
assert len(da._subindices['@c']) == 15
521+
522+
chunks_ids = [c.id for c in da['@c']]
523+
new_chunks = [
524+
Document(id=cid, embedding=np.ones(123) * i)
525+
for i, cid in enumerate(chunks_ids)
526+
]
527+
da['@c'] = new_chunks
528+
529+
res = da.find(np.random.random(123), on='@c')
530+
assert len(res) > 0
531+
532+
512533
@pytest.mark.parametrize('size', [1, 5])
513534
@pytest.mark.parametrize(
514535
'storage,config_gen',

0 commit comments

Comments
 (0)