Repository: dask/dask
Pull request: Preserve NumPy condition in da.asarray to preserve output shape
This PR resolves #4940 by allowing the `condition` argument to be a numpy array and leaving it as such.

- [ ] Tests added / passed
- [ ] Passes `flake8 dask`


Problem statement:
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?

Hints from issue 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>
> .
>

Commit message: allow condition to be ndarray and leave it as such

Patch to review:
diff --git a/dask/array/routines.py b/dask/array/routines.py
index 548700bb6f6..ab3201ddfd2 100644
--- a/dask/array/routines.py
+++ b/dask/array/routines.py
@@ -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 isinstance(condition, np.ndarray):
+        # Allow `condition` to be a numpy array, otherwise ensure `condition`
+        # is a dask array.
+        condition = asarray(condition)
+    condition = condition.astype(bool)
     a = asarray(a)
 
     if condition.ndim != 1:
diff --git a/dask/array/tests/test_routines.py b/dask/array/tests/test_routines.py
index d42fc4d6560..2c772a5d326 100644
--- 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()
+            elif isinstance(dc, np.ndarray):
+                # If condition is a numpy array then we expect the shape of the
+                # compressed axis to be a number, 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)


Code review comments:
diff_hunk: @@ -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 isinstance(condition, np.ndarray):
path: dask/array/routines.py
text: @user1:
We're starting to also support other array types, like `sparse.COO` and `cupy.ndarray`, so we try to avoid explicit type checks against `np.ndarray` if possible.  

I wonder in which cases this is necessary, perhaps when we pass in something like a list?  If so, I recommend the following:

```python
from dask.utils import is_arraylike  # up at the top of the file

if not is_arraylike(condition):
    condition = asarray(condition)
```