Repository: dask/dask

Issue:
Lazy implementation of compress
The current implementation of Dask Array's [`compress`]( http://dask.pydata.org/en/latest/array-api.html#dask.array.compress ) appears to be eagerly evaluated. Would be nice if this implementation were modified to be lazily evaluated. This would require using some unknown lengths and unknown chunksizes for axes that `compress` operated on.

Discussion:
jcrist: cc @jakirkham.
jcrist: I also added `da.extract` since it was quick to do. Supersedes #2546.

Patch:
diff --git a/dask/array/__init__.py b/dask/array/__init__.py
index 483dbd314..960270369 100644
--- a/dask/array/__init__.py
+++ b/dask/array/__init__.py
@@ -6,7 +6,7 @@ from .core import (Array, stack, concatenate, take, tensordot, transpose,
         roll, fromfunction, unique, store, squeeze, topk, bincount, tile,
         digitize, histogram, map_blocks, atop, to_hdf5, dot, cov, array,
         dstack, vstack, hstack, to_npy_stack, from_npy_stack, compress,
-        from_delayed, round, swapaxes, repeat, asarray, asanyarray)
+        extract, from_delayed, round, swapaxes, repeat, asarray, asanyarray)
 from .core import (around, isnull, notnull, isclose, eye, triu,
                    tril, diag, corrcoef)
 from .reshape import reshape
diff --git a/dask/array/core.py b/dask/array/core.py
index 3446e5723..18848bfc6 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -2451,8 +2451,6 @@ def take(a, indices, axis=0):
 
 @wraps(np.compress)
 def compress(condition, a, axis=None):
-    from .wrap import zeros
-
     if axis is None:
         a = a.ravel()
         axis = 0
@@ -2461,18 +2459,39 @@ def compress(condition, a, axis=None):
     if axis < 0:
         axis += a.ndim
 
-    condition = asarray(condition).astype(bool)
+    # Only coerce non-lazy values to numpy arrays
+    if not isinstance(condition, Array):
+        condition = np.array(condition, dtype=bool)
     if condition.ndim != 1:
         raise ValueError("Condition must be one dimensional")
-    if len(condition) < a.shape[axis]:
-        condition_xtr = zeros(a.shape[axis], dtype=bool, chunks=a.chunks[axis])
-        condition_xtr = condition_xtr[len(condition):]
-        condition = concatenate([condition, condition_xtr])
-
-    condition = np.array(condition)
-    slc = ((slice(None),) * axis + (condition, ) +
-           (slice(None),) * (a.ndim - axis - 1))
-    return a[slc]
+
+    if isinstance(condition, Array):
+        if len(condition) < a.shape[axis]:
+            a = a[tuple(slice(None, len(condition))
+                        if i == axis else slice(None)
+                        for i in range(a.ndim))]
+        inds = tuple(range(a.ndim))
+        out = atop(np.compress, inds, condition, (inds[axis],), a, inds,
+                   axis=axis, dtype=a.dtype)
+        out._chunks = tuple((np.NaN,) * len(c) if i == axis else c
+                            for i, c in enumerate(out.chunks))
+        return out
+    else:
+        # Optimized case when condition is known
+        if len(condition) < a.shape[axis]:
+            condition = condition.copy()
+            condition.resize(a.shape[axis])
+
+        slc = ((slice(None),) * axis + (condition, ) +
+               (slice(None),) * (a.ndim - axis - 1))
+        return a[slc]
+
+
+@wraps(np.extract)
+def extract(condition, arr):
+    if not isinstance(condition, Array):
+        condition = np.array(condition, dtype=bool)
+    return compress(condition.ravel(), arr.ravel())
 
 
 def _take_dask_array_from_numpy(a, indices, axis):
diff --git a/dask/diagnostics/profile.py b/dask/diagnostics/profile.py
index ad41e189e..41543c868 100644
--- a/dask/diagnostics/profile.py
+++ b/dask/diagnostics/profile.py
@@ -134,6 +134,10 @@ class ResourceProfiler(Callback):
     manually.
 
     >>> prof.clear()  # doctest: +SKIP
+
+    Note that when used as a context manager data will be collected throughout
+    the duration of the enclosed block. In contrast, when registered globally
+    data will only be collected while a dask scheduler is active.
     """
     def __init__(self, dt=1):
         self._tracker = _Tracker(dt)


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index 04d9ba159..f97cc2498 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -432,16 +432,20 @@ def test_compress():
     x = np.arange(25).reshape((5, 5))
     a = from_array(x, chunks=(2, 2))
 
-    assert_eq(np.compress([True, False, True, False, True], x),
-              da.compress([True, False, True, False, True], a))
-    assert_eq(np.compress([True, False, True, False, True], x, axis=0),
-              da.compress([True, False, True, False, True], a, axis=0))
-    assert_eq(np.compress([True, False, True, False, True], x, axis=1),
-              da.compress([True, False, True, False, True], a, axis=1))
-    assert_eq(np.compress([True, False], x),
-              da.compress([True, False], a))
-    assert_eq(np.compress([True, False], x, axis=1),
-              da.compress([True, False], a, axis=1))
+    c1 = np.array([True, False, True, False, True])
+    c2 = np.array([True, False])
+    c3 = [True, False]
+    dc1 = da.from_array(c1, chunks=3)
+    dc2 = da.from_array(c2, chunks=2)
+
+    for c, dc in [(c1, c1), (c2, c2), (c3, c3),
+                  (c1, dc1), (c2, dc2), (c3, dc2)]:
+        for axis in [None, 0, 1]:
+            res = da.compress(dc, a, axis=axis)
+            assert_eq(np.compress(c, x, axis=axis), res)
+            if isinstance(dc, Array):
+                axis = axis or 0
+                assert np.isnan(res.chunks[axis]).all()
 
     with pytest.raises(ValueError):
         da.compress([True, False], a, axis=100)
@@ -449,6 +453,25 @@ def test_compress():
         da.compress([[True], [False]], a, axis=100)
 
 
+def test_extract():
+    x = np.arange(25).reshape((5, 5))
+    a = from_array(x, chunks=(2, 2))
+
+    c1 = np.array([True, False, True, False, True])
+    c2 = np.array([[True, False], [True, False]])
+    c3 = np.array([True, False])
+    dc1 = da.from_array(c1, chunks=3)
+    dc2 = da.from_array(c2, chunks=(2, 1))
+    dc3 = da.from_array(c3, chunks=2)
+
+    for c, dc in [(c1, c1), (c2, c2), (c3, c3),
+                  (c1, dc1), (c2, dc2), (c3, dc3)]:
+        res = da.extract(dc, a)
+        assert_eq(np.extract(c, x), res)
+        if isinstance(dc, Array):
+            assert np.isnan(res.chunks[0]).all()
+
+
 def test_binops():
     a = Array(dict((('a', i), np.array([0])) for i in range(3)),
               'a', chunks=((1, 1, 1),), dtype='i8')