docarray · JohannesMessner · Feb 7, 2023 · Feb 2, 2023 · Feb 2, 2023 · Feb 2, 2023
diff --git a/docarray/array/array.py b/docarray/array/array.py
@@ -7,16 +7,22 @@
     Dict,
     Iterable,
     List,
+    Optional,
+    Sequence,
     Type,
     TypeVar,
     Union,
+    cast,
+    overload,
 )
 
+import numpy as np
 from typing_inspect import is_union_type
 
 from docarray.array.abstract_array import AnyDocumentArray
 from docarray.base_document import AnyDocument, BaseDocument
 from docarray.typing import NdArray
+from docarray.utils.misc import is_torch_available
 
 if TYPE_CHECKING:
     from pydantic import BaseConfig
@@ -29,6 +35,7 @@
 
 
 T = TypeVar('T', bound='DocumentArray')
+IndexIterType = Union[slice, Iterable[int], Iterable[bool], None]
 
 
 def _delegate_meth_to_data(meth_name: str) -> Callable:
@@ -48,6 +55,17 @@ def _delegate_meth(self, *args, **kwargs):
     return _delegate_meth
 
 
+def _is_np_int(item: Any) -> bool:
+    dtype = getattr(item, 'dtype', None)
+    ndim = getattr(item, 'ndim', None)
+    if dtype is not None and ndim is not None:
+        try:
+            return ndim == 0 and np.issubdtype(dtype, np.integer)
+        except TypeError:
+            return False
+    return False  # this is unreachable, but mypy wants it
+
+
 class DocumentArray(AnyDocumentArray):
     """
      DocumentArray is a container of Documents.
@@ -65,6 +83,7 @@ class DocumentArray(AnyDocumentArray):
     .. code-block:: python
         from docarray import BaseDocument, DocumentArray
         from docarray.typing import NdArray, ImageUrl
+        from typing import Optional
 
 
         class Image(BaseDocument):
@@ -78,33 +97,167 @@ class Image(BaseDocument):
 
 
     If your DocumentArray is homogeneous (i.e. follows the same schema), you can access
-    fields at the DocumentArray level (for example `da.tensor`). You can also set
-    fields, with `da.tensor = np.random.random([10, 100])`
+    fields at the DocumentArray level (for example `da.tensor` or `da.url`).
+    You can also set fields, with `da.tensor = np.random.random([10, 100])`:
+
+    .. code-block:: python
+        print(da.url)
+        # [ImageUrl('http://url.com/foo.png', host_type='domain'), ...]
+        import numpy as np
+
+        da.tensor = np.random.random([10, 100])
+        print(da.tensor)
+        # [NdArray([0.11299577, 0.47206767, 0.481723  , 0.34754724, 0.15016037,
+        #          0.88861321, 0.88317666, 0.93845579, 0.60486676, ... ]), ...]
+
+    You can index into a DocumentArray like a numpy array or torch tensor:
+
+    .. code-block:: python
+        da[0]  # index by position
+        da[0:5:2]  # index by slice
+        da[[0, 2, 3]]  # index by list of indices
+        da[da.tensor > 0.5]  # index by boolean mask
+
+
     """
 
     document_type: Type[BaseDocument] = AnyDocument
     __typed_da__: Dict[Type[BaseDocument], Type] = {}
 
     def __init__(
         self,
-        docs: Iterable[BaseDocument] = list(),
+        docs: Optional[Iterable[BaseDocument]] = None,
         tensor_type: Type['AbstractTensor'] = NdArray,
     ):
-        self._data = [doc_ for doc_ in docs]
+        self._data = list(docs) if docs is not None else []
         self.tensor_type = tensor_type
 
     def __len__(self):
         return len(self._data)
 
+    @overload
+    def __getitem__(self: T, item: int) -> BaseDocument:
+        ...
+
+    @overload
+    def __getitem__(self: T, item: IndexIterType) -> T:
+        ...
+
     def __getitem__(self, item):
+        item = self._normalize_index_item(item)
+
         if type(item) == slice:
             return self.__class__(self._data[item])
-        else:
+
+        if isinstance(item, int):
             return self._data[item]
 
+        if item is None:
+            return self
+
+        # _normalize_index_item() guarantees the line below is correct
+        head = item[0]  # type: ignore
+        if isinstance(head, bool):
+            return self._get_from_mask(item)
+        elif isinstance(head, int):
+            return self._get_from_indices(item)
+        else:
+            raise TypeError(f'Invalid type {type(head)} for indexing')
+
+    def __setitem__(self: T, key: IndexIterType, value: Union[T, BaseDocument]):
+        key_norm = self._normalize_index_item(key)
+
+        if isinstance(key_norm, int):
+            value_int = cast(BaseDocument, value)
+            self._data[key_norm] = value_int
+        elif isinstance(key_norm, slice):
+            value_slice = cast(T, value)
+            self._data[key_norm] = value_slice
+        else:
+            # _normalize_index_item() guarantees the line below is correct
+            head = key_norm[0]  # type: ignore
+            if isinstance(head, bool):
+                key_norm_ = cast(Iterable[bool], key_norm)
+                value_ = cast(Sequence[BaseDocument], value)  # this is no strictly true
+                # set_by_mask requires value_ to have getitem which
+                # _normalize_index_item() ensures
+                return self._set_by_mask(key_norm_, value_)
+            elif isinstance(head, int):
+                key_norm__ = cast(Iterable[int], key_norm)
+                return self._set_by_indices(key_norm__, value)
+            else:
+                raise TypeError(f'Invalid type {type(head)} for indexing')
+
     def __iter__(self):
         return iter(self._data)
 
+    @staticmethod
+    def _normalize_index_item(
+        item: Any,
+    ) -> Union[int, slice, Iterable[int], Iterable[bool], None]:
+        # basic index types
+        if item is None or isinstance(item, (int, slice, tuple, list)):
+            return item
+
+        # numpy index types
+        if _is_np_int(item):
+            return item.item()
+
+        index_has_getitem = hasattr(item, '__getitem__')
+        is_valid_bulk_index = index_has_getitem and isinstance(item, Iterable)
+        if not is_valid_bulk_index:
+            raise ValueError(f'Invalid index type {type(item)}')
+
+        if isinstance(item, np.ndarray) and (
+            item.dtype == np.bool_ or np.issubdtype(item.dtype, np.integer)
+        ):
+            return item.tolist()
+
+        # torch index types
+        torch_available = is_torch_available()
+        if torch_available:
+            import torch
+        else:
+            raise ValueError(f'Invalid index type {type(item)}')
+        allowed_torch_dtypes = [
+            torch.bool,
+            torch.int64,
+        ]
+        if isinstance(item, torch.Tensor) and (item.dtype in allowed_torch_dtypes):
+            return item.tolist()
+
+        return item
+
+    def _get_from_indices(self: T, item: Iterable[int]) -> T:
+        results = []
+        for ix in item:
+            results.append(self._data[ix])
+        return self.__class__(results)
+
+    def _set_by_indices(self: T, item: Iterable[int], value: Iterable[BaseDocument]):
+        # here we cannot use _get_offset_to_doc() because we need to change the doc
+        # that a given offset points to, not just retrieve it.
+        # Future optimization idea: _data could be List[DocContainer], where
+        # DocContainer points to the doc. Then we could use _get_offset_to_container()
+        # to swap the doc in the container.
+        for ix, doc_to_set in zip(item, value):
+            try:
+                self._data[ix] = doc_to_set
+            except KeyError:
+                raise IndexError(f'Index {ix} is out of range')
+
+    def _get_from_mask(self: T, item: Iterable[bool]) -> T:
+        return self.__class__(
+            (doc for doc, mask_value in zip(self, item) if mask_value)
+        )
+
+    def _set_by_mask(self: T, item: Iterable[bool], value: Sequence[BaseDocument]):
+        i_value = 0
+        for i, mask_value in zip(range(len(self)), item):
+            if mask_value:
+                self._data[i] = value[i_value]
+                i_value += 1
+
     append = _delegate_meth_to_data('append')
     extend = _delegate_meth_to_data('extend')
     insert = _delegate_meth_to_data('insert')

diff --git a/docarray/array/array_stacked.py b/docarray/array/array_stacked.py
@@ -6,10 +6,12 @@
     Iterable,
     List,
     Mapping,
+    Tuple,
     Type,
     TypeVar,
     Union,
     cast,
+    overload,
 )
 
 from docarray.array.abstract_array import AnyDocumentArray
@@ -34,6 +36,7 @@
     TorchTensor = None  # type: ignore
 
 T = TypeVar('T', bound='DocumentArrayStacked')
+IndexIterType = Union[slice, Iterable[int], Iterable[bool], None]
 
 
 class DocumentArrayStacked(AnyDocumentArray):
@@ -72,17 +75,20 @@ def from_document_array(self: T, docs: DocumentArray):
         self._columns = self._create_columns(docs, tensor_type=self.tensor_type)
 
     @classmethod
-    def _from_columns(
+    def _from_da_and_columns(
         cls: Type[T],
         docs: DocumentArray,
         columns: Mapping[str, Union['DocumentArrayStacked', AbstractTensor]],
     ) -> T:
+        """Create a DocumentArrayStacked from a DocumentArray
+        and an associated dict of columns"""
         # below __class_getitem__ is called explicitly instead
         # of doing DocumentArrayStacked[docs.document_type]
         # because mypy has issues with class[...] notation at runtime.
         # see bug here: https://github.com/python/mypy/issues/13026
         # as of 2023-01-05 it should be fixed on mypy master, though, see
         # here: https://github.com/python/typeshed/issues/4819#issuecomment-1354506442
+
         da_stacked = DocumentArray.__class_getitem__(cls.document_type)([]).stack()
         da_stacked._columns = columns
         da_stacked._docs = docs
@@ -199,25 +205,88 @@ def _set_array_attribute(
         else:
             setattr(self._docs, field, values)
 
-    def __getitem__(self, item):  # note this should handle slices
-        if isinstance(item, slice):
-            return self._get_slice(item)
+    @overload
+    def __getitem__(self: T, item: int) -> BaseDocument:
+        ...
+
+    @overload
+    def __getitem__(self: T, item: IndexIterType) -> T:
+        ...
+
+    def __getitem__(self, item):
+        if item is None:
+            return self  # PyTorch behaviour
+        # multiple docs case
+        if isinstance(item, (slice, Iterable)):
+            item_ = cast(Iterable, item)
+            return self._get_from_data_and_columns(item_)
+        # single doc case
         doc = self._docs[item]
         # NOTE: this could be speed up by using a cache
         for field in self._columns.keys():
             setattr(doc, field, self._columns[field][item])
         return doc
 
-    def _get_slice(self: T, item: slice) -> T:
-        """Return a slice of the DocumentArrayStacked
+    def __setitem__(
+        self: T, key: Union[int, IndexIterType], value: Union[T, BaseDocument]
+    ):
+        # multiple docs case
+        if isinstance(key, (slice, Iterable)):
+            return self._set_data_and_columns(key, value)
+        # single doc case
+        doc = self._docs[key]
+        for field in self._columns.keys():
+            setattr(doc, field, self._columns[field][key])
+        return doc
+
+    def _get_from_data_and_columns(self: T, item: Union[Tuple, Iterable]) -> T:
+        """Delegates the access to the data and the columns,
+        and combines into a stacked da.
 
-        :param item: the slice to apply
-        :return: a DocumentArrayStacked
+        :param item: the item used as index. Needs to be a valid index for both
+            DocumentArray (data) and column types (torch/tensorflow/numpy tensors)
+        :return: a DocumentArrayStacked, indexed according to `item`
         """
+        if isinstance(item, tuple):
+            item = list(item)
+        docs_indexed = self._docs[item]
+        columns_indexed = {k: col[item] for k, col in self._columns.items()}
+        columns_indexed_ = cast(Dict[str, Union[AbstractTensor, T]], columns_indexed)
+        return self._from_da_and_columns(docs_indexed, columns_indexed_)
+
+    def _set_data_and_columns(
+        self: T,
+        index_item: Union[Tuple, Iterable, slice],
+        value: Union[T, BaseDocument],
+    ):
+        """Delegates the setting to the data and the columns.
 
-        columns_sliced = {k: col[item] for k, col in self._columns.items()}
-        columns_sliced_ = cast(Dict[str, Union[AbstractTensor, T]], columns_sliced)
-        return self._from_columns(self._docs[item], columns_sliced_)
+        :param index_item: the key used as index. Needs to be a valid index for both
+            DocumentArray (data) and column types (torch/tensorflow/numpy tensors)
+        :value: the value to set at the `key` location
+        """
+        if isinstance(index_item, tuple):
+            index_item = list(index_item)
+        # set data and prepare columns
+        columns_to_set: Dict[str, Union[DocumentArrayStacked, AbstractTensor]]
+        if isinstance(value, DocumentArray):
+            self._docs[index_item] = value
+            columns_to_set = self._create_columns(value, self.tensor_type)
+        elif isinstance(value, BaseDocument):
+            self._docs[index_item] = value
+            columns_to_set = self._create_columns(
+                DocumentArray.__class_getitem__(self.document_type)([value]),
+                self.tensor_type,
+            )
+        elif isinstance(value, DocumentArrayStacked):
+            self._docs[index_item] = value._docs
+            columns_to_set = value._columns
+        else:
+            raise TypeError(f'Can not set a DocumentArrayStacked with {type(value)}')
+        # set columns
+        for col_key, col in self._columns.items():
+            # mypy is confused by a map that points to a union
+            self._columns[col_key][index_item] = columns_to_set[col_key]
 
     def __iter__(self):
         for i in range(len(self)):

diff --git a/docarray/base_document/mixins/proto.py b/docarray/base_document/mixins/proto.py
@@ -8,14 +8,6 @@
     from docarray.proto import DocumentProto, NodeProto
 
 
-try:
-    import torch  # noqa: F401
-except ImportError:
-    torch_imported = False
-else:
-    torch_imported = True
-
-
 T = TypeVar('T', bound='ProtoMixin')