Bag.compute returns list

mrocklin · Jun 2, 2015 · 338be88 · 338be88
1 parent f0b3449
commit 338be88
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/dask/bag/core.py b/dask/bag/core.py
@@ -572,8 +572,8 @@ def compute(self, **kwargs):
         results = get(self.dask, self._keys(), **kwargs)
         if isinstance(results[0], Iterable):
             results = toolz.concat(results)
-        if not isinstance(results, Iterator):
-            results = iter(results)
+        if isinstance(results, Iterator):
+            results = list(results)
         return results
 
     def concat(self):
@@ -591,7 +591,8 @@ def concat(self):
                         for i in range(self.npartitions))
         return Bag(merge(self.dask, dsk), name, self.npartitions)
 
-    __iter__ = compute
+    def __iter__(self):
+        return iter(self.compute())
 
     def groupby(self, grouper, npartitions=None):
         """ Group collection by key function

diff --git a/dask/bag/tests/test_bag.py b/dask/bag/tests/test_bag.py
@@ -426,3 +426,9 @@ def f(L):
     c = b.map(f)
 
     assert list(c) == [[2, 3, 4], [5, 6, 7]]
+
+
+def test_ensure_compute_output_is_concrete():
+    b = db.from_sequence([1, 2, 3])
+    result = b.map(lambda x: x + 1).compute()
+    assert not isinstance(result, Iterator)