feat(array): add traverse sugar

hanxiao · hanxiao · commit dfd045a771c2 · 2021-12-21T21:10:59.000+01:00
diff --git a/docarray/array/document.py b/docarray/array/document.py
@@ -98,7 +98,10 @@ def __getitem__(
         if isinstance(index, (int, np.generic)):
             return self._data[int(index)]
         elif isinstance(index, str):
-            return self._data[self._id2offset[index]]
+            if index.startswith('@'):
+                return self.traverse_flat(index[1:])
+            else:
+                return self._data[self._id2offset[index]]
         elif isinstance(index, slice):
             return DocumentArray(self._data[index])
         elif index is Ellipsis:
@@ -129,9 +132,14 @@ def __setitem__(
             self._data[index] = value
             self._id2offset[value.id] = index
         elif isinstance(index, str):
-            old_idx = self._id2offset.pop(index)
-            self._data[old_idx] = value
-            self._id2offset[value.id] = old_idx
+            if index.startswith('@'):
+                for _d, _v in zip(self.traverse_flat(index[1:]), value):
+                    _d._data = _v._data
+                self._rebuild_id2offset()
+            else:
+                old_idx = self._id2offset.pop(index)
+                self._data[old_idx] = value
+                self._id2offset[value.id] = old_idx
         elif isinstance(index, slice):
             self._data[index] = value
             self._rebuild_id2offset()
@@ -177,7 +185,10 @@ def __delitem__(self, index: 'DocumentArrayIndexType'):
             self._id2offset.pop(self._data[index].id)
             del self._data[index]
         elif isinstance(index, str):
-            del self._data[self._id2offset[index]]
+            if index.startswith('@'):
+                raise NotImplementedError('Delete elements along traversal paths is not implemented')
+            else:
+                del self._data[self._id2offset[index]]
             self._id2offset.pop(index)
         elif isinstance(index, slice):
             del self._data[index]
diff --git a/docarray/array/mixins/__init__.py b/docarray/array/mixins/__init__.py
@@ -45,6 +45,6 @@ class AllMixins(
     DataframeIOMixin,
     ABC,
 ):
-    """All plugins that can be used in :class:`DocumentArray` or :class:`DocumentArrayMemmap`. """
+    """All plugins that can be used in :class:`DocumentArray`. """
 
     ...
diff --git a/docarray/array/mixins/content.py b/docarray/array/mixins/content.py
@@ -54,7 +54,7 @@ def blobs(self) -> Optional['ArrayType']:
 
         .. warning:: This operation assumes all blobs have the same shape and dtype.
                  All dtype and shape values are assumed to be equal to the values of the
-                 first element in the DocumentArray / DocumentArrayMemmap
+                 first element in the DocumentArray
 
         :return: a :class:`ArrayType` of blobs
         """
diff --git a/docarray/array/mixins/empty.py b/docarray/array/mixins/empty.py
@@ -11,10 +11,10 @@ class EmptyMixin:
 
     @classmethod
     def empty(cls: Type['T'], size: int = 0) -> 'T':
-        """Create a :class:`DocumentArray` or :class:`DocumentArrayMemmap` object with :attr:`size` empty
+        """Create a :class:`DocumentArray`  object with :attr:`size` empty
         :class:`Document` objects.
 
         :param size: the number of empty Documents in this container
-        :return: a :class:`DocumentArray` or :class:`DocumentArrayMemmap` object
+        :return: a :class:`DocumentArray` object
         """
         return cls(Document() for _ in range(size))
diff --git a/docarray/array/mixins/group.py b/docarray/array/mixins/group.py
@@ -10,7 +10,7 @@
 
 class GroupMixin:
     """These helpers yield groups of :class:`DocumentArray` from
-    a source :class:`DocumentArray` or :class:`DocumentArrayMemmap`."""
+    a source :class:`DocumentArray`."""
 
     def split(self, tag: str) -> Dict[Any, 'DocumentArray']:
         """Split the `DocumentArray` into multiple DocumentArray according to the tag value of each `Document`.
diff --git a/docarray/array/mixins/match.py b/docarray/array/mixins/match.py
@@ -43,7 +43,7 @@ def match(
         - To invert the distance as score and make all values in range [0, 1],
             use ``dA.match(dB, normalization=(1, 0))``. Note, how ``normalization`` differs from the previous.
         - If a custom metric distance is provided. Make sure that it returns scores as distances and not similarity, meaning the smaller the better.
-        :param darray: the other DocumentArray or DocumentArrayMemmap to match against
+        :param darray: the other DocumentArray  to match against
         :param metric: the distance metric
         :param limit: the maximum number of matches, when not given defaults to 20.
         :param normalization: a tuple [a, b] to be used with min-max normalization,
@@ -136,7 +136,7 @@ def match(
     def _match(self, darray, cdist, limit, normalization, metric_name):
         """
         Computes the matches between self and `darray` loading `darray` into main memory.
-        :param darray: the other DocumentArray or DocumentArrayMemmap to match against
+        :param darray: the other DocumentArray or  to match against
         :param cdist: the distance metric
         :param limit: the maximum number of matches, when not given
                       all Documents in `darray` are considered as matches
@@ -173,7 +173,7 @@ def _match_online(
         """
         Computes the matches between self and `darray` loading `darray` into main memory in chunks of size `batch_size`.
 
-        :param darray: the other DocumentArray or DocumentArrayMemmap to match against
+        :param darray: the other DocumentArray or  to match against
         :param cdist: the distance metric
         :param limit: the maximum number of matches, when not given
                       all Documents in `another` are considered as matches
diff --git a/docarray/array/mixins/parallel.py b/docarray/array/mixins/parallel.py
@@ -10,7 +10,7 @@
 
 
 class ParallelMixin:
-    """Helper functions that provide parallel map to :class:`DocumentArray` or :class:`DocumentArrayMemmap`."""
+    """Helper functions that provide parallel map to :class:`DocumentArray`"""
 
     @overload
     def apply(
@@ -42,7 +42,7 @@ def apply(self: 'T', *args, **kwargs) -> 'T':
         # noqa: DAR102
         # noqa: DAR101
         # noqa: DAR201
-        :return: a new :class:`DocumentArray` or :class:`DocumentArrayMemmap`
+        :return: a new :class:`DocumentArray`
         """
         new_da = type(self)()
         new_da.extend(self.map(*args, **kwargs))
@@ -60,7 +60,7 @@ def map(
 
         .. seealso::
             - To process on a batch of elements, please use :meth:`.map_batch`;
-            - To return a :class:`DocumentArray`/:class:`DocumentArrayMemmap`, please use :meth:`.apply`.
+            - To return a :class:`DocumentArray`, please use :meth:`.apply`.
 
         :param func: a function that takes :class:`Document` as input and outputs anything. You can either modify elements
             in-place (only with `thread` backend) or work later on return elements.
@@ -116,7 +116,7 @@ def apply_batch(self: 'T', *args, **kwargs) -> 'T':
         # noqa: DAR102
         # noqa: DAR101
         # noqa: DAR201
-        :return: a new :class:`DocumentArray` or :class:`DocumentArrayMemmap`
+        :return: a new :class:`DocumentArray`
         """
         new_da = type(self)()
         for _b in self.map_batch(*args, **kwargs):
@@ -138,7 +138,7 @@ def map_batch(
 
         .. seealso::
             - To process single element, please use :meth:`.map`;
-            - To return :class:`DocumentArray` or :class:`DocumentArrayMemmap`, please use :meth:`.apply_batch`.
+            - To return :class:`DocumentArray`, please use :meth:`.apply_batch`.
 
         :param batch_size: Size of each generated batch (except the last one, which might be smaller, default: 32)
         :param shuffle: If set, shuffle the Documents before dividing into minibatches.
diff --git a/docarray/array/mixins/reduce.py b/docarray/array/mixins/reduce.py
@@ -23,7 +23,7 @@ def _reduce_doc_props(doc1: 'Document', doc2: 'Document'):
 
 class ReduceMixin:
     """
-    A mixin that provides reducing logic for :class:`DocumentArray` or :class:`DocumentArrayMemmap`
+    A mixin that provides reducing logic for :class:`DocumentArray`
     Reducing 2 or more DocumentArrays consists in merging all Documents into the same DocumentArray.
     If a Document belongs to 2 or more DocumentArrays, it is added once and data attributes are merged with priority to
     the Document belonging to the left-most DocumentArray. Matches and chunks are also reduced in the same way.
diff --git a/docarray/array/mixins/sample.py b/docarray/array/mixins/sample.py
@@ -22,9 +22,8 @@ def sample(self, k: int, seed: Optional[int] = None) -> 'DocumentArray':
             random.seed(seed)
         # NOTE, this could simplified to random.sample(self, k)
         # without getting indices and itemgetter etc.
-        # however it's only work on DocumentArray, not DocumentArrayMemmap.
-        indices = random.sample(range(len(self)), k)
-        sampled = operator.itemgetter(*indices)(self)
+        # however it's only work on DocumentArray.
+        sampled = random.sample(self, k)
 
         from ..document import DocumentArray
 
diff --git a/docarray/array/mixins/traverse.py b/docarray/array/mixins/traverse.py
@@ -34,7 +34,7 @@ def _check_traversal_path_type(tp):
 
 class TraverseMixin:
     """
-    A mixin used for traversing :class:`DocumentArray` or :class:`DocumentArrayMemmap`.
+    A mixin used for traversing :class:`DocumentArray`.
     """
 
     def traverse(
diff --git a/docarray/base.py b/docarray/base.py
@@ -1,6 +1,7 @@
 import copy as cp
 from dataclasses import fields
 from typing import TYPE_CHECKING, Optional, Tuple, Dict
+
 from .helper import typename
 
 if TYPE_CHECKING:
@@ -109,7 +110,3 @@ def __eq__(self, other):
         if type(self) is type(other):
             return self._data == other._data
         return False
-
-    def __iter__(self):
-        for f in self.non_empty_fields:
-            yield f, getattr(self, f)
diff --git a/docarray/document/mixins/sugar.py b/docarray/document/mixins/sugar.py
@@ -1,7 +1,7 @@
 from typing import overload, TYPE_CHECKING, Union, Callable, Optional, Tuple
 
 if TYPE_CHECKING:
-    from ... import DocumentArray, DocumentArrayMemmap
+    from ... import DocumentArray
     from ...types import AnyDNN, T, ArrayType
 
     import numpy as np
@@ -13,7 +13,7 @@ class SingletonSugarMixin:
     @overload
     def match(
         self: 'T',
-        darray: Union['DocumentArray', 'DocumentArrayMemmap'],
+        darray: 'DocumentArray',
         metric: Union[
             str, Callable[['ArrayType', 'ArrayType'], 'np.ndarray']
         ] = 'cosine',
@@ -35,7 +35,7 @@ def match(
             where you want to find for each element in `A` what are its nearest neighbours in `B`.
             Then you need :meth:`DocumentArray.match`
 
-        :param darray: the other DocumentArray or DocumentArrayMemmap to match against
+        :param darray: the other DocumentArray to match against
         :param metric: the distance metric
         :param limit: the maximum number of matches, when not given defaults to 20.
         :param normalization: a tuple [a, b] to be used with min-max normalization,
@@ -94,3 +94,14 @@ def embed(self: 'T', *args, **kwargs) -> 'T':
         _tmp = DocumentArray([self])
         _tmp.embed(*args, **kwargs)
         return self
+
+    def __getitem__(self, item: str):
+        if item.startswith('@'):
+            from ... import DocumentArray
+            return DocumentArray(self).traverse_flat(item[1:])
+
+    def __setitem__(self, key, value):
+        if key.startswith('@'):
+            from ... import DocumentArray
+            for _d, _v in zip(DocumentArray(self).traverse_flat(key[1:]), value):
+                _d._data = _v._data
diff --git a/docarray/proto/io/__init__.py b/docarray/proto/io/__init__.py
@@ -35,8 +35,9 @@ def parse_proto(pb_msg: 'DocumentProto') -> 'Document':
 
 def flush_proto(doc: 'Document') -> 'DocumentProto':
     pb_msg = DocumentProto()
-    for key, value in doc:
+    for key in doc.non_empty_fields:
         try:
+            value = getattr(doc, key)
             if key in ('blob', 'embedding'):
                 flush_ndarray(getattr(pb_msg, key), value)
             elif key in ('chunks', 'matches'):
