fix: apply suggestion

Signed-off-by: anna-charlotte <[email protected]>
docarray · anna-charlotte · Mar 3, 2023 · Feb 28, 2023 · Feb 28, 2023 · Feb 28, 2023
commit f6921e07b620ba0f69c8b500e51ce4fafacea45a
diff --git a/docarray/utils/map.py b/docarray/utils/map.py
@@ -72,7 +72,7 @@ def map_docs(
             yield x
 
 
-def map_batch(
+def map_docs_batch(
     da: T,
     func: Callable[[T], Union[T, T_doc]],
     batch_size: int,

diff --git a/tests/benchmark_tests/test_map.py b/tests/benchmark_tests/test_map.py
@@ -7,7 +7,7 @@
 from docarray import BaseDocument, DocumentArray
 from docarray.documents import Image
 from docarray.typing import NdArray
-from docarray.utils.map import map_batch, map_docs
+from docarray.utils.map import map_docs, map_docs_batch
 from tests.units.typing.test_bytes import IMAGE_PATHS
 
 pytestmark = [pytest.mark.benchmark, pytest.mark.slow]
@@ -66,7 +66,7 @@ def time_multiprocessing(num_workers: int) -> float:
             da = DocumentArray[MyMatrix]([MyMatrix(matrix=m) for m in matrices])
             start_time = time()
             list(
-                map_batch(
+                map_docs_batch(
                     da=da,
                     func=cpu_intensive_batch,
                     batch_size=8,
@@ -121,7 +121,7 @@ def time_multithreading_batch(num_workers: int) -> float:
         )
         start_time = time()
         list(
-            map_batch(
+            map_docs_batch(
                 da=da,
                 func=io_intensive_batch,
                 backend='thread',

diff --git a/tests/units/util/test_map.py b/tests/units/util/test_map.py
@@ -5,7 +5,7 @@
 from docarray import BaseDocument, DocumentArray
 from docarray.documents import Image
 from docarray.typing import ImageUrl, NdArray
-from docarray.utils.map import map_batch, map_docs
+from docarray.utils.map import map_docs, map_docs_batch
 from tests.units.typing.test_bytes import IMAGE_PATHS
 
 N_DOCS = 2
@@ -73,12 +73,14 @@ class MyImage(BaseDocument):
 @pytest.mark.slow
 @pytest.mark.parametrize('n_docs,batch_size', [(10, 5), (10, 8)])
 @pytest.mark.parametrize('backend', ['thread', 'process'])
-def test_map_batch(n_docs, batch_size, backend):
+def test_map_docs_batch(n_docs, batch_size, backend):
 
     da = DocumentArray[MyImage](
         [MyImage(url=IMAGE_PATHS['png']) for _ in range(n_docs)]
     )
-    it = map_batch(da=da, func=load_from_da, batch_size=batch_size, backend=backend)
+    it = map_docs_batch(
+        da=da, func=load_from_da, batch_size=batch_size, backend=backend
+    )
     assert isinstance(it, Generator)
 
     for batch in it: