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

Commit 0ea6846

Browse files
authored
refactor: contains method in the base class (#1701)
Signed-off-by: jupyterjazz <saba.sturua@jina.ai>
1 parent a3f6998 commit 0ea6846

9 files changed

Lines changed: 65 additions & 87 deletions

File tree

‎docarray/index/abstract.py‎

Lines changed: 25 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -230,6 +230,16 @@ def execute_query(self, query: Any, *args, **kwargs) -> Any:
230230
"""
231231
...
232232

233+
@abstractmethod
234+
def _doc_exists(self, doc_id: str) -> bool:
235+
"""
236+
Checks if a given document exists in the index.
237+
238+
:param doc_id: The id of a document to check.
239+
:return: True if the document exists in the index, False otherwise.
240+
"""
241+
...
242+
233243
@abstractmethod
234244
def _find(
235245
self,
@@ -403,6 +413,21 @@ def __delitem__(self, key: Union[str, Sequence[str]]):
403413
# delete data
404414
self._del_items(key)
405415

416+
def __contains__(self, item: BaseDoc) -> bool:
417+
"""
418+
Checks if a given document exists in the index.
419+
420+
:param item: The document to check.
421+
It must be an instance of BaseDoc or its subclass.
422+
:return: True if the document exists in the index, False otherwise.
423+
"""
424+
if safe_issubclass(type(item), BaseDoc):
425+
return self._doc_exists(str(item.id))
426+
else:
427+
raise TypeError(
428+
f"item must be an instance of BaseDoc or its subclass, not '{type(item).__name__}'"
429+
)
430+
406431
def configure(self, runtime_config=None, **kwargs):
407432
"""
408433
Configure the DocumentIndex.
@@ -1170,14 +1195,6 @@ def _get_root_doc_id(self, id: str, root: str, sub: str) -> str:
11701195
)
11711196
return self._get_root_doc_id(cur_root_id, root, '')
11721197

1173-
def __contains__(self, item: BaseDoc) -> bool:
1174-
"""Checks if a given BaseDoc item is contained in the index.
1175-
1176-
:param item: the given BaseDoc
1177-
:return: if the given BaseDoc item is contained in the index
1178-
"""
1179-
return False # Will be overridden by backends
1180-
11811198
def subindex_contains(self, item: BaseDoc) -> bool:
11821199
"""Checks if a given BaseDoc item is contained in the index or any of its subindices.
11831200

‎docarray/index/backends/elastic.py‎

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -669,16 +669,11 @@ def _format_response(self, response: Any) -> Tuple[List[Dict], List[Any]]:
669669
def _refresh(self, index_name: str):
670670
self._client.indices.refresh(index=index_name)
671671

672-
def __contains__(self, item: BaseDoc) -> bool:
673-
if safe_issubclass(type(item), BaseDoc):
674-
if len(item.id) == 0:
675-
return False
676-
ret = self._client_mget([item.id])
677-
return ret["docs"][0]["found"]
678-
else:
679-
raise TypeError(
680-
f"item must be an instance of BaseDoc or its subclass, not '{type(item).__name__}'"
681-
)
672+
def _doc_exists(self, doc_id: str) -> bool:
673+
if len(doc_id) == 0:
674+
return False
675+
ret = self._client_mget([doc_id])
676+
return ret["docs"][0]["found"]
682677

683678
###############################################
684679
# API Wrappers #

‎docarray/index/backends/hnswlib.py‎

Lines changed: 5 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -393,18 +393,11 @@ def _get_items(self, doc_ids: Sequence[str], out: bool = True) -> Sequence[TSche
393393
raise KeyError(f'No document with id {doc_ids} found')
394394
return out_docs
395395

396-
def __contains__(self, item: BaseDoc):
397-
if safe_issubclass(type(item), BaseDoc):
398-
hash_id = self._to_hashed_id(item.id)
399-
self._sqlite_cursor.execute(
400-
f"SELECT data FROM docs WHERE doc_id = '{hash_id}'"
401-
)
402-
rows = self._sqlite_cursor.fetchall()
403-
return len(rows) > 0
404-
else:
405-
raise TypeError(
406-
f"item must be an instance of BaseDoc or its subclass, not '{type(item).__name__}'"
407-
)
396+
def _doc_exists(self, doc_id: str) -> bool:
397+
hash_id = self._to_hashed_id(doc_id)
398+
self._sqlite_cursor.execute(f"SELECT data FROM docs WHERE doc_id = '{hash_id}'")
399+
rows = self._sqlite_cursor.fetchall()
400+
return len(rows) > 0
408401

409402
def num_docs(self) -> int:
410403
"""

‎docarray/index/backends/in_memory.py‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -431,13 +431,8 @@ def _text_search_batched(
431431
) -> _FindResultBatched:
432432
raise NotImplementedError(f'{type(self)} does not support text search.')
433433

434-
def __contains__(self, item: BaseDoc):
435-
if safe_issubclass(type(item), BaseDoc):
436-
return any(doc.id == item.id for doc in self._docs)
437-
else:
438-
raise TypeError(
439-
f"item must be an instance of BaseDoc or its subclass, not '{type(item).__name__}'"
440-
)
434+
def _doc_exists(self, doc_id: str) -> bool:
435+
return any(doc.id == doc_id for doc in self._docs)
441436

442437
def persist(self, file: Optional[str] = None) -> None:
443438
"""Persist InMemoryExactNNIndex into a binary file."""

‎docarray/index/backends/qdrant.py‎

Lines changed: 10 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -317,21 +317,16 @@ def num_docs(self) -> int:
317317
"""
318318
return self._client.count(collection_name=self.collection_name).count
319319

320-
def __contains__(self, item: BaseDoc) -> bool:
321-
if safe_issubclass(type(item), BaseDoc):
322-
response, _ = self._client.scroll(
323-
collection_name=self.index_name,
324-
scroll_filter=rest.Filter(
325-
must=[
326-
rest.HasIdCondition(has_id=[self._to_qdrant_id(item.id)]),
327-
],
328-
),
329-
)
330-
return len(response) > 0
331-
else:
332-
raise TypeError(
333-
f"item must be an instance of BaseDoc or its subclass, not '{type(item).__name__}'"
334-
)
320+
def _doc_exists(self, doc_id: str) -> bool:
321+
response, _ = self._client.scroll(
322+
collection_name=self.index_name,
323+
scroll_filter=rest.Filter(
324+
must=[
325+
rest.HasIdCondition(has_id=[self._to_qdrant_id(doc_id)]),
326+
],
327+
),
328+
)
329+
return len(response) > 0
335330

336331
def _del_items(self, doc_ids: Sequence[str]):
337332
items = self._get_items(doc_ids)

‎docarray/index/backends/redis.py‎

Lines changed: 1 addition & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -377,7 +377,7 @@ def _del_items(self, doc_ids: Sequence[str]) -> None:
377377
):
378378
self._client.delete(*batch)
379379

380-
def _doc_exists(self, doc_id) -> bool:
380+
def _doc_exists(self, doc_id: str) -> bool:
381381
"""
382382
Checks if a document exists in the index.
383383
@@ -610,18 +610,3 @@ def _text_search_batched(
610610
scores.append(results.scores)
611611

612612
return _FindResultBatched(documents=docs, scores=scores)
613-
614-
def __contains__(self, item: BaseDoc) -> bool:
615-
"""
616-
Checks if a given document exists in the index.
617-
618-
:param item: The document to check.
619-
It must be an instance of BaseDoc or its subclass.
620-
:return: True if the document exists in the index, False otherwise.
621-
"""
622-
if safe_issubclass(type(item), BaseDoc):
623-
return self._doc_exists(item.id)
624-
else:
625-
raise TypeError(
626-
f"item must be an instance of BaseDoc or its subclass, not '{type(item).__name__}'"
627-
)

‎docarray/index/backends/weaviate.py‎

Lines changed: 13 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -760,25 +760,20 @@ def _filter_by_parent_id(self, id: str) -> Optional[List[str]]:
760760
]
761761
return ids
762762

763-
def __contains__(self, item: BaseDoc) -> bool:
764-
if safe_issubclass(type(item), BaseDoc):
765-
result = (
766-
self._client.query.get(self.index_name, ['docarrayid'])
767-
.with_where(
768-
{
769-
"path": ['docarrayid'],
770-
"operator": "Equal",
771-
"valueString": f'{item.id}',
772-
}
773-
)
774-
.do()
775-
)
776-
docs = result["data"]["Get"][self.index_name]
777-
return docs is not None and len(docs) > 0
778-
else:
779-
raise TypeError(
780-
f"item must be an instance of BaseDoc or its subclass, not '{type(item).__name__}'"
763+
def _doc_exists(self, doc_id: str) -> bool:
764+
result = (
765+
self._client.query.get(self.index_name, ['docarrayid'])
766+
.with_where(
767+
{
768+
"path": ['docarrayid'],
769+
"operator": "Equal",
770+
"valueString": f'{doc_id}',
771+
}
781772
)
773+
.do()
774+
)
775+
docs = result["data"]["Get"][self.index_name]
776+
return docs is not None and len(docs) > 0
782777

783778
class QueryBuilder(BaseDocIndex.QueryBuilder):
784779
def __init__(self, document_index):

‎tests/index/base_classes/test_base_doc_store.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,9 @@ def python_type_to_db_type(self, x):
9797
def num_docs(self):
9898
return 3
9999

100+
def _doc_exists(self, doc_id: str) -> bool:
101+
return False
102+
100103
_index = _identity
101104
_del_items = _identity
102105
_get_items = _identity

‎tests/index/base_classes/test_configs.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,6 @@ class DBConfig(BaseDocIndex.DBConfig):
3535

3636
@dataclass
3737
class RuntimeConfig(BaseDocIndex.RuntimeConfig):
38-
3938
default_ef: int = 50
4039

4140

@@ -61,6 +60,7 @@ def python_type_to_db_type(self, x):
6160
_filter_batched = _identity
6261
_text_search = _identity
6362
_text_search_batched = _identity
63+
_doc_exists = _identity
6464

6565

6666
def test_defaults():

0 commit comments

Comments
 (0)