Repository: dask/dask

Issue:
Mask selection vs compress() - inconsistent shape
Using dask 1.2.2. In the following example:

```
In [16]: x = da.from_array(np.arange(5), chunks=2)                                                                                                                           

In [17]: x                                                                                                                                                                   
Out[17]: dask.array<array, shape=(5,), dtype=int64, chunksize=(2,)>

In [18]: t = np.array([1, 0, 0, 1, 0], dtype=bool)                                                                                                                           

In [19]: da.compress(t, x, axis=0)                                                                                                                                           
Out[19]: dask.array<getitem, shape=(nan,), dtype=int64, chunksize=(nan,)>

In [20]: x[t]                                                                                                                                                                
Out[20]: dask.array<getitem, shape=(2,), dtype=int64, chunksize=(1,)>
```

...I would expect ``da.compress(t, x, axis=0)`` to give the same result as ``x[t]``, however the shape is ``(nan,)`` in the first case and ``(2,)`` in the second. 

Is this expected behaviour?




Discussion:
Maybe the difference is because ``compress()`` calls ``asarray(condition)`` [here](https://github.com/dask/dask/blob/d8ff4c4d155d51cf275776c1ec8ef27c16d43f17/dask/array/routines.py#L1005). I.e., the ``condition`` array is coerced to a dask array, although it doesn't really need to be if it is a numpy array.

Perhaps we could add a check to allow ``condition`` to be a numpy array, and if so don't coerce to dask array? Happy to make a PR if that sounds reasonable.

cc @jakirkham as I know you rewrote ``compress()`` recently in #4548.
That sounds reasonable to me

On Sat, Jun 15, 2019 at 10:44 PM Alistair Miles <notifications@github.com>
wrote:

> Maybe the difference is because compress() calls asarray(condition) here
> <https://github.com/dask/dask/blob/d8ff4c4d155d51cf275776c1ec8ef27c16d43f17/dask/array/routines.py#L1005>.
> I.e., the condition array is coerced to a dask array, although it doesn't
> really need to be if it is a numpy array.
>
> Perhaps we could add a check to allow condition to be a numpy array, and
> if so don't coerce to dask array? Happy to make a PR if that sounds
> reasonable.
>
> cc @jakirkham <https://github.com/jakirkham> as I know you rewrote
> compress() recently in #4548 <https://github.com/dask/dask/pull/4548>.
>
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/4940?email_source=notifications&email_token=AACKZTHYURYJSGWV3U5JYT3P2VICDA5CNFSM4HYPVODKYY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGODXY7ZLI#issuecomment-502398125>,
> or mute the thread
> <https://github.com/notifications/unsubscribe-auth/AACKZTBYEAIVYE3UMS6USIDP2VICDANCNFSM4HYPVODA>
> .
>


Resolving patch:
diff --git a/dask/array/routines.py b/dask/array/routines.py
--- a/dask/array/routines.py
+++ b/dask/array/routines.py
@@ -14,7 +14,7 @@
 from ..core import flatten
 from ..base import tokenize
 from ..highlevelgraph import HighLevelGraph
-from ..utils import funcname, derived_from
+from ..utils import funcname, derived_from, is_arraylike
 from . import chunk
 from .creation import arange, diag, empty, indices
 from .utils import safe_wraps, validate_axis, meta_from_array, zeros_like_safe
@@ -1002,7 +1002,12 @@ def squeeze(a, axis=None):
 
 @derived_from(np)
 def compress(condition, a, axis=None):
-    condition = asarray(condition).astype(bool)
+
+    if not is_arraylike(condition):
+        # Allow `condition` to be anything array-like, otherwise ensure `condition`
+        # is a numpy array.
+        condition = np.asarray(condition)
+    condition = condition.astype(bool)
     a = asarray(a)
 
     if condition.ndim != 1:


Test patch:
diff --git a/dask/array/tests/test_routines.py b/dask/array/tests/test_routines.py
--- a/dask/array/tests/test_routines.py
+++ b/dask/array/tests/test_routines.py
@@ -950,8 +950,18 @@ def test_compress():
             res = da.compress(dc, a, axis=axis)
             assert_eq(np.compress(c, x, axis=axis), res)
             if isinstance(dc, da.Array):
+                # If condition is a dask array then we expect the shape of the
+                # compressed array to be nan, because we won't know that until
+                # the result is computed.
                 axis = axis or 0
+                assert np.isnan(res.shape[axis]).all()
                 assert np.isnan(res.chunks[axis]).all()
+            else:
+                # If condition is a not a dask array then we expect the shape of the
+                # compressed axis to be known, i.e., not nan.
+                axis = axis or 0
+                assert np.count_nonzero(dc) == res.shape[axis]
+                assert not np.isnan(res.chunks[axis]).any()
 
     with pytest.raises(ValueError):
         da.compress([True, False], a, axis=100)