perf(array): use less memory for apply (docarray#312)

alaeddine-13 · web-flow · commit e2f478ac5fe3 · 2022-04-30T08:18:26.000+02:00
diff --git 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(
