Repository: dask/distributed

Issue:
Support iterators in Map
We now support iterators/queues in scatter/gather.  To complete this story we should support iterators and queues in Map.  These should return iterators/queues to match the input type.  This should support infinite streams of data.



Discussion:
With this and the iterator/queue gather we could implement a joblib-like map with auto-batching fairly easy.


Resolving patch:
diff --git a/distributed/executor.py b/distributed/executor.py
--- a/distributed/executor.py
+++ b/distributed/executor.py
@@ -27,7 +27,7 @@
 from .client import (WrappedKey, unpack_remotedata, pack_data)
 from .core import read, write, connect, rpc, coerce_to_rpc
 from .scheduler import Scheduler
-from .utils import All, sync, funcname, ignoring
+from .utils import All, sync, funcname, ignoring, queue_to_iterator
 from .compatibility import Queue as pyQueue, Empty, isqueue
 
 logger = logging.getLogger(__name__)
@@ -505,6 +505,20 @@ def submit(self, func, *args, **kwargs):
 
         return Future(key, self)
 
+    def _threaded_map(self, q_out, func, qs_in, **kwargs):
+        """ Internal function for mapping Queue """
+        if isqueue(qs_in[0]):
+            get = pyQueue.get
+        elif isinstance(qs_in[0], Iterator):
+            get = next
+        else:
+            raise NotImplementedError()
+
+        while True:
+            args = [get(q) for q in qs_in]
+            f = self.submit(func, *args, **kwargs)
+            q_out.put(f)
+
     def map(self, func, *iterables, **kwargs):
         """ Map a function on a sequence of arguments
 
@@ -513,7 +527,7 @@ def map(self, func, *iterables, **kwargs):
         Parameters
         ----------
         func: callable
-        iterables: Iterables
+        iterables: Iterables, Iterators, or Queues
         pure: bool (defaults to True)
             Whether or not the function is pure.  Set ``pure=False`` for
             impure functions like ``np.random.random``.
@@ -527,12 +541,28 @@ def map(self, func, *iterables, **kwargs):
 
         Returns
         -------
-        list of futures
+        List, iterator, or Queue of futures, depending on the type of the
+        inputs.
 
         See also
         --------
         Executor.submit: Submit a single function
         """
+        if not callable(func):
+            raise TypeError("First input to map must be a callable function")
+
+        if (all(map(isqueue, iterables)) or
+            all(isinstance(i, Iterator) for i in iterables)):
+            q_out = pyQueue()
+            t = Thread(target=self._threaded_map, args=(q_out, func, iterables),
+                                                  kwargs=kwargs)
+            t.daemon = True
+            t.start()
+            if isqueue(iterables[0]):
+                return q_out
+            else:
+                return queue_to_iterator(q_out)
+
         pure = kwargs.pop('pure', True)
         workers = kwargs.pop('workers', None)
         allow_other_workers = kwargs.pop('allow_other_workers', False)
@@ -540,8 +570,6 @@ def map(self, func, *iterables, **kwargs):
         if allow_other_workers and workers is None:
             raise ValueError("Only use allow_other_workers= if using workers=")
 
-        if not callable(func):
-            raise TypeError("First input to map must be a callable function")
         iterables = [list(it) for it in iterables]
         if pure:
             keys = [funcname(func) + '-' + tokenize(func, kwargs, *args)
@@ -756,16 +784,10 @@ def scatter(self, data, workers=None, broadcast=False):
             t.daemon = True
             t.start()
 
-            if not isqueue(data) and isinstance(data, Iterator):
-                def _():
-                    while True:
-                        result = qout.get()
-                        if result == StopIteration:
-                            break
-                        yield result
-                return _()
-            else:
+            if isqueue(data):
                 return qout
+            else:
+                return queue_to_iterator(qout)
         else:
             return sync(self.loop, self._scatter, data, workers=workers,
                         broadcast=broadcast)
diff --git a/distributed/utils.py b/distributed/utils.py
--- a/distributed/utils.py
+++ b/distributed/utils.py
@@ -9,12 +9,14 @@
 import sys
 import tblib.pickling_support
 import tempfile
+from threading import Thread
 import traceback
 
 from dask import istask
 from toolz import memoize
 from tornado import gen
 
+from .compatibility import Queue
 
 logger = logging.getLogger(__name__)
 
@@ -255,6 +257,27 @@ def truncate_exception(e, n=10000):
         return e
 
 
+def queue_to_iterator(q):
+    while True:
+        result = q.get()
+        if result == StopIteration:
+            break
+        yield result
+
+def _dump_to_queue(seq, q):
+    for item in seq:
+        q.put(item)
+
+def iterator_to_queue(seq, maxsize=0):
+    q = Queue(maxsize=maxsize)
+
+    t = Thread(target=_dump_to_queue, args=(seq, q))
+    t.daemon = True
+    t.start()
+
+    return q
+
+
 import logging
 logging.basicConfig(format='%(name)s - %(levelname)s - %(message)s',
                     level=logging.INFO)


Test patch:
diff --git a/distributed/tests/test_executor.py b/distributed/tests/test_executor.py
--- a/distributed/tests/test_executor.py
+++ b/distributed/tests/test_executor.py
@@ -28,7 +28,7 @@
 from distributed.utils import ignoring, sync, tmp_text
 from distributed.utils_test import (cluster, cluster_center, slow,
         _test_cluster, _test_scheduler, loop, inc, dec, div, throws,
-        gen_cluster, gen_test)
+        gen_cluster, gen_test, double)
 
 
 @gen_cluster()
@@ -2108,3 +2108,52 @@ def test_traceback_clean(s, a, b):
             assert 'scheduler' not in tb.tb_frame.f_code.co_filename
             assert 'worker' not in tb.tb_frame.f_code.co_filename
             tb = tb.tb_next
+
+
+@gen_cluster()
+def test_map_queue(s, a, b):
+    from distributed.compatibility import Queue, isqueue
+    e = Executor((s.ip, s.port), start=False)
+    yield e._start()
+
+    q_1 = Queue(maxsize=2)
+    q_2 = e.map(inc, q_1)
+    assert isqueue(q_2)
+    q_3 = e.map(double, q_2)
+    assert isqueue(q_3)
+    q_4 = yield e._gather(q_3)
+    assert isqueue(q_4)
+
+    q_1.put(1)
+
+    f = q_4.get()
+    assert isinstance(f, Future)
+    result = yield f._result()
+    assert result == (1 + 1) * 2
+
+    yield e._shutdown()
+
+
+@gen_cluster()
+def test_map_iterator(s, a, b):
+    e = Executor((s.ip, s.port), start=False)
+    yield e._start()
+
+    x = iter([1, 2, 3])
+    y = iter([10, 20, 30])
+    f1 = e.map(add, x, y)
+    assert isinstance(f1, Iterator)
+
+    start = time()  # ensure that we compute eagerly
+    while not s.dask:
+        yield gen.sleep(0.01)
+        assert time() < start + 5
+
+    f2 = e.map(double, f1)
+    assert isinstance(f2, Iterator)
+
+    future = next(f2)
+    result = yield future._result()
+    assert result == (1 + 10) * 2
+
+    yield e._shutdown()
diff --git a/distributed/tests/test_utils.py b/distributed/tests/test_utils.py
--- a/distributed/tests/test_utils.py
+++ b/distributed/tests/test_utils.py
@@ -1,6 +1,4 @@
-from distributed.utils import (All, sync, is_kernel, ensure_ip,
-        truncate_exception, get_traceback)
-from distributed.utils_test import loop, inc, throws, div
+from collections import Iterator
 import pytest
 from threading import Thread
 import threading
@@ -9,6 +7,11 @@
 from tornado.locks import Event
 from time import time, sleep
 
+from distributed.utils import (All, sync, is_kernel, ensure_ip,
+        truncate_exception, get_traceback, queue_to_iterator, iterator_to_queue)
+from distributed.utils_test import loop, inc, throws, div
+from distributed.compatibility import Queue, isqueue
+
 def test_All(loop):
     @gen.coroutine
     def throws():
@@ -126,3 +129,22 @@ def c(x):
     except Exception as e:
         tb = get_traceback()
         assert type(tb).__name__ == 'traceback'
+
+
+def test_queue_to_iterator():
+    q = Queue()
+    q.put(1)
+    q.put(2)
+
+    seq = queue_to_iterator(q)
+    assert isinstance(seq, Iterator)
+    assert next(seq) == 1
+    assert next(seq) == 2
+
+
+def test_iterator_to_queue():
+    seq = iter([1, 2, 3])
+
+    q = iterator_to_queue(seq)
+    assert isqueue(q)
+    assert q.get() == 1
diff --git a/distributed/utils_test.py b/distributed/utils_test.py
--- a/distributed/utils_test.py
+++ b/distributed/utils_test.py
@@ -48,6 +48,10 @@ def throws(x):
     raise Exception()
 
 
+def double(x):
+    return x * 2
+
+
 def run_center(q):
     from distributed import Center
     from tornado.ioloop import IOLoop, PeriodicCallback