diff --git a/docarray/array/abstract_array.py b/docarray/array/abstract_array.py index 4812d69d4ca..da7cf397fc4 100644 --- a/docarray/array/abstract_array.py +++ b/docarray/array/abstract_array.py @@ -1,12 +1,18 @@ from abc import abstractmethod -from typing import Iterable, Type +from typing import TYPE_CHECKING, Dict, Iterable, List, Optional, Sequence, Type, Union from docarray.document import BaseDocument +if TYPE_CHECKING: + from docarray.typing import NdArray, TorchTensor -class AbstractDocumentArray(Iterable): + +class AbstractDocumentArray(Sequence): document_type: Type[BaseDocument] + _columns: Optional[ + Dict[str, Union['TorchTensor', 'AbstractDocumentArray', 'NdArray', None]] + ] # here columns are the holder of the data in tensor modes @abstractmethod def __init__(self, docs: Iterable[BaseDocument]): @@ -17,3 +23,11 @@ def __class_getitem__( cls, item: Type[BaseDocument] ) -> Type['AbstractDocumentArray']: ... + + @abstractmethod + def is_stacked(self) -> bool: + ... + + @abstractmethod + def _column_fields(self) -> List[str]: + ... diff --git a/docarray/array/array.py b/docarray/array/array.py index ba8f750463f..c71c8f54830 100644 --- a/docarray/array/array.py +++ b/docarray/array/array.py @@ -1,8 +1,31 @@ -from typing import Iterable, Type +from collections import defaultdict +from contextlib import contextmanager +from functools import wraps +from typing import DefaultDict, Dict, Iterable, List, Optional, Type, TypeVar, Union from docarray.array.abstract_array import AbstractDocumentArray from docarray.array.mixins import GetAttributeArrayMixin, ProtoArrayMixin from docarray.document import AnyDocument, BaseDocument, BaseNode +from docarray.typing import NdArray, TorchTensor + + +def _stacked_mode_blocker(func): + @wraps(func) + def wrapper(self, *args, **kwargs): + if self.is_stacked(): + raise RuntimeError( + f'Cannot call {func.__name__} when the document array is in stack mode' + ) + return func(self, *args, **kwargs) + + wrapper.__doc__ = ( + wrapper.__doc__ + ' \n This method is not available in stacked mode call' + ) + + return wrapper + + +T = TypeVar('T', bound='DocumentArray') class DocumentArray( @@ -13,9 +36,77 @@ class DocumentArray( BaseNode, ): """ - a DocumentArray is a list-like container of Document of the same schema + A DocumentArray is a container of Documents. :param docs: iterable of Document + + A DocumentArray is a list of Documents of any schema. However, many + DocumentArray features are only available if these Documents are + homogeneous and follow the same schema. To be precise, in this schema you can use + the `DocumentArray[MyDocument]` syntax where MyDocument is a Document class + (i.e. schema). This creates a DocumentArray that can only contain Documents of + the type 'MyDocument'. + + EXAMPLE USAGE + .. code-block:: python + from docarray import Document, DocumentArray + from docarray.typing import NdArray, ImageUrl + + + class Image(Document): + tensor: Optional[NdArray[100]] + url: ImageUrl + + + da = DocumentArray[Image]( + Image(url='http://url.com/foo.png') for _ in range(10) + ) # noqa: E510 + + + 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])` + + + A DocumentArray can be in one of two modes: unstacked mode and stacked mode. + + + **Unstacked mode (default)**: + In this case a DocumentArray is a list of Documents and each Document owns its data. + The getter and setter shown above return a list of the fields of each Document + (or DocumentArray if the field is a nested Document). This list/DocumentArray + is created on the fly. The setter sets the field of each Document to the value + of the list/DocumentArray/Tensor passed as parameters. + + This list-like behavior is not always optimal, especially when you want + to process data in batches or perform operations involving matrix computation. + This is where the stack mode of the DocumentArray comes in handy. + + In **stacked mode**, tensor-like fields of every Document are stored on the + DocumentArray level as one stacked tensor. This enables operations on the entire + batch without iterating over the DocumentArray. + In this mode the Documents in the DocumentArray no longer own the data but just + reference the data in the DocumentArray's tensor. + Operations like `da.append` are not supported because + they are too slow. For these operations you should use unstacked + mode instead. + + To switch from stacked to unstacked mode (or vice-versa), call `da.unstack()` or + `da.stack`. There are also two context managers for these modes: + `with da.stack_mode():` and `with da.unstack_mode():` + + see {meth}`.stack` and {meth}`.unstack` for more information. + + You should use unstacked mode if you want to: + + * Insert, append, delete, or shuffle Documents in your DocumentArray. + * Separate the DocumentArray into smaller batches later on. + + You should use stacked mode if you want to: + + * Process the entire DocumentArray as one batch. + * Use the DocumentArray in an ML model for training or inference. + """ document_type: Type[BaseDocument] = AnyDocument @@ -23,24 +114,264 @@ class DocumentArray( def __init__(self, docs: Iterable[BaseDocument]): super().__init__(doc_ for doc_ in docs) + self._columns: Optional[ + Dict[str, Union[TorchTensor, AbstractDocumentArray, NdArray, None]] + ] = None + def __class_getitem__(cls, item: Type[BaseDocument]): if not issubclass(item, BaseDocument): raise ValueError( f'DocumentArray[item] item should be a Document not a {item} ' ) - class _DocumenArrayTyped(DocumentArray): + class _DocumentArrayTyped(cls): # type: ignore document_type: Type[BaseDocument] = item - for field in _DocumenArrayTyped.document_type.__fields__.keys(): + for field in _DocumentArrayTyped.document_type.__fields__.keys(): def _property_generator(val: str): - return property(lambda self: self._get_documents_attribute(val)) + def _getter(self): + return self._get_array_attribute(val) + + def _setter(self, value): + self._set_array_attribute(val, value) - setattr(_DocumenArrayTyped, field, _property_generator(field)) + # need docstring for the property + return property(fget=_getter, fset=_setter) + + setattr(_DocumentArrayTyped, field, _property_generator(field)) # this generates property on the fly based on the schema of the item - _DocumenArrayTyped.__name__ = f'DocumentArray[{item.__name__}]' - _DocumenArrayTyped.__qualname__ = f'DocumentArray[{item.__name__}]' + _DocumentArrayTyped.__name__ = f'DocumentArray[{item.__name__}]' + _DocumentArrayTyped.__qualname__ = f'DocumentArray[{item.__name__}]' + + return _DocumentArrayTyped + + def __list_getitem__(self, item): + return super().__getitem__(item) + + def __getitem__(self, item): # note this should handle slices + if self.is_stacked(): + doc = super().__getitem__(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 + else: + return self.__list_getitem__(item) + + def __list_iter__(self): + return super().__iter__() + + def __iter__(self): + for i in range(len(self)): + yield self[i] + + def stack(self: T) -> T: + """ + Puts the DocumentArray into stacked mode. + :return: itself + + When entering stacked mode DocumentArray creates a column for: + + - Each Tensor field of each Document + - Nested Documents that contain at least one Tensor field. + + This is useful to perform operations on the whole + array at once. In stacked mode, accessing or setting the DocumentArray's fields + accesses or sets the columns of the array. + + EXAMPLE USAGE + .. code-block:: python + from docarray import Document, DocumentArray + from docarray.typing import NdArray + + + class Image(Document): + tensor: NdArray[100] + + + batch = DocumentArray[Image]( + [Image(tensor=np.zeros((100))) for _ in range(10)] + ) # noqa: E510 + + batch.stack() + + print(batch[0].tensor[0]) + # >>> 0 + + print(batch.tensor.shape) + # >>> (10, 3, 224, 224) + + batch.tensor = np.ones((10, 100)) + + print(batch[0].tensor[0]) + # >>> 1 + + batch.append(Image(tensor=np.zeros((100)))) + # >>> raise RuntimeError('Cannot call append when the document array is in + # >>> stack mode' + + see {meth}`.unstack` for more information on how to switch to unstack mode + """ + + if not self.is_stacked(): + + self._columns = dict() + + for field_name, field in self.document_type.__fields__.items(): + if ( + issubclass(field.type_, TorchTensor) + or issubclass(field.type_, BaseDocument) + or issubclass(field.type_, NdArray) + ): + self._columns[field_name] = None + + columns_to_stack: DefaultDict[ + str, Union[List[TorchTensor], List[NdArray], List[BaseDocument]] + ] = defaultdict( # type: ignore + list # type: ignore + ) # type: ignore + + for doc in self.__list_iter__(): + for field_to_stack in self._column_fields(): + columns_to_stack[field_to_stack].append( + getattr(doc, field_to_stack) + ) + setattr(doc, field_to_stack, None) + + for field_to_stack, to_stack in columns_to_stack.items(): + + type_ = self.document_type.__fields__[field_to_stack].type_ + if issubclass(type_, BaseDocument): + self._columns[field_to_stack] = DocumentArray[type_]( # type: ignore # noqa: E501 + to_stack + ).stack() + else: + self._columns[field_to_stack] = type_.__docarray_stack__(to_stack) + + return self + + def unstack(self: T) -> T: + """ + Puts the DocumentArray into unstacked mode. + :return: itself + + + Calling unstack will unstack all columns of the DocumentArray and restore + the data of each Document in the DocumentArray. + + EXAMPLE USAGE + .. code-block:: python + from docarray import Document, DocumentArray, Text + from docarray.typing import NdArray + + + class Image(Document): + tensor: NdArray[100] + + + batch = DocumentArray[Image]( + [Image(tensor=np.zeros((100))) for _ in range(10)] + ) # noqa: E510 + + batch.stack() + batch.unstack() + print(type(batch.tensor)) + # >>> list + print(len(batch.tensor)) + # >>> 10 + + batch.append(Image(tensor=np.zeros((100)))) + # this work just fine + + batch.tensor = np.ones((10, 100)) + # this iterates over the ndarray and assigns each row to a Document + + print(batch[0].tensor[0]) + # >>> 1 + + see {meth}`.stack` for more information on switching to stack mode + """ + if self.is_stacked() and self._columns: + + for i, doc in enumerate(self.__list_iter__()): + for field in self._column_fields(): + val = self._columns[field] + if val is not None: + setattr(doc, field, val[i]) + else: + raise RuntimeError('Internal error, one of the column is None') + # NOTE: here we might need to copy the tensor + # see here + # https://discuss.pytorch.org/t/what-happened-to-a-view-of-a-tensor-when-the-original-tensor-is-deleted/167294 # noqa: E501 + + for field in self._column_fields(): + # list needed here otherwise we are modifying the dict while iterating + del self._columns[field] + + self._columns = None + + return self + + def a(self, x: Union[TorchTensor, NdArray, AbstractDocumentArray, None]): + if x is not None: + x[0] + + @contextmanager + def stacked_mode(self): + """ + Context manager to put the DocumentArray in stacked mode and unstack it when + exiting the context manager. + + EXAMPLE USAGE + .. code-block:: python + with da.stacked_mode(): + ... + """ + try: + yield self.stack() + finally: + self.unstack() + + @contextmanager + def unstacked_mode(self): + """ + Context manager to put the DocumentArray in unstacked mode and stack it when + exiting the context manager. + + EXAMPLE USAGE + .. code-block:: python + with da.unstacked_mode(): + ... + """ + try: + yield self.unstack() + finally: + self.stack() + + def is_stacked(self) -> bool: + """ + Return True if the document array is in stack mode + """ + return self._columns is not None + + def _column_fields(self) -> List[str]: + """ + return the list of fields that are columns of the DocumentArray + :return: the list of keys of the columns + """ + if self.is_stacked() and self._columns is not None: + # need to repeat myself here bc of mypy + return list(self._columns.keys()) + else: + return [] - return _DocumenArrayTyped + append = _stacked_mode_blocker(list.append) + extend = _stacked_mode_blocker(list.extend) + clear = _stacked_mode_blocker(list.clear) + insert = _stacked_mode_blocker(list.insert) + pop = _stacked_mode_blocker(list.pop) + remove = _stacked_mode_blocker(list.remove) + reverse = _stacked_mode_blocker(list.reverse) + sort = _stacked_mode_blocker(list.sort) diff --git a/docarray/array/mixins/attribute.py b/docarray/array/mixins/attribute.py index 712b5b6524e..148111db7f5 100644 --- a/docarray/array/mixins/attribute.py +++ b/docarray/array/mixins/attribute.py @@ -1,25 +1,42 @@ -from typing import List, Union +from typing import TYPE_CHECKING, List, Union from docarray.array.abstract_array import AbstractDocumentArray from docarray.document import BaseDocument +if TYPE_CHECKING: + from docarray.typing import NdArray, TorchTensor + class GetAttributeArrayMixin(AbstractDocumentArray): """Helpers that provide attributes getter in bulk""" - def _get_documents_attribute( - self, field: str - ) -> Union[List, AbstractDocumentArray]: + def _get_array_attribute( + self, + field: str, + ) -> Union[List, AbstractDocumentArray, 'TorchTensor', 'NdArray']: """Return all values of the fields from all docs this array contains :param field: name of the fields to extract :return: Returns a list of the field value for each document in the array like container """ - field_type = self.__class__.document_type._get_nested_document_class(field) - if issubclass(field_type, BaseDocument): + if ( + self.is_stacked() + and field in self._column_fields() + and self._columns is not None + ): + attributes = self._columns[field] + if attributes is not None: + return attributes + else: + raise ValueError( + f'The DocumentArray is in stacked mode, but no stacked data is ' + f'present for {field}. This is inconsistent' + ) + + elif issubclass(field_type, BaseDocument): # calling __class_getitem__ ourselves is a hack otherwise mypy complain # most likely a bug in mypy though # bug reported here https://github.com/python/mypy/issues/14111 @@ -28,3 +45,24 @@ def _get_documents_attribute( ) else: return [getattr(doc, field) for doc in self] + + def _set_array_attribute( + self, + field: str, + values: Union[List, AbstractDocumentArray, 'TorchTensor', 'NdArray'], + ): + """Set all Documents in this DocumentArray using the passed values + + :param field: name of the fields to extract + :values: the values to set at the DocumentArray level + """ + if ( + self.is_stacked() + and self._columns is not None + and field in self._column_fields() + and not isinstance(values, List) + ): + self._columns[field] = values + else: + for doc, value in zip(self, values): + setattr(doc, field, value) diff --git a/docarray/typing/tensor/abstract_tensor.py b/docarray/typing/tensor/abstract_tensor.py index a792f329896..f96c075a39c 100644 --- a/docarray/typing/tensor/abstract_tensor.py +++ b/docarray/typing/tensor/abstract_tensor.py @@ -1,6 +1,6 @@ import abc from abc import ABC -from typing import TYPE_CHECKING, Any, Generic, Tuple, Type, TypeVar +from typing import TYPE_CHECKING, Any, Generic, List, Tuple, Type, TypeVar, Union from docarray.typing.abstract_type import AbstractType @@ -91,3 +91,9 @@ def validate( def __class_getitem__(cls, item: Any): target_shape = cls.__validate_getitem__(item) return cls._create_parametrized_type(target_shape) + + @classmethod + @abc.abstractmethod + def __docarray_stack__(cls: Type[T], seq: Union[List[T], Tuple[T]]) -> T: + """Stack a sequence of tensors into a single tensor.""" + ... diff --git a/docarray/typing/tensor/ndarray.py b/docarray/typing/tensor/ndarray.py index 1e0ee9174fd..af5092f086c 100644 --- a/docarray/typing/tensor/ndarray.py +++ b/docarray/typing/tensor/ndarray.py @@ -191,3 +191,10 @@ def _flush_tensor_to_proto(pb_msg: 'NdArrayProto', value: 'NdArray'): pb_msg.dense.ClearField('shape') pb_msg.dense.shape.extend(list(value.shape)) pb_msg.dense.dtype = value.dtype.str + + @classmethod + def __docarray_stack__( + cls: Type[T], seq: Union[List['NdArray'], Tuple['NdArray']] + ) -> T: + """Stack a sequence of ndarrays into a single ndarray.""" + return cls.from_ndarray(np.stack(seq)) diff --git a/docarray/typing/tensor/torch_tensor.py b/docarray/typing/tensor/torch_tensor.py index b564908be75..47522fd834c 100644 --- a/docarray/typing/tensor/torch_tensor.py +++ b/docarray/typing/tensor/torch_tensor.py @@ -1,6 +1,17 @@ import warnings from copy import copy -from typing import TYPE_CHECKING, Any, Dict, Generic, Tuple, Type, TypeVar, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Generic, + List, + Tuple, + Type, + TypeVar, + Union, + cast, +) import numpy as np import torch # type: ignore @@ -73,6 +84,8 @@ class MyDoc(Document): """ + __parametrized_meta__ = metaTorchAndNode + @classmethod def __get_validators__(cls): # one or more validators may be yielded which will be called in the @@ -80,8 +93,6 @@ def __get_validators__(cls): # the value returned from the previous validator yield cls.validate - __parametrized_meta__ = metaTorchAndNode - @classmethod def __validate_shape__(cls, t: T, shape: Tuple[int]) -> T: # type: ignore if t.shape == shape: @@ -212,3 +223,10 @@ def _flush_tensor_to_proto(pb_msg: 'NdArrayProto', value: 'TorchTensor'): pb_msg.dense.ClearField('shape') pb_msg.dense.shape.extend(list(value_np.shape)) pb_msg.dense.dtype = value_np.dtype.str + + @classmethod + def __docarray_stack__( + cls: Type[T], seq: Union[Tuple['TorchTensor'], List['TorchTensor']] + ) -> T: + """Stack a sequence of ndarray into a single ndarray.""" + return cls.from_native_torch_tensor(torch.stack(seq)) # type: ignore diff --git a/tests/integrations/array/test_torch_train.py b/tests/integrations/array/test_torch_train.py new file mode 100644 index 00000000000..0b25bd190f0 --- /dev/null +++ b/tests/integrations/array/test_torch_train.py @@ -0,0 +1,35 @@ +from typing import Optional + +import torch + +from docarray import Document, DocumentArray +from docarray.typing import TorchTensor + + +def test_torch_train(): + class Mmdoc(Document): + text: str + tensor: Optional[TorchTensor[3, 224, 224]] + + N = 10 + + batch = DocumentArray[Mmdoc](Mmdoc(text=f'hello{i}') for i in range(N)) + + batch.tensor = torch.zeros(N, 3, 224, 224) + batch.stack() + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(3, 16, 3) + + def forward(self, x): + return self.conv(x) + + model = Model() + opt = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9) + + for _ in range(2): + loss = model(batch.tensor).sum() + loss.backward() + opt.step() diff --git a/tests/units/array/test_mixins/test_attribute.py b/tests/units/array/test_mixins/test_attribute.py index 1a168bcd0f0..4472a91ff57 100644 --- a/tests/units/array/test_mixins/test_attribute.py +++ b/tests/units/array/test_mixins/test_attribute.py @@ -16,19 +16,39 @@ class Mmdoc(BaseDocument): (Mmdoc(text=f'hello{i}', tensor=np.zeros((3, 224, 224))) for i in range(N)) ) - tensors = da._get_documents_attribute('tensor') + tensors = da._get_array_attribute('tensor') assert len(tensors) == N for tensor in tensors: assert tensor.shape == (3, 224, 224) - texts = da._get_documents_attribute('text') + texts = da._get_array_attribute('text') assert len(texts) == N for i, text in enumerate(texts): assert text == f'hello{i}' +def test_set_attributes(): + class InnerDoc(BaseDocument): + text: str + + class Mmdoc(BaseDocument): + inner: InnerDoc + + N = 10 + + da = DocumentArray[Mmdoc]( + (Mmdoc(inner=InnerDoc(text=f'hello{i}')) for i in range(N)) + ) + + list_docs = [InnerDoc(text=f'hello{i}') for i in range(N)] + da._set_array_attribute('inner', list_docs) + + for doc, list_doc in zip(da, list_docs): + assert doc.inner is list_doc + + def test_get_bulk_attributes(): class Mmdoc(BaseDocument): text: str diff --git a/tests/units/array/test_stack_mod.py b/tests/units/array/test_stack_mod.py new file mode 100644 index 00000000000..3fe20c07ddf --- /dev/null +++ b/tests/units/array/test_stack_mod.py @@ -0,0 +1,178 @@ +from typing import Optional + +import numpy as np +import pytest +import torch + +from docarray import Document, DocumentArray +from docarray.typing import NdArray, TorchTensor + + +def test_stack(): + class Image(Document): + tensor: TorchTensor[3, 224, 224] + + batch = DocumentArray[Image]( + [Image(tensor=torch.zeros(3, 224, 224)) for _ in range(10)] + ) + + batch.stack() + + assert (batch._columns['tensor'] == torch.zeros(10, 3, 224, 224)).all() + assert (batch.tensor == torch.zeros(10, 3, 224, 224)).all() + assert batch._columns['tensor'].data_ptr() == batch.tensor.data_ptr() + + for doc, tensor in zip(batch, batch.tensor): + assert doc.tensor.data_ptr() == tensor.data_ptr() + + for i in range(len(batch)): + assert batch[i].tensor.data_ptr() == batch.tensor[i].data_ptr() + + +def test_stack_mod_nested_document(): + class Image(Document): + tensor: TorchTensor[3, 224, 224] + + class MMdoc(Document): + img: Image + + batch = DocumentArray[MMdoc]( + [MMdoc(img=Image(tensor=torch.zeros(3, 224, 224))) for _ in range(10)] + ) + + batch.stack() + + assert ( + batch._columns['img']._columns['tensor'] == torch.zeros(10, 3, 224, 224) + ).all() + + assert (batch.img.tensor == torch.zeros(10, 3, 224, 224)).all() + + assert ( + batch._columns['img']._columns['tensor'].data_ptr() + == batch.img.tensor.data_ptr() + ) + + +def test_unstack(): + class Image(Document): + tensor: TorchTensor[3, 224, 224] + + batch = DocumentArray[Image]( + [Image(tensor=torch.zeros(3, 224, 224)) for _ in range(10)] + ) + + batch.stack() + + batch.unstack() + + for doc in batch: + assert (doc.tensor == torch.zeros(3, 224, 224)).all() + + +def test_unstack_nested_document(): + class Image(Document): + tensor: TorchTensor[3, 224, 224] + + class MMdoc(Document): + img: Image + + batch = DocumentArray[MMdoc]( + [MMdoc(img=Image(tensor=torch.zeros(3, 224, 224))) for _ in range(10)] + ) + + batch.stack() + + batch.unstack() + + +def test_stack_runtime_error(): + class Image(Document): + tensor: TorchTensor[3, 224, 224] + + batch = DocumentArray[Image]( + [Image(tensor=torch.zeros(3, 224, 224)) for _ in range(10)] + ) + + batch.stack() + + with pytest.raises(RuntimeError): + batch.append([]) + + +def test_context_stack(): + class Image(Document): + tensor: TorchTensor[3, 224, 224] + + batch = DocumentArray[Image]( + [Image(tensor=torch.zeros(3, 224, 224)) for _ in range(10)] + ) + + assert not (batch.is_stacked()) + with batch.stacked_mode(): + assert batch.is_stacked() + + assert not (batch.is_stacked()) + + +def test_context_not_stack(): + class Image(Document): + tensor: TorchTensor[3, 224, 224] + + batch = DocumentArray[Image]( + [Image(tensor=torch.zeros(3, 224, 224)) for _ in range(10)] + ) + + batch.stack() + + assert batch.is_stacked() + with batch.unstacked_mode(): + assert not (batch.is_stacked()) + + assert batch.is_stacked() + + +def test_stack_setter(): + class Image(Document): + tensor: TorchTensor[3, 224, 224] + + batch = DocumentArray[Image]( + [Image(tensor=torch.zeros(3, 224, 224)) for _ in range(10)] + ) + + batch.stack() + + batch.tensor = torch.ones(10, 3, 224, 224) + + assert (batch.tensor == torch.ones(10, 3, 224, 224)).all() + + +def test_stack_optional(): + class Image(Document): + tensor: Optional[TorchTensor[3, 224, 224]] + + batch = DocumentArray[Image]( + [Image(tensor=torch.zeros(3, 224, 224)) for _ in range(10)] + ) + + batch.stack() + + assert (batch._columns['tensor'] == torch.zeros(10, 3, 224, 224)).all() + assert (batch.tensor == torch.zeros(10, 3, 224, 224)).all() + + +def test_stack_numpy(): + class Image(Document): + tensor: NdArray[3, 224, 224] + + batch = DocumentArray[Image]( + [Image(tensor=np.zeros((3, 224, 224))) for _ in range(10)] + ) + + batch.stack() + + assert (batch._columns['tensor'] == np.zeros((10, 3, 224, 224))).all() + assert (batch.tensor == np.zeros((10, 3, 224, 224))).all() + assert batch.tensor.ctypes.data == batch._columns['tensor'].ctypes.data + + batch.unstack()