Repository: dask/dask

Issue:
dask.Bag slow and very large memory consumption with join
In a certain application I am passing a (moderately) large argument to Bag.join. But I am encountering a few issues:
- The graph does not appear in the diagnostics page (presumably the graph is being communicated)
- The client consumes GB of memory

I suspect this has to do with the fact that tasks are serialized individually and this large argument gets serialized multiple times.

This is the code to reproduce the issue: 

```
import dask, distributed
from dask.distributed import Client
from dask.sizeof import sizeof
import dask.bag as db
import time

dask.__version__
# '0.17.1'
distributed.__version__
# '1.21.0'
c = Client()
data = list(range(1000000))
sizeof(data) # 37 Mb

start = time.time()
bag = db.from_sequence(range(100000), npartitions=100).join(data, lambda x: x)
print("Graph creation time", time.time() - start)
# Graph creation time 1.0362498760223389

start = time.time()
bag.count().compute()
print("Graph computation time", time.time() - start)

# This consumes GB of memory
```


Discussion:
First, have you considered using dask.dataframe?  The join operations there will be much faster.  It may be that your data doesn't fit nicely into dask.dataframe, but I thought I'd check.

Resolving patch:
diff --git a/dask/bag/core.py b/dask/bag/core.py
--- a/dask/bag/core.py
+++ b/dask/bag/core.py
@@ -918,21 +918,65 @@ def std(self, ddof=0):
     def join(self, other, on_self, on_other=None):
         """ Joins collection with another collection.
 
-        Other collection must be an Iterable, and not a Bag.
+        Other collection must be one of the following:
 
+        1.  An iterable.  We recommend tuples over lists for internal
+            performance reasons.
+        2.  A delayed object, pointing to a tuple.  This is recommended if the
+            other collection is sizable and you're using the distributed
+            scheduler.  Dask is able to pass around data wrapped in delayed
+            objects with greater sophistication.
+        3.  A Bag with a single partition
+
+        You might also consider Dask Dataframe, whose join operations are much
+        more heavily optimized.
+
+        Parameters
+        ----------
+        other: Iterable, Delayed, Bag
+            Other collection on which to join
+        on_self: callable
+            Function to call on elements in this collection to determine a
+            match
+        on_other: callable (defaults to on_self)
+            Function to call on elements in the other collection to determine a
+            match
+
+        Examples
+        --------
         >>> people = from_sequence(['Alice', 'Bob', 'Charlie'])
         >>> fruit = ['Apple', 'Apricot', 'Banana']
         >>> list(people.join(fruit, lambda x: x[0]))  # doctest: +SKIP
         [('Apple', 'Alice'), ('Apricot', 'Alice'), ('Banana', 'Bob')]
         """
-        assert isinstance(other, Iterable)
-        assert not isinstance(other, Bag)
+        name = 'join-' + tokenize(self, other, on_self, on_other)
+        dsk = {}
+        if isinstance(other, Bag):
+            if other.npartitions == 1:
+                dsk.update(other.dask)
+                dsk['join-%s-other' % name] = (list, other._keys()[0])
+                other = other._keys()[0]
+            else:
+                msg = ("Multi-bag joins are not implemented. "
+                       "We recommend Dask dataframe if appropriate")
+                raise NotImplementedError(msg)
+        elif isinstance(other, Delayed):
+            dsk.update(other.dask)
+            other = other._key
+        elif isinstance(other, Iterable):
+            other = other
+        else:
+            msg = ("Joined argument must be single-partition Bag, "
+                   " delayed object, or Iterable, got %s" %
+                   type(other).__name)
+            raise TypeError(msg)
+
         if on_other is None:
             on_other = on_self
-        name = 'join-' + tokenize(self, other, on_self, on_other)
-        dsk = dict(((name, i), (list, (join, on_other, other,
-                                       on_self, (self.name, i))))
-                   for i in range(self.npartitions))
+
+        dsk.update({(name, i): (list, (join, on_other, other,
+                                       on_self, (self.name, i)))
+                   for i in range(self.npartitions)})
         return type(self)(merge(self.dask, dsk), name, self.npartitions)
 
     def product(self, other):


Test patch:
diff --git a/dask/bag/tests/test_bag.py b/dask/bag/tests/test_bag.py
--- a/dask/bag/tests/test_bag.py
+++ b/dask/bag/tests/test_bag.py
@@ -32,6 +32,15 @@
 b = Bag(dsk, 'x', 3)
 
 
+def assert_eq(a, b):
+    if hasattr(a, 'compute'):
+        a = a.compute(get=dask.local.get_sync)
+    if hasattr(b, 'compute'):
+        b = b.compute(get=dask.local.get_sync)
+
+    assert a == b
+
+
 def iseven(x):
     return x % 2 == 0
 
@@ -353,12 +362,18 @@ def test_var():
     assert float(b.var()) == 2.0
 
 
-def test_join():
-    c = b.join([1, 2, 3], on_self=isodd, on_other=iseven)
-    assert list(c) == list(join(iseven, [1, 2, 3], isodd, list(b)))
-    assert (list(b.join([1, 2, 3], isodd)) ==
-            list(join(isodd, [1, 2, 3], isodd, list(b))))
-    assert c.name == b.join([1, 2, 3], on_self=isodd, on_other=iseven).name
+@pytest.mark.parametrize('transform', [
+    identity,
+    dask.delayed,
+    lambda x: db.from_sequence(x, npartitions=1)
+])
+def test_join(transform):
+    other = transform([1, 2, 3])
+    c = b.join(other, on_self=isodd, on_other=iseven)
+    assert_eq(c, list(join(iseven, [1, 2, 3], isodd, list(b))))
+    assert_eq(b.join(other, isodd),
+              list(join(isodd, [1, 2, 3], isodd, list(b))))
+    assert c.name == b.join(other, on_self=isodd, on_other=iseven).name
 
 
 def test_foldby():