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:
Should add this seems useful for implementing things like `nonzero`. ( https://github.com/dask/dask/pull/2539 )
AFAICT this comes down to these [lines of code]( https://github.com/dask/dask/blob/0.15.1/dask/array/core.py#L2461-L2466 ). TBH this comes down to two issues. First there seems to be there is no [`resize`]( https://docs.scipy.org/doc/numpy/reference/generated/numpy.ndarray.resize.html#numpy.ndarray.resize ) implemented for Dask Arrays, which could potentially be worked around. Second this needs to use [slices]( https://github.com/dask/dask/blob/0.15.1/dask/array/core.py#L2468-L2470 ). Though it is less clear how to fix this latter constraint.
> First there seems to be there is no [`resize`]( https://docs.scipy.org/doc/numpy/reference/generated/numpy.ndarray.resize.html#numpy.ndarray.resize ) implemented for Dask Arrays, which could potentially be worked around.

Addressing this part of the problem in PR ( https://github.com/dask/dask/pull/2545 ).
> Second this needs to use [slices]( https://github.com/dask/dask/blob/0.15.1/dask/array/core.py#L2468-L2470 ).

It might be possible to get around this if we could slice one Dask Array with another. ( https://github.com/dask/dask/issues/2412 ) There appears to be an attempt to allow this exact sort of slicing in PR ( https://github.com/dask/dask/pull/2430 ), but it looks stalled.
Maybe it will be easier to tackle this problem if we focus simply on the 1-D array case as is seen with [`extract`]( https://docs.scipy.org/doc/numpy/reference/generated/numpy.extract.html ). Attempting to add `extract` for Dask Array's in PR ( https://github.com/dask/dask/pull/2546 ). While this implementation is not lazy yet, hopefully it is easier to make a lazy implementation of `extract`. That way, it can be leveraged by `compress`.
I tried my best to address this in PR ( https://github.com/dask/dask/pull/2548 ), but am not really thrilled with the result. Perhaps someone can use that to build a better solution.

Resolving patch:
diff --git a/dask/array/__init__.py b/dask/array/__init__.py
--- a/dask/array/__init__.py
+++ b/dask/array/__init__.py
@@ -6,7 +6,7 @@
         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
--- 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):


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
--- 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')