Allow for map_async to run items in parallel

Use an asyncio.Queue of the tasks to ensure that arrival and departure
order of elements match. Asserts back pressure when a new value arrives
via update but the work queue is full.

Because asyncio.Queue cannot peak, the parallelism factor is not
precise as the worker callback can have either zero or one task in hand
but it must free up a slot in the queue to do so. Under pressure, the
parallelism will generally be `(parallelism + 1)` instead of
`parallelism` as given in the `__init__` as one Future will be in the
awaited in the worker callback while the queue fills up from update
calls.

Author: dwskoog

streamz/core.py

@@ -720,15 +720,16 @@ def update(self, x, who=None, metadata=None):
 
 @Stream.register_api()
 class map_async(Stream):
-    """ Apply an async function to every element in the stream
+    """ Apply an async function to every element in the stream, preserving order
+    even when evaluating multiple inputs in parallel.
 
     Parameters
     ----------
     func: async callable
     *args :
         The arguments to pass to the function.
-    buffer_size:
-        The max size of the input buffer, default value is unlimited
+    parallelism:
+        The maximum number of parallel Tasks for evaluating func, default value is 1
     **kwargs:
         Keyword arguments to pass to func
 
@@ -747,32 +748,31 @@ class map_async(Stream):
     4
     6
     8
-
     """
-    def __init__(self, upstream, func, *args, buffer_size=0, **kwargs):
+    def __init__(self, upstream, func, *args, parallelism=1, **kwargs):
         self.func = func
         stream_name = kwargs.pop('stream_name', None)
         self.kwargs = kwargs
         self.args = args
-        self.input_queue = asyncio.Queue(maxsize=buffer_size)
+        self.work_queue = asyncio.Queue(maxsize=parallelism)
 
         Stream.__init__(self, upstream, stream_name=stream_name, ensure_io_loop=True)
-        self.input_task = self._create_task(self.input_callback())
+        self.work_task = self._create_task(self.work_callback())
 
     def update(self, x, who=None, metadata=None):
-        coro = self.func(x, *self.args, **self.kwargs)
-        self._retain_refs(metadata)
-        return self._create_task(self.input_queue.put((coro, metadata)))
+        return self._create_task(self._insert_job(x, metadata))
 
     def _create_task(self, coro):
+        if gen.is_future(coro):
+            return coro
         return self.loop.asyncio_loop.create_task(coro)
 
-    async def input_callback(self):
+    async def work_callback(self):
         while True:
             try:
-                coro, metadata = await self.input_queue.get()
-                self.input_queue.task_done()
-                result = await coro
+                task, metadata = await self.work_queue.get()
+                self.work_queue.task_done()
+                result = await task
             except Exception as e:
                 logger.exception(e)
                 raise
@@ -782,6 +782,21 @@ async def input_callback(self):
                     await asyncio.gather(*results)
                 self._release_refs(metadata)
 
+    async def _wait_for_work_slot(self):
+        while self.work_queue.full():
+            await asyncio.sleep(0)
+
+    async def _insert_job(self, x, metadata):
+        try:
+            await self._wait_for_work_slot()
+            coro = self.func(x, *self.args, **self.kwargs)
+            task = self._create_task(coro)
+            await self.work_queue.put((task, metadata))
+            self._retain_refs(metadata)
+        except Exception as e:
+            logger.exception(e)
+            raise
+
 
 @Stream.register_api()
 class starmap(Stream):

streamz/tests/test_core.py

@@ -133,11 +133,13 @@ def add_tor(x=0, y=0):
         return x + y
 
     async def add_native(x=0, y=0):
+        await asyncio.sleep(0.1)
         return x + y
 
     source = Stream(asynchronous=True)
-    L = source.map_async(add_tor, y=1).map_async(add_native, y=2).sink_to_list()
+    L = source.map_async(add_tor, y=1).map_async(add_native, parallelism=2, y=2).buffer(1).sink_to_list()
 
+    start = time()
     yield source.emit(0)
     yield source.emit(1)
     yield source.emit(2)
@@ -146,6 +148,7 @@ def fail_func():
         assert L == [3, 4, 5]
 
     yield await_for(lambda: L == [3, 4, 5], 1, fail_func=fail_func)
+    assert (time() - start) == pytest.approx(0.1, abs=4e-3)
 
 
 @pytest.mark.asyncio
@@ -155,11 +158,13 @@ def add_tor(x=0, y=0):
         return x + y
 
     async def add_native(x=0, y=0):
+        await asyncio.sleep(0.1)
         return x + y
 
     source = Stream(asynchronous=True)
-    L = source.map_async(add_tor, y=1).map_async(add_native, y=2).sink_to_list()
+    L = source.map_async(add_tor, y=1).map_async(add_native, parallelism=2, y=2).sink_to_list()
 
+    start = time()
     await source.emit(0)
     await source.emit(1)
     await source.emit(2)
@@ -168,6 +173,7 @@ def fail_func():
         assert L == [3, 4, 5]
 
     await await_for(lambda: L == [3, 4, 5], 1, fail_func=fail_func)
+    assert (time() - start) == pytest.approx(0.1, abs=4e-3)
 
 
 def test_map_args():