diff --git a/docarray/array/mixins/parallel.py b/docarray/array/mixins/parallel.py index acc588bde13..4b215d33617 100644 --- a/docarray/array/mixins/parallel.py +++ b/docarray/array/mixins/parallel.py @@ -60,11 +60,8 @@ def apply(self: 'T', *args, **kwargs) -> 'T': # noqa: DAR201 :return: a new :class:`DocumentArray` """ - from ... import DocumentArray - - new_da = DocumentArray(self.map(*args, **kwargs)) - self.clear() - self.extend(new_da) + for doc in self.map(*args, **kwargs): + self[doc.id] = doc return self def map( @@ -157,13 +154,8 @@ def apply_batch(self: 'T', *args, **kwargs) -> 'T': # noqa: DAR201 :return: a new :class:`DocumentArray` """ - from ... import DocumentArray - - new_da = DocumentArray() for _b in self.map_batch(*args, **kwargs): - new_da.extend(_b) - self.clear() - self.extend(new_da) + self[[doc.id for doc in _b]] = _b return self def map_batch(