docarray · maxwelljin · Jun 14, 2023 · Jun 9, 2023 · Jun 12, 2023 · Jun 12, 2023
diff --git a/docarray/index/abstract.py b/docarray/index/abstract.py
@@ -1166,3 +1166,29 @@ def _get_root_doc_id(self, id: str, root: str, sub: str) -> str:
                 id, fields[0], '__'.join(fields[1:])
             )
             return self._get_root_doc_id(cur_root_id, root, '')
+
+    def __contains__(self, item: BaseDoc) -> bool:
+        """Checks if a given BaseDoc item is contained in the index.
+
+        :param item: the given BaseDoc
+        :return: if the given BaseDoc item is contained in the index
+        """
+        return False  # Will be overridden by backends
+
+    def subindex_contains(self, item: BaseDoc) -> bool:
+        """Checks if a given BaseDoc item is contained in the index or any of its subindices.
+
+        :param item: the given BaseDoc
+        :return: if the given BaseDoc item is contained in the index/subindices
+        """
+        if self.num_docs() == 0:
+            return False
+
+        if safe_issubclass(type(item), BaseDoc):
+            return self.__contains__(item) or any(
+                index.subindex_contains(item) for index in self._subindices.values()
+            )
+        else:
+            raise TypeError(
+                f"item must be an instance of BaseDoc or its subclass, not '{type(item).__name__}'"
+            )
diff --git a/docarray/index/backends/elastic.py b/docarray/index/backends/elastic.py
@@ -30,6 +30,7 @@
 from docarray.typing import AnyTensor
 from docarray.typing.tensor.abstract_tensor import AbstractTensor
 from docarray.typing.tensor.ndarray import NdArray
+from docarray.utils._internal._typing import safe_issubclass
 from docarray.utils._internal.misc import import_library
 from docarray.utils.find import _FindResult, _FindResultBatched
 
@@ -670,6 +671,17 @@ def _format_response(self, response: Any) -> Tuple[List[Dict], List[Any]]:
     def _refresh(self, index_name: str):
         self._client.indices.refresh(index=index_name)
 
+    def __contains__(self, item: BaseDoc) -> bool:
+        if safe_issubclass(type(item), BaseDoc):
+            if len(item.id) == 0:
+                return False
+            ret = self._client_mget([item.id])
+            return ret["docs"][0]["found"]
+        else:
+            raise TypeError(
+                f"item must be an instance of BaseDoc or its subclass, not '{type(item).__name__}'"
+            )
+
     ###############################################
     # API Wrappers                                #
     ###############################################

diff --git a/docarray/index/backends/hnswlib.py b/docarray/index/backends/hnswlib.py
@@ -37,8 +37,8 @@
 from docarray.proto import DocProto
 from docarray.typing.tensor.abstract_tensor import AbstractTensor
 from docarray.typing.tensor.ndarray import NdArray
-from docarray.utils._internal.misc import import_library, is_np_int
 from docarray.utils._internal._typing import safe_issubclass
+from docarray.utils._internal.misc import import_library, is_np_int
 from docarray.utils.find import _FindResult, _FindResultBatched
 
 if TYPE_CHECKING:
@@ -392,6 +392,19 @@ def _get_items(self, doc_ids: Sequence[str], out: bool = True) -> Sequence[TSche
             raise KeyError(f'No document with id {doc_ids} found')
         return out_docs
 
+    def __contains__(self, item: BaseDoc):
+        if safe_issubclass(type(item), BaseDoc):
+            hash_id = self._to_hashed_id(item.id)
+            self._sqlite_cursor.execute(
+                f"SELECT data FROM docs WHERE doc_id = '{hash_id}'"
+            )
+            rows = self._sqlite_cursor.fetchall()
+            return len(rows) > 0
+        else:
+            raise TypeError(
+                f"item must be an instance of BaseDoc or its subclass, not '{type(item).__name__}'"
+            )
+
     def num_docs(self) -> int:
         """
         Get the number of documents.

diff --git a/docarray/index/backends/in_memory.py b/docarray/index/backends/in_memory.py
@@ -433,6 +433,14 @@ def _text_search_batched(
     ) -> _FindResultBatched:
         raise NotImplementedError(f'{type(self)} does not support text search.')
 
+    def __contains__(self, item: BaseDoc):
+        if safe_issubclass(type(item), BaseDoc):
+            return any(doc.id == item.id for doc in self._docs)
+        else:
+            raise TypeError(
+                f"item must be an instance of BaseDoc or its subclass, not '{type(item).__name__}'"
+            )
+
     def persist(self, file: str = 'in_memory_index.bin') -> None:
         """Persist InMemoryExactNNIndex into a binary file."""
         self._docs.save_binary(file=file)

diff --git a/docarray/index/backends/qdrant.py b/docarray/index/backends/qdrant.py
@@ -30,6 +30,7 @@
 )
 from docarray.typing import NdArray
 from docarray.typing.tensor.abstract_tensor import AbstractTensor
+from docarray.utils._internal._typing import safe_issubclass
 from docarray.utils._internal.misc import import_library, torch_imported
 from docarray.utils.find import _FindResult
 
@@ -315,6 +316,22 @@ def num_docs(self) -> int:
         """
         return self._client.count(collection_name=self.collection_name).count
 
+    def __contains__(self, item: BaseDoc) -> bool:
+        if safe_issubclass(type(item), BaseDoc):
+            response, _ = self._client.scroll(
+                collection_name=self.index_name,
+                scroll_filter=rest.Filter(
+                    must=[
+                        rest.HasIdCondition(has_id=[self._to_qdrant_id(item.id)]),
+                    ],
+                ),
+            )
+            return len(response) > 0
+        else:
+            raise TypeError(
+                f"item must be an instance of BaseDoc or its subclass, not '{type(item).__name__}'"
+            )
+
     def _del_items(self, doc_ids: Sequence[str]):
         items = self._get_items(doc_ids)
         if len(items) < len(doc_ids):

diff --git a/docarray/index/backends/weaviate.py b/docarray/index/backends/weaviate.py
@@ -31,6 +31,7 @@
 from docarray.typing import AnyTensor
 from docarray.typing.tensor.abstract_tensor import AbstractTensor
 from docarray.typing.tensor.ndarray import NdArray
+from docarray.utils._internal._typing import safe_issubclass
 from docarray.utils._internal.misc import import_library
 from docarray.utils.find import FindResult, _FindResult
 
@@ -762,6 +763,26 @@ def _filter_by_parent_id(self, id: str) -> Optional[List[str]]:
         ]
         return ids
 
+    def __contains__(self, item: BaseDoc) -> bool:
+        if safe_issubclass(type(item), BaseDoc):
+            result = (
+                self._client.query.get(self.index_name, ['docarrayid'])
+                .with_where(
+                    {
+                        "path": ['docarrayid'],
+                        "operator": "Equal",
+                        "valueString": f'{item.id}',
+                    }
+                )
+                .do()
+            )
+            docs = result["data"]["Get"][self.index_name]
+            return docs is not None and len(docs) > 0
+        else:
+            raise TypeError(
+                f"item must be an instance of BaseDoc or its subclass, not '{type(item).__name__}'"
+            )
+
     class QueryBuilder(BaseDocIndex.QueryBuilder):
         def __init__(self, document_index):
             self._queries = [

diff --git a/docarray/utils/_internal/_typing.py b/docarray/utils/_internal/_typing.py
@@ -3,6 +3,7 @@
 from typing_extensions import get_origin
 from typing_inspect import get_args, is_typevar, is_union_type
 
+from docarray.typing.id import ID
 from docarray.typing.tensor.abstract_tensor import AbstractTensor
 
 
@@ -45,6 +46,6 @@ def safe_issubclass(x: type, a_tuple: type) -> bool:
     :return: A boolean value - 'True' if 'x' is a subclass of 'A_tuple', 'False' otherwise.
              Note that if the origin of 'x' is a list or tuple, the function immediately returns 'False'.
     """
-    if (get_origin(x) in (list, tuple, dict, set)) or is_typevar(x):
+    if (get_origin(x) in (list, tuple, dict, set)) or is_typevar(x) or x == ID:
         return False
     return issubclass(x, a_tuple)
diff --git a/tests/index/elastic/v7/test_find.py b/tests/index/elastic/v7/test_find.py
@@ -323,3 +323,22 @@ class MyDoc(BaseDoc):
 
     docs, _ = index.execute_query(query)
     assert [doc['id'] for doc in docs] == ['7', '6', '5', '4']
+
+
+def test_contain():
+    class SimpleSchema(BaseDoc):
+        tens: NdArray[10]
+
+    index = ElasticV7DocIndex[SimpleSchema]()
+    index_docs = [SimpleDoc(tens=np.zeros(10)) for _ in range(10)]
+
+    assert (index_docs[0] in index) is False
+
+    index.index(index_docs)
+
+    for doc in index_docs:
+        assert (doc in index) is True
+
+    index_docs_new = [SimpleDoc(tens=np.zeros(10)) for _ in range(10)]
+    for doc in index_docs_new:
+        assert (doc in index) is False
diff --git a/tests/index/elastic/v7/test_subindex.py b/tests/index/elastic/v7/test_subindex.py
@@ -170,3 +170,32 @@ def test_subindex_del(index):
     assert index._subindices['docs'].num_docs() == 20
     assert index._subindices['list_docs'].num_docs() == 20
     assert index._subindices['list_docs']._subindices['docs'].num_docs() == 100
+
+
+def test_subindex_contain(index):
+    # Checks for individual simple_docs within list_docs
+    for i in range(4):
+        doc = index[f'{i + 1}']
+        for simple_doc in doc.list_docs:
+            assert index.subindex_contains(simple_doc) is True
+            for nested_doc in simple_doc.docs:
+                assert index.subindex_contains(nested_doc) is True
+
+    invalid_doc = SimpleDoc(
+        id='non_existent',
+        simple_tens=np.zeros(10),
+        simple_text='invalid',
+    )
+    assert index.subindex_contains(invalid_doc) is False
+
+    # Checks for an empty doc
+    empty_doc = SimpleDoc(
+        id='',
+        simple_tens=np.zeros(10),
+        simple_text='',
+    )
+    assert index.subindex_contains(empty_doc) is False
+
+    # Empty index
+    empty_index = ElasticV7DocIndex[MyDoc]()
+    assert (empty_doc in empty_index) is False
diff --git a/tests/index/elastic/v8/test_find.py b/tests/index/elastic/v8/test_find.py
@@ -342,3 +342,22 @@ class MyDoc(BaseDoc):
 
     index = ElasticDocIndex[MyDoc]()
     assert index.index_name == MyDoc.__name__.lower()
+
+
+def test_contain():
+    class SimpleSchema(BaseDoc):
+        tens: NdArray[10]
+
+    index = ElasticDocIndex[SimpleSchema]()
+    index_docs = [SimpleDoc(tens=np.zeros(10)) for _ in range(10)]
+
+    assert (index_docs[0] in index) is False
+
+    index.index(index_docs)
+
+    for doc in index_docs:
+        assert (doc in index) is True
+
+    index_docs_new = [SimpleDoc(tens=np.zeros(10)) for _ in range(10)]
+    for doc in index_docs_new:
+        assert (doc in index) is False
diff --git a/tests/index/elastic/v8/test_subindex.py b/tests/index/elastic/v8/test_subindex.py
@@ -166,9 +166,30 @@ def test_subindex_filter(index):
         assert doc.id.split('-')[-1] == '0'
 
 
-def test_subindex_del(index):
-    del index['0']
-    assert index.num_docs() == 4
-    assert index._subindices['docs'].num_docs() == 20
-    assert index._subindices['list_docs'].num_docs() == 20
-    assert index._subindices['list_docs']._subindices['docs'].num_docs() == 100
+def test_subindex_contain(index):
+    # Checks for individual simple_docs within list_docs
+    for i in range(4):
+        doc = index[f'{i + 1}']
+        for simple_doc in doc.list_docs:
+            assert index.subindex_contains(simple_doc) is True
+            for nested_doc in simple_doc.docs:
+                assert index.subindex_contains(nested_doc) is True
+
+    invalid_doc = SimpleDoc(
+        id='non_existent',
+        simple_tens=np.zeros(10),
+        simple_text='invalid',
+    )
+    assert index.subindex_contains(invalid_doc) is False
+
+    # Checks for an empty doc
+    empty_doc = SimpleDoc(
+        id='',
+        simple_tens=np.zeros(10),
+        simple_text='',
+    )
+    assert index.subindex_contains(empty_doc) is False
+
+    # Empty index
+    empty_index = ElasticDocIndex[MyDoc]()
+    assert (empty_doc in empty_index) is False
diff --git a/tests/index/hnswlib/test_find.py b/tests/index/hnswlib/test_find.py
@@ -329,3 +329,19 @@ class SimpleSchema(BaseDoc):
     docs, scores = index.execute_query(q)
 
     assert len(docs) == expected_docs
+
+
+def test_contain(tmp_path):
+    class SimpleSchema(BaseDoc):
+        tens: NdArray[10] = Field(space="cosine")
+
+    index = HnswDocumentIndex[SimpleSchema](work_dir=str(tmp_path))
+    index_docs = [SimpleDoc(tens=np.zeros(10)) for _ in range(10)]
+    index.index(index_docs)
+
+    for doc in index_docs:
+        assert (doc in index) is True
+
+    index_docs_new = [SimpleDoc(tens=np.zeros(10)) for _ in range(10)]
+    for doc in index_docs_new:
+        assert (doc in index) is False
diff --git a/tests/index/hnswlib/test_subindex.py b/tests/index/hnswlib/test_subindex.py
@@ -154,3 +154,32 @@ def test_subindex_del(index):
     assert index._subindices['docs'].num_docs() == 20
     assert index._subindices['list_docs'].num_docs() == 20
     assert index._subindices['list_docs']._subindices['docs'].num_docs() == 100
+
+
+def test_subindex_contain(index):
+    # Checks for individual simple_docs within list_docs
+    for i in range(4):
+        doc = index[f'{i + 1}']
+        for simple_doc in doc.list_docs:
+            assert index.subindex_contains(simple_doc) is True
+            for nested_doc in simple_doc.docs:
+                assert index.subindex_contains(nested_doc) is True
+
+    invalid_doc = SimpleDoc(
+        id='non_existent',
+        simple_tens=np.zeros(10),
+        simple_text='invalid',
+    )
+    assert index.subindex_contains(invalid_doc) is False
+
+    # Checks for an empty doc
+    empty_doc = SimpleDoc(
+        id='',
+        simple_tens=np.zeros(10),
+        simple_text='',
+    )
+    assert index.subindex_contains(empty_doc) is False
+
+    # Empty index
+    empty_index = HnswDocumentIndex[MyDoc]()
+    assert (empty_doc in empty_index) is False
diff --git a/tests/index/in_memory/test_in_memory.py b/tests/index/in_memory/test_in_memory.py
@@ -341,3 +341,9 @@ class MyDoc(BaseDoc):
     del doc_index['0']
     assert doc_index.num_docs() == 9
     assert doc_index._subindices['docs'].num_docs() == 90
+
+
+def test_document_contain(doc_index):
+    num_docs = doc_index.num_docs()
+    for i in range(num_docs):
+        assert (doc_index._docs[i] in doc_index) is True
diff --git a/tests/index/in_memory/test_subindex.py b/tests/index/in_memory/test_subindex.py
@@ -152,3 +152,32 @@ def test_subindex_del(index):
     assert index._subindices['docs'].num_docs() == 20
     assert index._subindices['list_docs'].num_docs() == 20
     assert index._subindices['list_docs']._subindices['docs'].num_docs() == 100
+
+
+def test_subindex_contain(index):
+    # Checks for individual simple_docs within list_docs
+    for i in range(4):
+        doc = index[f'{i + 1}']
+        for simple_doc in doc.list_docs:
+            assert index.subindex_contains(simple_doc) is True
+            for nested_doc in simple_doc.docs:
+                assert index.subindex_contains(nested_doc) is True
+
+    invalid_doc = SimpleDoc(
+        id='non_existent',
+        simple_tens=np.zeros(10),
+        simple_text='invalid',
+    )
+    assert index.subindex_contains(invalid_doc) is False
+
+    # Checks for an empty doc
+    empty_doc = SimpleDoc(
+        id='',
+        simple_tens=np.zeros(10),
+        simple_text='',
+    )
+    assert index.subindex_contains(empty_doc) is False
+
+    # Empty index
+    empty_index = InMemoryExactNNIndex[MyDoc]()
+    assert (empty_doc in empty_index) is False