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:
mrocklin: Merging tomorrow if there are no further comments

Patch:
diff --git a/dask/array/core.py b/dask/array/core.py
index 6f132e1ad..a4990fe95 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -1754,9 +1754,9 @@ class Array(Base):
         return cumprod(self, axis, dtype, out=out)
 
     @derived_from(np.ndarray)
-    def squeeze(self):
+    def squeeze(self, axis=None):
         from .routines import squeeze
-        return squeeze(self)
+        return squeeze(self, axis)
 
     def rechunk(self, chunks, threshold=None, block_size_limit=None):
         """ See da.rechunk for docstring """
diff --git a/dask/array/routines.py b/dask/array/routines.py
index 8f1ea653a..da414204b 100644
--- a/dask/array/routines.py
+++ b/dask/array/routines.py
@@ -5,7 +5,6 @@ import warnings
 from collections import Iterable
 from distutils.version import LooseVersion
 from functools import wraps, partial
-from itertools import product
 from numbers import Integral
 from operator import getitem
 
@@ -836,20 +835,23 @@ def ravel(array):
 
 @wraps(np.squeeze)
 def squeeze(a, axis=None):
-    if 1 not in a.shape:
-        return a
     if axis is None:
         axis = tuple(i for i, d in enumerate(a.shape) if d == 1)
-    b = a.map_blocks(partial(np.squeeze, axis=axis), dtype=a.dtype)
-    chunks = tuple(bd for bd in b.chunks if bd != (1,))
+    elif not isinstance(axis, tuple):
+        axis = (axis,)
+
+    if any(a.shape[i] != 1 for i in axis):
+        raise ValueError("cannot squeeze axis with size other than one")
+
+    for i in axis:
+        if not (-a.ndim <= i < a.ndim):
+            raise ValueError("%i out of bounds for %i-D array" % (i, a.ndim))
 
-    name = 'squeeze-' + tokenize(a, axis)
-    old_keys = list(product([b.name], *[range(len(bd)) for bd in b.chunks]))
-    new_keys = list(product([name], *[range(len(bd)) for bd in chunks]))
+    axis = tuple(i % a.ndim for i in axis)
 
-    dsk = {n: b.dask[o] for o, n in zip(old_keys, new_keys)}
+    sl = tuple(0 if i in axis else slice(None) for i, s in enumerate(a.shape))
 
-    return Array(sharedict.merge(b.dask, (name, dsk)), name, chunks, dtype=a.dtype)
+    return a[sl]
 
 
 def topk(k, x):
diff --git a/dask/bag/core.py b/dask/bag/core.py
index 9ad1ff047..3b30bebf8 100644
--- a/dask/bag/core.py
+++ b/dask/bag/core.py
@@ -918,21 +918,65 @@ class Bag(Base):
     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):
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
index 12a4cda26..56632fa08 100644
--- a/docs/source/changelog.rst
+++ b/docs/source/changelog.rst
@@ -19,10 +19,10 @@ DataFrame
 - Add ``dd.tseries.Resampler.agg`` (:pr:`3202`) `Richard Postelnik`_
 - Support operations that mix dataframes and arrays (:pr:`3230`) `Matthew Rocklin`_
 
-
 Bag
 +++
 
+- Support joining against single-partitioned bags and delayed objects (:pr:`3254`) `Matthew Rocklin`_
 
 Core
 ++++
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 3edb01201..a045f7771 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -190,7 +190,7 @@ distributed memory clusters.
    scheduling.rst
    distributed.rst
 
-**Diagnosing Perofmrance**
+**Diagnosing Performance**
 
 Parallel code can be tricky to debug and profile. Dask provides a several tools
 to help make debugging and profiling graph execution easier.


Test patch:
diff --git a/dask/array/tests/test_routines.py b/dask/array/tests/test_routines.py
index f37b9f817..f084adbd2 100644
--- a/dask/array/tests/test_routines.py
+++ b/dask/array/tests/test_routines.py
@@ -710,13 +710,35 @@ def test_ravel():
     assert_eq(np.ravel(x), da.ravel(a))
 
 
-def test_squeeze():
-    x = da.ones((10, 1), chunks=(3, 1))
+@pytest.mark.parametrize('is_func', [True, False])
+@pytest.mark.parametrize('axis', [None, 0, -1, (0, -1)])
+def test_squeeze(is_func, axis):
+    a = np.arange(10)[None, :, None, None]
+    d = da.from_array(a, chunks=(1, 3, 1, 1))
+
+    if is_func:
+        a_s = np.squeeze(a, axis=axis)
+        d_s = da.squeeze(d, axis=axis)
+    else:
+        a_s = a.squeeze(axis=axis)
+        d_s = d.squeeze(axis=axis)
+
+    assert_eq(d_s, a_s)
+    assert same_keys(d_s, da.squeeze(d, axis=axis))
 
-    assert_eq(x.squeeze(), x.compute().squeeze())
+    if axis is None:
+        axis = tuple(range(a.ndim))
+    else:
+        axis = axis if isinstance(axis, tuple) else (axis,)
+        axis = tuple(i % a.ndim for i in axis)
+    axis = tuple(
+        i for i, c in enumerate(d.chunks) if i in axis and len(c) == 1
+    )
 
-    assert x.squeeze().chunks == ((3, 3, 3, 1),)
-    assert same_keys(x.squeeze(), x.squeeze())
+    exp_d_s_chunks = tuple(
+        c for i, c in enumerate(d.chunks) if i not in axis
+    )
+    assert d_s.chunks == exp_d_s_chunks
 
 
 def test_vstack():
diff --git a/dask/bag/tests/test_bag.py b/dask/bag/tests/test_bag.py
index 004139fba..cd9bf68d2 100644
--- a/dask/bag/tests/test_bag.py
+++ b/dask/bag/tests/test_bag.py
@@ -32,6 +32,15 @@ L = list(range(5)) * 3
 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():