Repository: jina-ai/serve
Pull request: refactor: refactor DocumentArray
**Changes proposed**
Have `DocumentArray` be a `view` over a `List or RepeatedContainer of DocumentProro` instead of over an `iterable` of `Document`.

But from the user perspective return `Document` instead of `DocumentProto` when iterating or accessing elements in it.

**Benefits**
We had to revert #2648 because `DocumentArray` and `MatchArray` when instantiated from `query.matches` where behaving very differently because the internal `self._doc_views` can hold multiple types

Problem statement:
`DocumentArray` inconsistency
**Describe the bug**
`DocumentArray` holds a `view` of a list of different `DocumentViews` that can end up living in the same `Array,_docs_view` elements of the native jina types and `Proto` types. Which can lead to problems when applying functions as `sort`, `map`, etc ...

This also creates the problem of potential different behavior between `DocumentArray` and `MatchArray` or `ChunkArray`

**Potential solution**
Let's force the conversion to `List[Document]` when instantiating `DocumentArray`.

Hints from issue discussion:

Commit message: Merge branch 'master' of https://github.com/jina-ai/jina into fix-2670

Patch to review:
diff --git a/jina/types/arrays/document.py b/jina/types/arrays/document.py
index 7aab3915178ee..c1bc7a26f430d 100644
--- a/jina/types/arrays/document.py
+++ b/jina/types/arrays/document.py
@@ -13,11 +13,13 @@
     Optional,
     Generator,
     BinaryIO,
+    TypeVar,
 )
 
 from .traversable import TraversableSequence
 from ...helper import typename, cached_property
-from ...proto.jina_pb2 import DocumentProto, DocumentArrayProto
+from ..document import Document
+from ...proto import jina_pb2
 
 try:
     # when protobuf using Cpp backend
@@ -32,6 +34,13 @@
 
 __all__ = ['DocumentArray', 'DocumentArrayGetAttrMixin']
 
+DocumentArraySourceType = TypeVar(
+    'DocumentArraySourceType',
+    jina_pb2.DocumentArrayProto,
+    List[Document],
+    List[jina_pb2.DocumentProto],
+)
+
 if False:
     from ..document import Document
 
@@ -98,24 +107,48 @@ class DocumentArray(
     It gives an efficient view of a list of Document. One can iterate over it like
     a generator but ALSO modify it, count it, get item, or union two 'DocumentArray's using the '+' and '+=' operators.
 
-    :param doc_views: A list of :class:`Document`
-    :type doc_views: Optional[Union['RepeatedContainer', Iterable['Document']]]
+    It is supposed to act as a view containing a pointer to a `RepeatedContainer` of `DocumentProto` while offering `Document` Jina native types
+    when getting items or iterating over it
+
+    :param docs: the document array to construct from. One can also give `DocumentArrayProto` directly, then depending on the ``copy``,
+                it builds a view or a copy from it. It also can accept a List
+    :param copy: when ``docs`` is given as a :class:`DocumentArrayProto` object, build a
+                view (i.e. weak reference) from it or a deep copy from it.
     """
 
-    def __init__(
-        self,
-        doc_views: Optional[Union['RepeatedContainer', Iterable['Document']]] = None,
-    ):
+    def __init__(self, docs: Optional[DocumentArraySourceType] = None):
         super().__init__()
-        if doc_views is not None:
-            from .memmap import DocumentArrayMemmap
-
-            if isinstance(doc_views, (Generator, DocumentArrayMemmap)):
-                self._doc_views = list(doc_views)
+        self._pb_body = []
+        if docs is not None:
+            if isinstance(docs, jina_pb2.DocumentArrayProto):
+                # This would happen when loading from file or memmap
+                self._pb_body = docs.docs
+            elif isinstance(docs, RepeatedContainer):
+                # This would happen when `doc.matches` or `doc.chunks`
+                self._pb_body = docs
+            elif isinstance(docs, DocumentArray):
+                # This would happen in the client
+                self._pb_body = docs._pb_body
+            elif isinstance(docs, list):
+                # This would happen in the client
+                for doc in docs:
+                    if isinstance(doc, Document):
+                        self._pb_body.append(doc.proto)
+                    elif isinstance(doc, jina_pb2.DocumentProto):
+                        self._pb_body.append(doc)
+                    else:
+                        raise ValueError(f'Unexpected element in an input list')
             else:
-                self._doc_views = doc_views
-        else:
-            self._doc_views = []
+                from .memmap import DocumentArrayMemmap
+
+                if isinstance(docs, (Generator, DocumentArrayMemmap)):
+                    docs = list(docs)
+                    for doc in docs:
+                        self.append(doc)
+                else:
+                    raise ValueError(
+                        f'DocumentArray got an unexpected input {type(docs)}'
+                    )
 
     def insert(self, index: int, doc: 'Document') -> None:
         """
@@ -124,7 +157,7 @@ def insert(self, index: int, doc: 'Document') -> None:
         :param index: Position of the insertion.
         :param doc: The doc needs to be inserted.
         """
-        self._doc_views.insert(index, doc.proto)
+        self._pb_body.insert(index, doc.proto)
 
     def __setitem__(self, key, value: 'Document'):
         if isinstance(key, int):
@@ -136,11 +169,11 @@ def __setitem__(self, key, value: 'Document'):
 
     def __delitem__(self, index: Union[int, str, slice]):
         if isinstance(index, int):
-            del self._doc_views[index]
+            del self._pb_body[index]
         elif isinstance(index, str):
             del self[self._id_to_index[index]]
         elif isinstance(index, slice):
-            del self._doc_views[index]
+            del self._pb_body[index]
         else:
             raise IndexError(
                 f'do not support this index type {typename(index)}: {index}'
@@ -148,21 +181,18 @@ def __delitem__(self, index: Union[int, str, slice]):
 
     def __eq__(self, other):
         return (
-            type(self._doc_views) is type(other._doc_views)
-            and self._doc_views == other._doc_views
+            type(self._pb_body) is type(other._pb_body)
+            and self._pb_body == other._pb_body
         )
 
     def __len__(self):
-        return len(self._doc_views)
+        return len(self._pb_body)
 
     def __iter__(self) -> Iterator['Document']:
         from ..document import Document
 
-        for d in self._doc_views:
-            if not isinstance(d, Document):
-                yield Document(d)
-            else:
-                yield d
+        for d in self._pb_body:
+            yield Document(d)
 
     def __contains__(self, item: str):
         return item in self._id_to_index
@@ -171,15 +201,11 @@ def __getitem__(self, item: Union[int, str, slice]):
         from ..document import Document
 
         if isinstance(item, int):
-            view = self._doc_views[item]
-            if not isinstance(view, Document):
-                return Document(view)
-            else:
-                return view
+            return Document(self._pb_body[item])
         elif isinstance(item, str):
             return self[self._id_to_index[item]]
         elif isinstance(item, slice):
-            return DocumentArray(self._doc_views[item])
+            return DocumentArray(self._pb_body[item])
         else:
             raise IndexError(f'do not support this index type {typename(item)}: {item}')
 
@@ -202,7 +228,7 @@ def append(self, doc: 'Document'):
 
         :param doc: The doc needs to be appended.
         """
-        self._doc_views.append(doc.proto)
+        self._pb_body.append(doc.proto)
 
     def extend(self, iterable: Iterable['Document']) -> None:
         """
@@ -215,37 +241,45 @@ def extend(self, iterable: Iterable['Document']) -> None:
 
     def clear(self):
         """Clear the data of :class:`DocumentArray`"""
-        del self._doc_views[:]
+        while len(self._pb_body) > 0:
+            self._pb_body.pop()
 
     def reverse(self):
         """In-place reverse the sequence."""
-        if isinstance(self._doc_views, RepeatedContainer):
-            size = len(self._doc_views)
-            hi_idx = size - 1
-            for i in range(int(size / 2)):
-                tmp = DocumentProto()
-                tmp.CopyFrom(self._doc_views[hi_idx])
-                self._doc_views[hi_idx].CopyFrom(self._doc_views[i])
-                self._doc_views[i].CopyFrom(tmp)
-                hi_idx -= 1
-        elif isinstance(self._doc_views, list):
-            self._doc_views.reverse()
+        size = len(self._pb_body)
+        hi_idx = size - 1
+        for i in range(int(size / 2)):
+            tmp = jina_pb2.DocumentProto()
+            tmp.CopyFrom(self._pb_body[hi_idx])
+            self._pb_body[hi_idx].CopyFrom(self._pb_body[i])
+            self._pb_body[i].CopyFrom(tmp)
+            hi_idx -= 1
 
     @cached_property
     def _id_to_index(self):
         """Returns a doc_id to index in list
 
         .. # noqa: DAR201"""
-        return {d.id: i for i, d in enumerate(self._doc_views)}
+        return {d.id: i for i, d in enumerate(self._pb_body)}
 
-    def sort(self, *args, **kwargs):
+    def sort(self, key=None, *args, **kwargs):
         """
         Sort the items of the :class:`DocumentArray` in place.
 
+        :param key: key callable to sort based upon
         :param args: variable set of arguments to pass to the sorting underlying function
         :param kwargs: keyword arguments to pass to the sorting underlying function
         """
-        self._doc_views.sort(*args, **kwargs)
+
+        def overriden_key(proto):
+            # Function to override the `proto` and wrap it around a `Document` to enable sorting via `Document-like` interface
+            d = Document(proto)
+            return key(d)
+
+        if key:
+            self._pb_body.sort(key=overriden_key, *args, **kwargs)
+        else:
+            self._pb_body.sort(*args, **kwargs)
 
     def __bool__(self):
         """To simulate ```l = []; if l: ...```
@@ -257,22 +291,19 @@ def __bool__(self):
     def __str__(self):
         from ..document import Document
 
-        if hasattr(self._doc_views, '__len__'):
-            content = f'{self.__class__.__name__} has {len(self._doc_views)} items'
+        content = f'{self.__class__.__name__} has {len(self._pb_body)} items'
 
-            if len(self._doc_views) > 3:
-                content += ' (showing first three)'
-        else:
-            content = 'unknown length array'
+        if len(self._pb_body) > 3:
+            content += ' (showing first three)'
 
         content += ':\n'
-        content += ',\n'.join(str(Document(d)) for d in self._doc_views[:3])
+        content += ',\n'.join(str(Document(d)) for d in self._pb_body[:3])
 
         return content
 
     def __repr__(self):
         content = ' '.join(
-            f'{k}={v}' for k, v in {'length': len(self._doc_views)}.items()
+            f'{k}={v}' for k, v in {'length': len(self._pb_body)}.items()
         )
         content += f' at {id(self)}'
         content = content.strip()
@@ -326,12 +357,9 @@ def save_binary(self, file: Union[str, BinaryIO]) -> None:
             file_ctx = open(file, 'wb')
 
         with file_ctx as fp:
-            dap = DocumentArrayProto()
-            if self._doc_views:
-                if isinstance(self._doc_views[0], DocumentProto):
-                    dap.docs.extend(self._doc_views)
-                else:
-                    dap.docs.extend([d.proto for d in self._doc_views])
+            dap = jina_pb2.DocumentArrayProto()
+            if self._pb_body:
+                dap.docs.extend(self._pb_body)
             fp.write(dap.SerializeToString())
 
     def save_json(self, file: Union[str, TextIO]) -> None:
@@ -387,7 +415,7 @@ def load_binary(cls, file: Union[str, BinaryIO]) -> 'DocumentArray':
         else:
             file_ctx = open(file, 'rb')
 
-        dap = DocumentArrayProto()
+        dap = jina_pb2.DocumentArrayProto()
 
         with file_ctx as fp:
             dap.ParseFromString(fp.read())
diff --git a/jina/types/document/__init__.py b/jina/types/document/__init__.py
index ab24d7346bd67..b0a70ac2f9cd4 100644
--- a/jina/types/document/__init__.py
+++ b/jina/types/document/__init__.py
@@ -25,8 +25,6 @@
 from jina.types.struct import StructView
 
 from .converters import png_to_buffer, to_datauri, guess_mime, to_image_blob
-from ..arrays.chunk import ChunkArray
-from ..arrays.match import MatchArray
 from ..mixin import ProtoTypeMixin
 from ..ndarray.generic import NdArray, BaseSparseNdArray
 from ..score import NamedScore
@@ -42,6 +40,9 @@
 from ...proto import jina_pb2
 
 if False:
+    from ..arrays.chunk import ChunkArray
+    from ..arrays.match import MatchArray
+
     from scipy.sparse import coo_matrix
 
     # fix type-hint complain for sphinx and flake
@@ -655,6 +656,9 @@ def matches(self) -> 'MatchArray':
 
         :return: the array of matches attached to this document
         """
+        # Problem with cyclic dependency
+        from ..arrays.match import MatchArray
+
         return MatchArray(self._pb_body.matches, reference_doc=self)
 
     @matches.setter
@@ -672,6 +676,9 @@ def chunks(self) -> 'ChunkArray':
 
         :return: the array of chunks of this document
         """
+        # Problem with cyclic dependency
+        from ..arrays.chunk import ChunkArray
+
         return ChunkArray(self._pb_body.chunks, reference_doc=self)
 
     @chunks.setter
diff --git a/jina/types/request/mixin.py b/jina/types/request/mixin.py
index 97f189040dc54..b1753c6ecbd86 100644
--- a/jina/types/request/mixin.py
+++ b/jina/types/request/mixin.py
@@ -1,3 +1,5 @@
+from abc import abstractmethod
+
 from ..arrays import DocumentArray
 from ...proto import jina_pb2
 
@@ -5,6 +7,14 @@
 class DocsPropertyMixin:
     """Mixin class of docs property."""
 
+    @abstractmethod
+    def body(self):
+        """Requred body property
+
+        .. # noqa: DAR201
+        """
+        ...
+
     @property
     def docs(self) -> 'DocumentArray':
         """Get the :class: `DocumentArray` with sequence `body.docs` as content.
@@ -17,6 +27,14 @@ def docs(self) -> 'DocumentArray':
 class GroundtruthPropertyMixin:
     """Mixin class of groundtruths property."""
 
+    @abstractmethod
+    def body(self):
+        """Requred body property
+
+        .. # noqa: DAR201
+        """
+        ...
+
     @property
     def groundtruths(self) -> 'DocumentArray':
         """Get the groundtruths in :class: `DocumentArray` type.
@@ -29,6 +47,14 @@ def groundtruths(self) -> 'DocumentArray':
 class CommandMixin:
     """Mixin class of command property."""
 
+    @abstractmethod
+    def proto(self):
+        """Requred proto property
+
+        .. # noqa: DAR201
+        """
+        ...
+
     @property
     def command(self) -> str:
         """Get the command.
diff --git a/tests/unit/types/arrays/test_documentarray.py b/tests/unit/types/arrays/test_documentarray.py
index 70b465b255e13..fa1b08bf61fda 100644
--- a/tests/unit/types/arrays/test_documentarray.py
+++ b/tests/unit/types/arrays/test_documentarray.py
@@ -1,8 +1,9 @@
 import os
 from copy import deepcopy
 
-import numpy as np
 import pytest
+
+import numpy as np
 from scipy.sparse import coo_matrix
 
 from jina import Document, DocumentArray
@@ -249,10 +250,35 @@ def test_documentarray_filter():
         assert d.score.value > 2
 
 
-def test_da_with_graphs():
-    da = DocumentArray([GraphDocument() for _ in range(6)])
-
+def test_da_with_different_inputs():
+    docs = [Document() for _ in range(10)]
+    da = DocumentArray(
+        [docs[i] if (i % 2 == 0) else docs[i].proto for i in range(len(docs))]
+    )
+    assert len(da) == 10
     for d in da:
-        assert isinstance(d, GraphDocument) is True
+        assert isinstance(d, Document)
+
+
+def test_da_sort_by_document_interface_not_in_proto():
+    docs = [Document(embedding=np.array([1] * (10 - i))) for i in range(10)]
+    da = DocumentArray(
+        [docs[i] if (i % 2 == 0) else docs[i].proto for i in range(len(docs))]
+    )
+    assert len(da) == 10
+    assert da[0].embedding.shape == (10,)
+
+    da.sort(key=lambda d: d.embedding.shape[0])
+    assert da[0].embedding.shape == (1,)
+
+
+def test_da_reverse():
+    docs = [Document(embedding=np.array([1] * (10 - i))) for i in range(10)]
+    da = DocumentArray(
+        [docs[i] if (i % 2 == 0) else docs[i].proto for i in range(len(docs))]
+    )
+    assert len(da) == 10
+    assert da[0].embedding.shape == (10,)
 
-    assert isinstance(da[0], GraphDocument) is True
+    da.reverse()
+    assert da[0].embedding.shape == (1,)
diff --git a/tests/unit/types/arrays/test_matcharray.py b/tests/unit/types/arrays/test_matcharray.py
index dd29509c3ad7f..46970249a3b0c 100644
--- a/tests/unit/types/arrays/test_matcharray.py
+++ b/tests/unit/types/arrays/test_matcharray.py
@@ -1,5 +1,7 @@
 import pytest
 
+import numpy as np
+
 from jina.types.arrays.match import MatchArray
 from jina.types.document import Document
 from jina.types.request import Request
@@ -34,7 +36,7 @@ def matches(document_factory):
             document_factory.create(3, 'test 3'),
         ]
     )
-    return req.proto.data.docs
+    return req.docs
 
 
 @pytest.fixture
@@ -61,3 +63,25 @@ def test_mime_type_not_reassigned():
     d.mime_type = 'text/plain'
     r = d.matches.append(m)
     assert r.mime_type == ''
+
+
+def test_matches_sort_by_document_interface_in_proto():
+    docs = [Document(weight=(10 - i)) for i in range(10)]
+    query = Document()
+    query.matches = docs
+    assert len(query.matches) == 10
+    assert query.matches[0].weight == 10
+
+    query.matches.sort(key=lambda m: m.weight)
+    assert query.matches[0].weight == 1
+
+
+def test_matches_sort_by_document_interface_not_in_proto():
+    docs = [Document(embedding=np.array([1] * (10 - i))) for i in range(10)]
+    query = Document()
+    query.matches = docs
+    assert len(query.matches) == 10
+    assert query.matches[0].embedding.shape == (10,)
+
+    query.matches.sort(key=lambda m: m.embedding.shape[0])
+    assert query.matches[0].embedding.shape == (1,)


Code review comments:
diff_hunk: @@ -98,24 +107,48 @@ class DocumentArray(
     It gives an efficient view of a list of Document. One can iterate over it like
     a generator but ALSO modify it, count it, get item, or union two 'DocumentArray's using the '+' and '+=' operators.
 
-    :param doc_views: A list of :class:`Document`
-    :type doc_views: Optional[Union['RepeatedContainer', Iterable['Document']]]
+    It is supposed to act as a view containing a pointer to a `RepeatedContainer` of `DocumentProto` while offering `Document` Jina native types
+    when getting items or iterating over it
+
+    :param docs: the document array to construct from. One can also give `DocumentArrayProto` directly, then depending on the ``copy``,
+                it builds a view or a copy from it. It also can accept a List
+    :param copy: when ``docs`` is given as a :class:`DocumentArrayProto` object, build a
+                view (i.e. weak reference) from it or a deep copy from it.
path: jina/types/arrays/document.py
text: @user1:
```suggestion
```
i dont think this is implemented anywhere

@user1:
if copy feature is not needed anywhere, then i suggest remove this docstring and ready to merge

@author:
oh yes, leftover