Repository: dask/dask

Issue:
std() fails on masked array of all masked elements
```python
In [1]: import numpy as np                                                                                                                                                                                 

In [2]: import dask.array as da                                                                                                                                                                            

In [3]: all_masked = np.ma.masked_array([1, 2], mask=[True, True])                                                                                                                                         

In [4]: all_masked.std()                                                                                                                                                                                   
Out[4]: masked

In [5]: dask_all_masked = da.ma.masked_array([1, 2], mask=[True, True])                                                                                                                                    

In [6]: dask_all_masked.std().compute()                                                                                                                                                                    
---------------------------------------------------------------------------
MaskError                                 Traceback (most recent call last)
<ipython-input-6-dd8d98362ed2> in <module>
----> 1 dask_all_masked.std().compute()

~/dl/forks/dask/dask/base.py in compute(self, **kwargs)
    173         dask.base.compute
    174         """
--> 175         (result,) = compute(self, traverse=False, **kwargs)
    176         return result
    177 

~/dl/forks/dask/dask/base.py in compute(*args, **kwargs)
    444     keys = [x.__dask_keys__() for x in collections]
    445     postcomputes = [x.__dask_postcompute__() for x in collections]
--> 446     results = schedule(dsk, keys, **kwargs)
    447     return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
    448 

~/dl/forks/dask/dask/threaded.py in get(dsk, result, cache, num_workers, pool, **kwargs)
     80         get_id=_thread_get_id,
     81         pack_exception=pack_exception,
---> 82         **kwargs
     83     )
     84 

~/dl/forks/dask/dask/local.py in get_async(apply_async, num_workers, dsk, result, cache, get_id, rerun_exceptions_locally, pack_exception, raise_exception, callbacks, dumps, loads, **kwargs)
    489                         _execute_task(task, data)  # Re-execute locally
    490                     else:
--> 491                         raise_exception(exc, tb)
    492                 res, worker_id = loads(res_info)
    493                 state["cache"][key] = res

~/dl/forks/dask/dask/compatibility.py in reraise(exc, tb)
    128         if exc.__traceback__ is not tb:
    129             raise exc.with_traceback(tb)
--> 130         raise exc
    131 
    132     import pickle as cPickle

~/dl/forks/dask/dask/local.py in execute_task(key, task_info, dumps, loads, get_id, pack_exception)
    231     try:
    232         task, data = loads(task_info)
--> 233         result = _execute_task(task, data)
    234         id = get_id()
    235         result = dumps((result, id))

~/dl/forks/dask/dask/core.py in _execute_task(arg, cache, dsk)
    116     elif istask(arg):
    117         func, args = arg[0], arg[1:]
--> 118         args2 = [_execute_task(a, cache) for a in args]
    119         return func(*args2)
    120     elif not ishashable(arg):

~/dl/forks/dask/dask/core.py in <listcomp>(.0)
    116     elif istask(arg):
    117         func, args = arg[0], arg[1:]
--> 118         args2 = [_execute_task(a, cache) for a in args]
    119         return func(*args2)
    120     elif not ishashable(arg):

~/dl/forks/dask/dask/core.py in _execute_task(arg, cache, dsk)
    117         func, args = arg[0], arg[1:]
    118         args2 = [_execute_task(a, cache) for a in args]
--> 119         return func(*args2)
    120     elif not ishashable(arg):
    121         return arg

~/dl/forks/dask/dask/array/reductions.py in moment_agg(pairs, order, ddof, dtype, sum, axis, computing_meta, **kwargs)
    731             denominator = np.nan
    732     else:
--> 733         denominator[denominator < 0] = np.nan
    734 
    735     return divide(M, denominator, dtype=dtype)

~/.local/share/virtualenvs/dask-dyglJSq3/lib/python3.7/site-packages/numpy/ma/core.py in __setitem__(self, indx, value)
   3288         """
   3289         if self is masked:
-> 3290             raise MaskError('Cannot alter the masked element.')
   3291         _data = self._data
   3292         _mask = self._mask

MaskError: Cannot alter the masked element.
```

Discussion:
gjoseph92: @mrocklin 2d8dd51 is my initial thought on how to resolve #5338, but I really don't know what to do about `sqrt`'s handling of masked constants. Looking for some ideas here.
mrocklin: @jcrist do you have a moment to look through this?  I've never really touched masked arrays before.
mrocklin: Also, I wonder if @bjlittle could recommend someone to help review this
mrocklin: cc'ing also @DPeterK in case he knows someone comfortable with reviewing masked array code.
TomAugspurger: Hows this going @gjoseph92? Is there anything in particular that others like @bjlittle or @DPeterK can help out with?
mrocklin: @TomAugspurger , this is currently blocked on people able to review masked array code.  

It seems like we no longer have anyone qualified who is able to review these things.  I wonder if we should stop supporting masked arrays.  @jacobtomlinson do you know if this is still used within UK Met?  If so, do you know anyone there who we could lean on to help?
jacobtomlinson: AFAIK it is and I think @DPeterK is the right person to ping, he's just on vacation right now.
gjoseph92: @TomAugspurger I've been out on vacation myself, so just catching up on this.

The core issue is that dask's data model for masked arrays doesn't quite line up with NumPy's (and frankly, kind of abuses dask's)—NumPy uses a separate subclass, while dask holds them in a plain dask array, and just knows that the functions in the graph will produce masked arrays. I don't believe there's a definitive way to tell if a dask array is masked or not without actually computing it. I think `_meta` isn't reliable; I remember some cases where `_meta` would be explicitly instantiated as a `np.array` regardless of input type, though of course I can't remember where now. More generally, I'd guess dask just doesn't play well with any ndarray subclasses, MaskedArray included.

Besides making a separate dask MaskedArray subclass (lot of effort), we could at least expose dask versions of all the `np.ma` ufuncs, so you'd have access to those versions when you needed them. But then you'd also want masked versions of the reductions (`da.masked_std` would need to use `da.masked_sqrt`, for example). Still a lot of effort and code duplication.

We could also make the dask ufuncs automatically switch between masked and non-masked versions on the type of the input array (like I mentioned [here](https://github.com/dask/dask/pull/5339#discussion_r319151225)). But that seems bad, because users can't control it, and it changes current behavior out from under them.

Or, to get this merged, I can just remove the test for `std` or `@pytest.mark.xfail` it. The fact that masked arrays likely need their own versions of many functions to behave correctly is its own, much larger issue.
mrocklin: This appears to be stalled.  @jcrist you're probably the best equiped person to handle this, but I can understand that it might not be high in your priority list.  I thought I'd make you aware of it in case you have some extra time.

@gjoseph92 I apologize that no one has been found who is able to review this work so far.
jcrist: Thanks @gjoseph92. I've updated this PR to make the required fixes a bit more localized. Thanks for being patient here, this one slipped through the cracks (apologies).
jcrist: Hmmm, tests are failing for numpy 1.15. Debugging.
jcrist: Was a bug in the test, fixed now.
gjoseph92: Wow thanks @jcrist! Wasn't expecting this to get solved, but we've still been running into it, so I really appreciate the help. `safe_sqrt` is a good solution here; looks good to me.
jcrist: Failure is unrelated. cc @jrbourbeau or @TomAugspurger for a quick review, but I think this is good-to-go.

Patch:
diff --git a/dask/array/reductions.py b/dask/array/reductions.py
index 29685ba6c..f21130bff 100644
--- a/dask/array/reductions.py
+++ b/dask/array/reductions.py
@@ -15,7 +15,6 @@ from .core import _concatenate2, Array, handle_out, implements
 from .blockwise import blockwise
 from ..blockwise import lol_tuples
 from .creation import arange, diagonal
-from .ufunc import sqrt
 from .utils import full_like_safe, validate_axis, compute_meta, is_arraylike
 from .wrap import zeros, ones
 from .numpy_compat import ma_divide, divide as np_divide
@@ -727,7 +726,7 @@ def moment_agg(
     if isinstance(denominator, Number):
         if denominator < 0:
             denominator = np.nan
-    else:
+    elif denominator is not np.ma.masked:
         denominator[denominator < 0] = np.nan
 
     return divide(M, denominator, dtype=dtype)
@@ -812,9 +811,30 @@ with ignoring(AttributeError):
     nanvar = derived_from(np)(nanvar)
 
 
+def _sqrt(a):
+    o = np.sqrt(a)
+    if isinstance(o, np.ma.masked_array) and not o.shape and o.mask.all():
+        return np.ma.masked
+    return o
+
+
+def safe_sqrt(a):
+    """A version of sqrt that properly handles scalar masked arrays.
+
+    To mimic ``np.ma`` reductions, we need to convert scalar masked arrays that
+    have an active mask to the ``np.ma.masked`` singleton. This is properly
+    handled automatically for reduction code, but not for ufuncs. We implement
+    a simple version here, since calling `np.ma.sqrt` everywhere is
+    significantly more expensive.
+    """
+    if hasattr(a, "_elemwise"):
+        return a._elemwise(_sqrt, a)
+    return _sqrt(a)
+
+
 @derived_from(np)
 def std(a, axis=None, dtype=None, keepdims=False, ddof=0, split_every=None, out=None):
-    result = sqrt(
+    result = safe_sqrt(
         var(
             a,
             axis=axis,
@@ -834,7 +854,7 @@ def std(a, axis=None, dtype=None, keepdims=False, ddof=0, split_every=None, out=
 def nanstd(
     a, axis=None, dtype=None, keepdims=False, ddof=0, split_every=None, out=None
 ):
-    result = sqrt(
+    result = safe_sqrt(
         nanvar(
             a,
             axis=axis,
diff --git a/dask/dataframe/methods.py b/dask/dataframe/methods.py
index 92e0f0468..5d2ed4ff6 100644
--- a/dask/dataframe/methods.py
+++ b/dask/dataframe/methods.py
@@ -7,6 +7,7 @@ from tlz import partition
 
 from .utils import (
     is_series_like,
+    is_index_like,
     is_dataframe_like,
     PANDAS_GT_0250,
     hash_object_dispatch,
@@ -271,7 +272,7 @@ def unique(x, series_name=None):
     out = x.unique()
     # out can be either an np.ndarray or may already be a series
     # like object.  When out is an np.ndarray, it must be wrapped.
-    if not is_series_like(out):
+    if not (is_series_like(out) or is_index_like(out)):
         out = pd.Series(out, name=series_name)
     return out
 


Test patch:
diff --git a/dask/array/tests/test_masked.py b/dask/array/tests/test_masked.py
index 9c4260b31..adeb81cc5 100644
--- a/dask/array/tests/test_masked.py
+++ b/dask/array/tests/test_masked.py
@@ -225,12 +225,15 @@ def test_filled():
 
 def assert_eq_ma(a, b):
     res = a.compute()
-    assert type(res) == type(b)
-    if hasattr(res, "mask"):
-        np.testing.assert_equal(res.mask, b.mask)
-        a = da.ma.filled(a)
-        b = np.ma.filled(b)
-    assert_eq(a, b, equal_nan=True)
+    if res is np.ma.masked:
+        assert res is b
+    else:
+        assert type(res) == type(b)
+        if hasattr(res, "mask"):
+            np.testing.assert_equal(res.mask, b.mask)
+            a = da.ma.filled(a)
+            b = np.ma.filled(b)
+        assert_eq(a, b, equal_nan=True)
 
 
 @pytest.mark.parametrize("dtype", ("i8", "f8"))
@@ -261,6 +264,20 @@ def test_reductions(dtype, reduction):
     )
 
 
+@pytest.mark.parametrize("dtype", ("i8", "f8"))
+@pytest.mark.parametrize(
+    "reduction", ["sum", "prod", "mean", "var", "std", "min", "max", "any", "all"]
+)
+def test_reductions_allmasked(dtype, reduction):
+    x = np.ma.masked_array([1, 2], mask=True)
+    dx = da.from_array(x, asarray=False)
+
+    dfunc = getattr(da, reduction)
+    func = getattr(np, reduction)
+
+    assert_eq_ma(dfunc(dx), func(x))
+
+
 @pytest.mark.parametrize("reduction", ["argmin", "argmax"])
 def test_arg_reductions(reduction):
     x = np.random.random((10, 10, 10))
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
index e17797a20..5b14fc4ff 100644
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -1037,6 +1037,8 @@ def test_unique():
 
     assert_eq(ddf.x.unique(split_every=2), pd.Series(pdf.x.unique(), name="x"))
     assert_eq(ddf.y.unique(split_every=2), pd.Series(pdf.y.unique(), name="y"))
+    assert_eq(ddf.index.unique(), pdf.index.unique())
+
     assert ddf.x.unique(split_every=2)._name != ddf.x.unique()._name