Merge pull request #349 from nils-braun/bugfix/336-problematic-dicts

Wrap the items passed to `client.scatter` with a dict

Author: martindurant

streamz/dask.py

@@ -5,6 +5,7 @@
 from tornado import gen
 
 from dask.compatibility import apply
+from dask.base import tokenize
 from distributed.client import default_client
 
 from .core import Stream
@@ -103,7 +104,14 @@ def update(self, x, who=None, metadata=None):
         client = default_client()
 
         self._retain_refs(metadata)
-        future = yield client.scatter(x, asynchronous=True)
+        # We need to make sure that x is treated as it is by dask
+        # However, client.scatter works internally different for
+        # lists and dicts. So we always use a dict here to be sure
+        # we know the format exactly. The key will be taken as the
+        # dask identifier of the data.
+        tokenized_x = f"{type(x).__name__}-{tokenize(x)}"
+        future_as_dict = yield client.scatter({tokenized_x: x}, asynchronous=True)
+        future = future_as_dict[tokenized_x]
         f = yield self._emit(future, metadata=metadata)
         self._release_refs(metadata)
 

streamz/tests/test_dask.py

@@ -28,6 +28,28 @@ def test_map(c, s, a, b):
     assert all(isinstance(f, Future) for f in futures_L)
 
 
+@gen_cluster(client=True)
+def test_map_on_dict(c, s, a, b):
+    # dask treats dicts differently, so we have to make sure
+    # the user sees no difference in the streamz api.
+    # Regression test against #336
+    def add_to_dict(d):
+        d["x"] = d["i"]
+        return d
+
+    source = Stream(asynchronous=True)
+    futures = source.scatter().map(add_to_dict)
+    L = futures.gather().sink_to_list()
+
+    for i in range(5):
+        yield source.emit({"i": i})
+
+    assert len(L) == 5
+    for i, item in enumerate(sorted(L, key=lambda x: x["x"])):
+        assert item["x"] == i
+        assert item["i"] == i
+
+
 @gen_cluster(client=True)
 def test_scan(c, s, a, b):
     source = Stream(asynchronous=True)