docarray · JoanFM · Nov 28, 2022 · Nov 17, 2022 · Nov 17, 2022 · Nov 21, 2022
diff --git a/docarray/array/mixins/find.py b/docarray/array/mixins/find.py
@@ -1,15 +1,13 @@
 import abc
-from typing import overload, Optional, Union, Dict, List, Tuple, Callable, TYPE_CHECKING
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Union, overload
 
 import numpy as np
-
 from docarray.math import ndarray
 from docarray.score import NamedScore
 
 if TYPE_CHECKING:  # pragma: no cover
-    from docarray.typing import T, ArrayType
-
     from docarray import Document, DocumentArray
+    from docarray.typing import ArrayType, T
 
 
 class FindMixin:
@@ -99,6 +97,7 @@ def find(
         filter: Union[Dict, str, None] = None,
         only_id: bool = False,
         index: str = 'text',
+        return_root: Optional[bool] = False,
         on: Optional[str] = None,
         **kwargs,
     ) -> Union['DocumentArray', List['DocumentArray']]:
@@ -126,14 +125,17 @@ def find(
                       parameter is ignored. By default, the Document `text` attribute will be used for search,
                       otherwise the tag field specified by `index` will be used. You can only use this parameter if the
                       storage backend supports searching by text.
+        :param return_root: if set, then the root-level DocumentArray will be returned
         :param on: specifies a subindex to search on. If set, the returned DocumentArray will be retrieved from the given subindex.
         :param kwargs: other kwargs.
 
         :return: a list of DocumentArrays containing the closest Document objects for each of the queries in `query`.
         """
+        from docarray import Document, DocumentArray
+
         index_da = self._get_index(subindex_name=on)
         if index_da is not self:
-            return index_da.find(
+            results = index_da.find(
                 query,
                 metric,
                 limit,
@@ -144,7 +146,15 @@ def find(
                 index,
                 on=None,
             )
-        from docarray import Document, DocumentArray
+
+            if return_root:
+                da = self._get_root_docs(results)
+                for d, s in zip(da, results[:, 'scores']):
+                    d.scores = s
+
+                return da
+
+            return results
 
         if isinstance(query, dict):
             if filter is None:
@@ -301,3 +311,15 @@ def _find_by_text(self, *args, **kwargs):
         raise NotImplementedError(
             f'Search by text is not supported with this backend {self.__class__.__name__}'
         )
+
+    def _get_root_docs(self, docs: 'DocumentArray') -> 'DocumentArray':
+        """Get the root documents of the current DocumentArray.
+
+        :return: a `DocumentArray` containing the root documents.
+        """
+
+        if not all(docs[:, 'tags___root_id_']):
+            raise ValueError(
+                f'Not all Documents in this subindex have the "_root_id_" attribute set in all `tags`.'
+            )
+        return self[docs[:, 'tags___root_id_']]
diff --git a/docarray/array/mixins/setitem.py b/docarray/array/mixins/setitem.py
@@ -63,6 +63,10 @@ def __setitem__(
         index: 'DocumentArrayIndexType',
         value: Union['Document', Sequence['Document']],
     ):
+        from docarray.helper import check_root_id
+
+        if self._is_subindex:
+            check_root_id(self, value)
 
         self._update_subindices_set(index, value)
         # set by offset

diff --git a/docarray/array/storage/annlite/backend.py b/docarray/array/storage/annlite/backend.py
@@ -31,6 +31,7 @@ class AnnliteConfig:
     max_connection: Optional[int] = None
     n_components: Optional[int] = None
     columns: Optional[Union[List[Tuple[str, str]], Dict[str, str]]] = None
+    root_id: bool = True
 
 
 class BackendMixin(BaseBackendMixin):
@@ -104,7 +105,7 @@ def _init_storage(
 
         self._annlite = AnnLite(self.n_dim, lock=False, **filter_dict(config))
 
-        super()._init_storage()
+        super()._init_storage(**kwargs)
 
         if _docs is None:
             return

diff --git a/docarray/array/storage/annlite/seqlike.py b/docarray/array/storage/annlite/seqlike.py
@@ -20,7 +20,7 @@ def _extend(self, values: Iterable['Document']) -> None:
         self._offset2ids.extend([doc.id for doc in docs])
 
     def _append(self, value: 'Document'):
-        self.extend([value])
+        self._extend([value])
 
     def __eq__(self, other):
         """In annlite backend, data are considered as identical if configs point to the same database source"""

diff --git a/docarray/array/storage/base/backend.py b/docarray/array/storage/base/backend.py
@@ -17,9 +17,11 @@ def _init_storage(
         self,
         _docs: Optional['DocumentArraySourceType'] = None,
         copy: bool = False,
+        _is_subindex: bool = False,
         *args,
         **kwargs,
     ):
+        self._is_subindex = _is_subindex
         self._load_offset2ids()
 
     def _init_subindices(
@@ -40,7 +42,9 @@ def _init_subindices(
                 config_joined = self._ensure_unique_config(
                     config, config_subindex, config_joined, name
                 )
-                self._subindices[name] = self.__class__(config=config_joined)
+                self._subindices[name] = self.__class__(
+                    config=config_joined, _is_subindex=True
+                )
                 if _docs:
                     from docarray import DocumentArray
 

diff --git a/docarray/array/storage/base/getsetdel.py b/docarray/array/storage/base/getsetdel.py
@@ -200,13 +200,25 @@ def _update_subindices_set(self, set_index, docs):
             _check_valid_values_nested_set(self[set_index], docs)
             if set_index in subindices:
                 subindex_da = subindices[set_index]
+
                 subindex_da.clear()
                 subindex_da.extend(docs)
         else:  # root level set, update subindices iteratively
             for subindex_selector, subindex_da in subindices.items():
                 old_ids = DocumentArray(self[set_index])[subindex_selector, 'id']
                 del subindex_da[old_ids]
-                subindex_da.extend(DocumentArray(docs)[subindex_selector])
+
+                value = DocumentArray(docs)
+
+                if (
+                    getattr(subindex_da, '_config', None)  # checks if in-memory da
+                    and subindex_da._config.root_id
+                ):
+                    for v in value:
+                        for doc in DocumentArray(v)[subindex_selector]:
+                            doc.tags['_root_id_'] = v.id
+
+                subindex_da.extend(value[subindex_selector])
 
     def _set_docs(self, ids, docs: Iterable['Document']):
         docs = list(docs)

diff --git a/docarray/array/storage/base/seqlike.py b/docarray/array/storage/base/seqlike.py
@@ -1,5 +1,6 @@
+import warnings
 from abc import abstractmethod
-from typing import Iterator, Iterable, MutableSequence
+from typing import Iterable, Iterator, MutableSequence
 
 from docarray import Document, DocumentArray
 
@@ -10,7 +11,15 @@ class BaseSequenceLikeMixin(MutableSequence[Document]):
     def _update_subindices_append_extend(self, value):
         if getattr(self, '_subindices', None):
             for selector, da in self._subindices.items():
-                docs_selector = DocumentArray(value)[selector]
+
+                value = DocumentArray(value)
+
+                if getattr(da, '_config', None) and da._config.root_id:
+                    for v in value:
+                        for doc in DocumentArray(v)[selector]:
+                            doc.tags['_root_id_'] = v.id
+
+                docs_selector = value[selector]
                 if len(docs_selector) > 0:
                     da.extend(docs_selector)
 
@@ -63,6 +72,12 @@ def __bool__(self):
         return len(self) > 0
 
     def extend(self, values: Iterable['Document'], **kwargs) -> None:
+
+        from docarray.helper import check_root_id
+
+        if self._is_subindex:
+            check_root_id(self, values)
+
         self._extend(values, **kwargs)
         self._update_subindices_append_extend(values)
 

diff --git a/docarray/array/storage/elastic/backend.py b/docarray/array/storage/elastic/backend.py
@@ -46,6 +46,7 @@ class ElasticConfig:
     ef_construction: Optional[int] = None
     m: Optional[int] = None
     columns: Optional[Union[List[Tuple[str, str]], Dict[str, str]]] = None
+    root_id: bool = True
 
 
 _banned_indexname_chars = ['[', ' ', '"', '*', '\\', '<', '|', ',', '>', '/', '?', ']']
@@ -100,7 +101,7 @@ def _init_storage(
         self._build_offset2id_index()
 
         # Note super()._init_storage() calls _load_offset2ids which calls _get_offset2ids_meta
-        super()._init_storage()
+        super()._init_storage(**kwargs)
 
         if _docs is None:
             return

diff --git a/docarray/array/storage/memory/find.py b/docarray/array/storage/memory/find.py
@@ -180,3 +180,19 @@ def _get_dist(da: 'DocumentArray'):
         idx = np.take_along_axis(top_inds, permutation, axis=1)
 
         return dist, idx
+
+    def _get_root_docs(self, docs: 'DocumentArray') -> 'DocumentArray':
+        """Get the root documents of the current DocumentArray.
+
+        :return: a `DocumentArray` containing the root documents.
+        """
+        from docarray import DocumentArray
+
+        root_da_flat = self[...]
+        da = DocumentArray()
+        for doc in docs:
+            result = doc
+            while getattr(result, 'parent_id', None):
+                result = root_da_flat[result.parent_id]
+            da.append(result)
+        return da
diff --git a/docarray/array/storage/milvus/backend.py b/docarray/array/storage/milvus/backend.py
@@ -93,6 +93,7 @@ class MilvusConfig:
     batch_size: int = -1
     columns: Optional[Union[List[Tuple[str, str]], Dict[str, str]]] = None
     list_like: bool = True
+    root_id: bool = True
 
 
 class BackendMixin(BaseBackendMixin):
@@ -134,7 +135,7 @@ def _init_storage(
         self._collection = self._create_or_reuse_collection()
         self._offset2id_collection = self._create_or_reuse_offset2id_collection()
         self._build_index()
-        super()._init_storage()
+        super()._init_storage(**kwargs)
 
         # To align with Sqlite behavior; if `docs` is not `None` and table name
         # is provided, :class:`DocumentArraySqlite` will clear the existing

diff --git a/docarray/array/storage/qdrant/backend.py b/docarray/array/storage/qdrant/backend.py
@@ -51,6 +51,7 @@ class QdrantConfig:
     full_scan_threshold: Optional[int] = None
     m: Optional[int] = None
     columns: Optional[Union[List[Tuple[str, str]], Dict[str, str]]] = None
+    root_id: bool = True
 
 
 class BackendMixin(BaseBackendMixin):
@@ -128,7 +129,7 @@ def _init_storage(
 
         self._initialize_qdrant_schema()
 
-        super()._init_storage()
+        super()._init_storage(**kwargs)
 
         if docs is None and config.collection_name:
             return

diff --git a/docarray/array/storage/redis/backend.py b/docarray/array/storage/redis/backend.py
@@ -35,6 +35,7 @@ class RedisConfig:
     block_size: Optional[int] = None
     initial_cap: Optional[int] = None
     columns: Optional[Union[List[Tuple[str, str]], Dict[str, str]]] = None
+    root_id: bool = True
 
 
 class BackendMixin(BaseBackendMixin):
@@ -87,7 +88,7 @@ def _init_storage(
         self._client = self._build_client()
         self._build_index()
 
-        super()._init_storage()
+        super()._init_storage(**kwargs)
 
         if _docs is None:
             return

diff --git a/docarray/array/storage/sqlite/backend.py b/docarray/array/storage/sqlite/backend.py
@@ -29,6 +29,7 @@ class SqliteConfig:
     conn_config: Dict = field(default_factory=dict)
     journal_mode: str = 'WAL'
     synchronous: str = 'OFF'
+    root_id: bool = True
 
 
 class BackendMixin(BaseBackendMixin):
@@ -101,7 +102,7 @@ def _init_storage(
         self._connection.commit()
         self._config = config
         self._list_like = config.list_like
-        super()._init_storage()
+        super()._init_storage(**kwargs)
 
         if _docs is None:
             return

diff --git a/docarray/array/storage/weaviate/backend.py b/docarray/array/storage/weaviate/backend.py
@@ -52,6 +52,7 @@ class WeaviateConfig:
     # weaviate python client parameters
     batch_size: Optional[int] = field(default=50)
     dynamic_batching: Optional[bool] = field(default=False)
+    root_id: bool = True
 
     def __post_init__(self):
         if isinstance(self.timeout_config, list):

diff --git a/docarray/helper.py b/docarray/helper.py
@@ -12,7 +12,7 @@
 import hubble
 
 if TYPE_CHECKING:  # pragma: no cover
-    from docarray import DocumentArray
+    from docarray import Document, DocumentArray
 
 __resources_path__ = os.path.join(
     os.path.dirname(
@@ -491,6 +491,35 @@ def _get_array_info(da: 'DocumentArray'):
     return is_homo, _nested_in, _nested_items, attr_counter, all_attrs_names
 
 
+def check_root_id(da: 'DocumentArray', value: Union['Document', Sequence['Document']]):
+
+    from docarray import Document
+    from docarray.array.memory import DocumentArrayInMemory
+
+    if not (
+        isinstance(value, Document)
+        or (isinstance(value, Sequence) and isinstance(value[0], Document))
+    ):
+        return
+
+    if isinstance(value, Document):
+        value = [value]
+
+    if isinstance(da, DocumentArrayInMemory):
+        if not all([getattr(doc, 'parent_id', None) for doc in value]):
+            warnings.warn(
+                "Not all documents have parent_id set. This may cause unexpected behavior.",
+                UserWarning,
+            )
+    elif da._config.root_id and not all(
+        [doc.tags.get('_root_id_', None) for doc in value]
+    ):
+        warnings.warn(
+            "root_id is enabled but not all documents have _root_id_ set. This may cause unexpected behavior.",
+            UserWarning,
+        )
+
+
 def login(interactive: Optional[bool] = None, force: bool = False, **kwargs):
     """Login to Jina AI Cloud account.
     :param interactive: If set to true, login will support notebook environments, otherwise the enviroment will be inferred.

diff --git a/docs/advanced/document-store/annlite.md b/docs/advanced/document-store/annlite.md
@@ -48,6 +48,7 @@ The following configs can be set:
 | `max_connection`  | The number of bi-directional links created for every new element during construction.                                                  | `None`, defaults to the default value in the AnnLite package* |
 | `n_components`    | The output dimension of PCA model. Should be a positive number and less than `n_dim` if it's not `None`                                | `None`, defaults to the default value in the AnnLite package* |
 | `list_like`       | Controls if ordering of Documents is persisted in the Database. Disabling this breaks list-like features, but can improve performance. | True                                                          |
+| `root_id`         | Boolean flag indicating whether to store `root_id` in the tags of chunk level Documents                                                | True                                                          |
 
 *You can check the default values in [the AnnLite source code](https://github.com/jina-ai/annlite/blob/main/annlite/core/index/hnsw/index.py)
 

diff --git a/docs/advanced/document-store/elasticsearch.md b/docs/advanced/document-store/elasticsearch.md
@@ -404,6 +404,7 @@ The following configs can be set:
 | `tag_indices`     | List of tags to index                                                                                                                  | False                                                  |
 | `batch_size`      | Batch size used to handle storage refreshes/updates                                                                                    | 64                                                     |
 | `list_like`       | Controls if ordering of Documents is persisted in the Database. Disabling this breaks list-like features, but can improve performance. | True                                                   |
+| `root_id`         | Boolean flag indicating whether to store `root_id` in the tags of chunk level Documents                                                | True                                                   |
 
 ```{tip}
 You can read more about HNSW parameters and their default values [here](https://www.elastic.co/guide/en/elasticsearch/reference/current/dense-vector.html#dense-vector-params)

diff --git a/docs/advanced/document-store/milvus.md b/docs/advanced/document-store/milvus.md
@@ -128,10 +128,11 @@ The following configs can be set:
 | `index_params`      | A dictionary of parameters used for index building. The [allowed parameters](https://milvus.io/docs/v2.1.x/index.md) depend on the index type.                                                                                    | {'M': 4, 'efConstruction': 200} (assumes HNSW index) |
 | `collection_config` | Configuration for the Milvus collection. Passed as **kwargs during collection creation (`Collection(...)`).                                                                                                                       | {}                                                   |
 | `serialize_config`  | [Serialization config of each Document](../../../fundamentals/document/serialization.md)                                                                                                                                          | {}                                                   |
- | `consistency_level` | [Consistency level](https://milvus.io/docs/v2.1.x/consistency.md#Consistency-levels) for Milvus database operations. Can be 'Session', 'Strong', 'Bounded' or 'Eventually'.                                                       | 'Session'                                            |
+| `consistency_level` | [Consistency level](https://milvus.io/docs/v2.1.x/consistency.md#Consistency-levels) for Milvus database operations. Can be 'Session', 'Strong', 'Bounded' or 'Eventually'.                                                       | 'Session'                                            |
 | `batch_size`        | Default batch size for CRUD operations.                                                                                                                                                                                           | -1 (no batching)                                     |
 | `columns`           | Additional columns to be stored in the datbase, taken from Document `tags`.                                                                                                                                                       | None                                                 |
 | `list_like`         | Controls if ordering of Documents is persisted in the Database. Disabling this breaks list-like features, but can improve performance.                                                                                            | True                                                 |
+| `root_id`           | Boolean flag indicating whether to store `root_id` in the tags of chunk level Documents                                                                                                                                           | True                                                 |
 
 ## Minimal example
 

diff --git a/docs/advanced/document-store/qdrant.md b/docs/advanced/document-store/qdrant.md
@@ -94,6 +94,8 @@ The following configs can be set:
 | `m`                   | Number of edges per node in the index graph. Larger = more accurate search, more space required                                          | `None`, defaults to the default value in Qdrant* |
 | `columns`             | Other fields to store in Document                                                                                                        | `None`                                           |
 | `list_like`           | Controls if ordering of Documents is persisted in the Database. Disabling this breaks list-like features, but can improve performance.   | True                                             |
+| `root_id`             | Boolean flag indicating whether to store `root_id` in the tags of chunk level Documents                                                  | True                                             |
+
 
 *You can read more about the HNSW parameters and their default values [here](https://qdrant.tech/documentation/indexing/#vector-index)